ClassCompilerSelector(DrillConfig config, OptionManager sessionOptions) {
OptionValue value = sessionOptions.getOption(JAVA_COMPILER_OPTION);
this.policy =
CompilerPolicy.valueOf(
(value != null)
? value.string_val.toUpperCase()
: config.getString(JAVA_COMPILER_CONFIG).toUpperCase());
value = sessionOptions.getOption(JAVA_COMPILER_JANINO_MAXSIZE_OPTION);
this.janinoThreshold =
(value != null) ? value.num_val : config.getLong(JAVA_COMPILER_JANINO_MAXSIZE_CONFIG);
value = sessionOptions.getOption(JAVA_COMPILER_DEBUG_OPTION);
boolean debug =
(value != null) ? value.bool_val : config.getBoolean(JAVA_COMPILER_DEBUG_CONFIG);
this.janinoClassCompiler =
(policy == CompilerPolicy.JANINO || policy == CompilerPolicy.DEFAULT)
? new JaninoClassCompiler(QueryClassLoader.this, debug)
: null;
this.jdkClassCompiler =
(policy == CompilerPolicy.JDK || policy == CompilerPolicy.DEFAULT)
? new JDKClassCompiler(QueryClassLoader.this, debug)
: null;
}
public ExternalSortBatch(ExternalSort popConfig, FragmentContext context, RecordBatch incoming)
throws OutOfMemoryException {
super(popConfig, context, true);
this.incoming = incoming;
DrillConfig config = context.getConfig();
Configuration conf = new Configuration();
conf.set("fs.default.name", config.getString(ExecConstants.EXTERNAL_SORT_SPILL_FILESYSTEM));
try {
this.fs = FileSystem.get(conf);
} catch (IOException e) {
throw new RuntimeException(e);
}
SPILL_BATCH_GROUP_SIZE = config.getInt(ExecConstants.EXTERNAL_SORT_SPILL_GROUP_SIZE);
SPILL_THRESHOLD = config.getInt(ExecConstants.EXTERNAL_SORT_SPILL_THRESHOLD);
dirs = Iterators.cycle(config.getStringList(ExecConstants.EXTERNAL_SORT_SPILL_DIRS));
oAllocator = oContext.getAllocator();
copierAllocator =
oAllocator.newChildAllocator(
oAllocator.getName() + ":copier",
PriorityQueueCopier.INITIAL_ALLOCATION,
PriorityQueueCopier.MAX_ALLOCATION);
FragmentHandle handle = context.getHandle();
fileName =
String.format(
"%s_majorfragment%s_minorfragment%s_operator%s",
QueryIdHelper.getQueryId(handle.getQueryId()),
handle.getMajorFragmentId(),
handle.getMinorFragmentId(),
popConfig.getOperatorId());
}
@Test
public void testSubstringNegative(
@Injectable final DrillbitContext bitContext,
@Injectable UserServer.UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(
FileUtils.getResourceAsFile("/functions/testSubstringNegative.json"),
Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
NullableVarCharVector c1 =
exec.getValueVectorById(
new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarCharVector.class);
NullableVarCharVector.Accessor a1;
a1 = c1.getAccessor();
int count = 0;
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
if (!a1.isNull(i)) {
NullableVarCharHolder holder = new NullableVarCharHolder();
a1.get(i, holder);
// when offset is negative, substring return empty string.
assertEquals("", holder.toString());
++count;
}
}
assertEquals(50, count);
}
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}
@Test
public void test(final DrillbitContext bitContext) throws Exception {
final DrillConfig c = DrillConfig.create();
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getConfig();
result = c;
bitContext.getCache();
result = new LocalCache();
}
};
bitContext.getCache().run();
StoragePluginRegistry r = new StoragePluginRegistry(bitContext);
SchemaPlus plus = Frameworks.createRootSchema();
r.init();
r.getSchemaFactory().registerSchemas(null, plus);
printSchema(plus, 0);
}
public FunctionRegistry(DrillConfig config) {
try {
Set<Class<? extends CallProvider>> providerClasses =
PathScanner.scanForImplementations(
CallProvider.class,
config.getStringList(CommonConstants.LOGICAL_FUNCTION_SCAN_PACKAGES));
Map<String, FunctionDefinition> funcs = new HashMap<String, FunctionDefinition>();
for (Class<? extends CallProvider> c : providerClasses) {
CallProvider p = c.newInstance();
FunctionDefinition[] defs = p.getFunctionDefintions();
for (FunctionDefinition d : defs) {
for (String rn : d.getRegisteredNames()) {
FunctionDefinition d2 = funcs.put(rn, d);
// logger.debug("Registering function {}", d);
if (d2 != null) {
throw new ExceptionInInitializerError(
String.format(
"Failure while registering functions. The function %s tried to register with the name %s but the function %s already registered with that name.",
d.getName(), rn, d2.getName()));
}
}
}
}
funcMap = funcs;
} catch (Exception e) {
throw new RuntimeException("Failure while setting up FunctionRegistry.", e);
}
}
public static synchronized Class<?>[] getSubTypes(DrillConfig config) {
List<String> packages =
config.getStringList(CommonConstants.STORAGE_PLUGIN_CONFIG_SCAN_PACKAGES);
Class<?>[] sec = PathScanner.scanForImplementationsArr(FormatPluginConfig.class, packages);
logger.debug("Adding Format Plugin Configs including {}", (Object) sec);
return sec;
}
public class TestSimpleProjection {
static final org.slf4j.Logger logger =
org.slf4j.LoggerFactory.getLogger(TestSimpleProjection.class);
DrillConfig c = DrillConfig.create();
@Test
public void project(
@Injectable final DrillbitContext bitContext, @Injectable UserClientConnection connection)
throws Exception {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry("test");
bitContext.getAllocator();
result = BufferAllocator.getAllocator(c);
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile("/project/test1.json"), Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(
bitContext, FragmentHandle.getDefaultInstance(), connection, null, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
BigIntVector c1 =
exec.getValueVectorById(
new SchemaPath("col1", ExpressionPosition.UNKNOWN), BigIntVector.class);
BigIntVector c2 =
exec.getValueVectorById(
new SchemaPath("col2", ExpressionPosition.UNKNOWN), BigIntVector.class);
int x = 0;
BigIntVector.Accessor a1, a2;
a1 = c1.getAccessor();
a2 = c2.getAccessor();
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
assertEquals(a1.get(i) + 1, a2.get(i));
x += a1.get(i);
}
System.out.println(x);
}
}
@After
public void tearDown() throws Exception {
// pause to get logger to catch up.
Thread.sleep(1000);
}
}
private void setupSortMemoryAllocations(final PhysicalPlan plan) {
// look for external sorts
final List<ExternalSort> sortList = new LinkedList<>();
for (final PhysicalOperator op : plan.getSortedOperators()) {
if (op instanceof ExternalSort) {
sortList.add((ExternalSort) op);
}
}
// if there are any sorts, compute the maximum allocation, and set it on them
if (sortList.size() > 0) {
final OptionManager optionManager = queryContext.getOptions();
final long maxWidthPerNode =
optionManager.getOption(ExecConstants.MAX_WIDTH_PER_NODE_KEY).num_val;
long maxAllocPerNode =
Math.min(
DrillConfig.getMaxDirectMemory(),
queryContext.getConfig().getLong(ExecConstants.TOP_LEVEL_MAX_ALLOC));
maxAllocPerNode =
Math.min(
maxAllocPerNode,
optionManager.getOption(ExecConstants.MAX_QUERY_MEMORY_PER_NODE_KEY).num_val);
final long maxSortAlloc = maxAllocPerNode / (sortList.size() * maxWidthPerNode);
logger.debug("Max sort alloc: {}", maxSortAlloc);
for (final ExternalSort externalSort : sortList) {
externalSort.setMaxAllocation(maxSortAlloc);
}
}
}
public WorkspaceSchemaFactory(
DrillConfig drillConfig,
PStoreProvider provider,
FileSystemPlugin plugin,
String schemaName,
String storageEngineName,
DrillFileSystem fileSystem,
WorkspaceConfig config,
List<FormatMatcher> formatMatchers)
throws ExecutionSetupException, IOException {
this.fs = fileSystem;
this.plugin = plugin;
this.drillConfig = drillConfig;
this.config = config;
this.mapper = drillConfig.getMapper();
this.fileMatchers = Lists.newArrayList();
this.dirMatchers = Lists.newArrayList();
this.storageEngineName = storageEngineName;
this.schemaName = schemaName;
// setup cache
if (storageEngineName == null) {
this.knownViews = null;
// this.knownPaths = null;
} else {
this.knownViews =
provider.getPStore(
PStoreConfig //
.newJacksonBuilder(drillConfig.getMapper(), String.class) //
.name(Joiner.on('.').join("storage.views", storageEngineName, schemaName)) //
.build());
// this.knownPaths = provider.getPTable(PTableConfig //
// .newJacksonBuilder(drillConfig.getMapper(), String.class) //
// .name(Joiner.on('.').join("storage.cache", storageEngineName, schemaName)) //
// .build());
}
for (FormatMatcher m : formatMatchers) {
if (m.supportDirectoryReads()) {
dirMatchers.add(m);
}
fileMatchers.add(m);
}
}
/**
* Use reflection to scan for implementations of {@see FormatPlugin}.
*
* @param config - Drill configuration object, used to find the packages to scan
* @return - list of classes that implement the interface.
*/
public static synchronized Class<?>[] getSubTypes(final DrillConfig config) {
final List<String> packages =
config.getStringList(CommonConstants.STORAGE_PLUGIN_CONFIG_SCAN_PACKAGES);
final Class<?>[] pluginClasses =
PathScanner.scanForImplementationsArr(FormatPluginConfig.class, packages);
final String lineBrokenList =
pluginClasses.length == 0 ? "" : "\n\t- " + Joiner.on("\n\t- ").join(pluginClasses);
logger.debug(
"Found {} format plugin configuration classes: {}.", pluginClasses.length, lineBrokenList);
return pluginClasses;
}
public <T> T getWith(DrillConfig config, Class<T> c) {
try {
// logger.debug("Read tree {}", root);
return config.getMapper().treeToValue(root, c);
} catch (JsonProcessingException e) {
throw new LogicalPlanParsingException(
String.format(
"Failure while trying to convert late bound json options to type of %s. Reference was originally located at line %d, column %d.",
c.getCanonicalName(), location.getLineNr(), location.getColumnNr()),
e);
}
}
public UserServer(
DrillConfig config,
ScanResult classpathScan,
BufferAllocator alloc,
EventLoopGroup eventLoopGroup,
UserWorker worker,
Executor executor)
throws DrillbitStartupException {
super(
UserRpcConfig.getMapping(config, executor), alloc.getAsByteBufAllocator(), eventLoopGroup);
this.worker = worker;
this.alloc = alloc;
// TODO: move this up
if (config.getBoolean(ExecConstants.USER_AUTHENTICATION_ENABLED)) {
authenticator = UserAuthenticatorFactory.createAuthenticator(config, classpathScan);
} else {
authenticator = null;
}
if (config.getBoolean(ExecConstants.IMPERSONATION_ENABLED)) {
impersonationManager = new InboundImpersonationManager();
} else {
impersonationManager = null;
}
}
public DrillFunctionImplementationRegistry(DrillConfig config) {
FunctionConverter converter = new FunctionConverter();
Set<Class<? extends DrillFunc>> providerClasses =
PathScanner.scanForImplementations(
DrillFunc.class, config.getStringList(ExecConstants.FUNCTION_PACKAGES));
for (Class<? extends DrillFunc> clazz : providerClasses) {
DrillFuncHolder holder = converter.getHolder(clazz);
if (holder != null) {
// register handle for each name the function can be referred to
String[] names = holder.getRegisteredNames();
for (String name : names) methods.put(name, holder);
} else {
logger.warn("Unable to initialize function for class {}", clazz.getName());
}
}
}
public DrillbitEndpoint start() throws DrillbitStartupException, UnknownHostException {
int userPort =
userServer.bind(config.getInt(ExecConstants.INITIAL_USER_PORT), allowPortHunting);
String address =
useIP
? InetAddress.getLocalHost().getHostAddress()
: InetAddress.getLocalHost().getCanonicalHostName();
DrillbitEndpoint partialEndpoint =
DrillbitEndpoint.newBuilder()
.setAddress(address)
// .setAddress("localhost")
.setUserPort(userPort)
.build();
partialEndpoint = controller.start(partialEndpoint);
return dataPool.start(partialEndpoint);
}
@Test
@Ignore
public void testParseParquetPhysicalPlan() throws Exception {
RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
DrillConfig config = DrillConfig.create();
try (Drillbit bit1 = new Drillbit(config, serviceSet);
DrillClient client = new DrillClient(config, serviceSet.getCoordinator()); ) {
bit1.run();
client.connect();
List<QueryDataBatch> results =
client.runQuery(
org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
Resources.toString(Resources.getResource(fileName), Charsets.UTF_8));
RecordBatchLoader loader = new RecordBatchLoader(bit1.getContext().getAllocator());
int count = 0;
for (QueryDataBatch b : results) {
System.out.println(String.format("Got %d results", b.getHeader().getRowCount()));
count += b.getHeader().getRowCount();
loader.load(b.getHeader().getDef(), b.getData());
for (VectorWrapper vw : loader) {
System.out.print(vw.getValueVector().getField().toExpr() + ": ");
ValueVector vv = vw.getValueVector();
for (int i = 0; i < vv.getAccessor().getValueCount(); i++) {
Object o = vv.getAccessor().getObject(i);
if (o instanceof byte[]) {
System.out.print(" [" + new String((byte[]) o) + "]");
} else {
System.out.print(" [" + vv.getAccessor().getObject(i) + "]");
}
// break;
}
System.out.println();
}
loader.clear();
b.release();
}
client.close();
System.out.println(String.format("Got %d total results", count));
}
}
@Test
@Ignore
public void testParseParquetPhysicalPlanRemote() throws Exception {
DrillConfig config = DrillConfig.create();
try (DrillClient client = new DrillClient(config); ) {
client.connect();
ParquetResultsListener listener = new ParquetResultsListener();
Stopwatch watch = new Stopwatch();
watch.start();
client.runQuery(
org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
Resources.toString(Resources.getResource(fileName), Charsets.UTF_8),
listener);
System.out.println(
String.format(
"Got %d total records in %d seconds",
listener.await(), watch.elapsed(TimeUnit.SECONDS)));
client.close();
}
}
@Test
@Ignore
public void testPerformance(
@Injectable final DrillbitContext bitContext,
@Injectable UserServer.UserClientConnection connection)
throws Exception {
DrillConfig c = DrillConfig.create();
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(
bitContext, BitControl.PlanFragment.getDefaultInstance(), connection, registry);
// new NonStrictExpectations() {
// {
// context.getAllocator(); result = BufferAllocator.getAllocator(DrillConfig.create());
// }
// };
final String fileName = "/tmp/parquet_test_performance.parquet";
HashMap<String, FieldInfo> fields = new HashMap<>();
ParquetTestProperties props =
new ParquetTestProperties(1, 20 * 1000 * 1000, DEFAULT_BYTES_PER_PAGE, fields);
populateFieldInfoMap(props);
// generateParquetFile(fileName, props);
Configuration dfsConfig = new Configuration();
List<Footer> footers = ParquetFileReader.readFooters(dfsConfig, new Path(fileName));
Footer f = footers.iterator().next();
List<SchemaPath> columns = Lists.newArrayList();
columns.add(new SchemaPath("_MAP.integer", ExpressionPosition.UNKNOWN));
columns.add(new SchemaPath("_MAP.bigInt", ExpressionPosition.UNKNOWN));
columns.add(new SchemaPath("_MAP.f", ExpressionPosition.UNKNOWN));
columns.add(new SchemaPath("_MAP.d", ExpressionPosition.UNKNOWN));
columns.add(new SchemaPath("_MAP.b", ExpressionPosition.UNKNOWN));
columns.add(new SchemaPath("_MAP.bin", ExpressionPosition.UNKNOWN));
columns.add(new SchemaPath("_MAP.bin2", ExpressionPosition.UNKNOWN));
int totalRowCount = 0;
FileSystem fs = new CachedSingleFileSystem(fileName);
BufferAllocator allocator = new TopLevelAllocator();
for (int i = 0; i < 25; i++) {
ParquetRecordReader rr =
new ParquetRecordReader(
context,
256000,
fileName,
0,
fs,
new CodecFactoryExposer(dfsConfig),
f.getParquetMetadata(),
columns);
TestOutputMutator mutator = new TestOutputMutator(allocator);
rr.setup(mutator);
Stopwatch watch = new Stopwatch();
watch.start();
int rowCount = 0;
while ((rowCount = rr.next()) > 0) {
totalRowCount += rowCount;
}
System.out.println(
String.format("Time completed: %s. ", watch.elapsed(TimeUnit.MILLISECONDS)));
rr.cleanup();
}
allocator.close();
System.out.println(String.format("Total row count %s", totalRowCount));
}
public <T> T getListWith(DrillConfig config, TypeReference<T> t) throws IOException {
ObjectMapper mapper = config.getMapper();
return mapper.treeAsTokens(root).readValueAs(t);
// return mapper.treeToValue(root, mapper.getTypeFactory().constructCollectionType(List.class,
// c));
}
public class PlanningBase extends ExecTest {
static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(PlanningBase.class);
@Rule public final TestRule TIMEOUT = TestTools.getTimeoutRule(10000);
@Mocked DrillbitContext dbContext;
@Mocked QueryContext context;
private final DrillConfig config = DrillConfig.create();
protected void testSqlPlanFromFile(String file) throws Exception {
testSqlPlan(getFile(file));
}
protected void testSqlPlan(String sqlCommands) throws Exception {
String[] sqlStrings = sqlCommands.split(";");
final DistributedCache cache = new LocalCache();
cache.run();
final SystemOptionManager opt = new SystemOptionManager(cache);
opt.init();
final OptionManager sess = new SessionOptionManager(opt);
new NonStrictExpectations() {
{
dbContext.getMetrics();
result = new MetricRegistry();
dbContext.getAllocator();
result = new TopLevelAllocator();
dbContext.getConfig();
result = config;
dbContext.getOptionManager();
result = opt;
dbContext.getCache();
result = cache;
}
};
final StoragePluginRegistry registry = new StoragePluginRegistry(dbContext);
registry.init();
final FunctionImplementationRegistry functionRegistry =
new FunctionImplementationRegistry(config);
final SchemaPlus root = Frameworks.createRootSchema(false);
registry.getSchemaFactory().registerSchemas(new UserSession(null, null, null), root);
new NonStrictExpectations() {
{
context.getNewDefaultSchema();
result = root;
context.getStorage();
result = registry;
context.getFunctionRegistry();
result = functionRegistry;
context.getSession();
result = new UserSession(null, null, null);
context.getCurrentEndpoint();
result = DrillbitEndpoint.getDefaultInstance();
context.getActiveEndpoints();
result = ImmutableList.of(DrillbitEndpoint.getDefaultInstance());
context.getPlannerSettings();
result = new PlannerSettings(sess);
context.getOptions();
result = sess;
context.getConfig();
result = config;
context.getCache();
result = cache;
}
};
for (String sql : sqlStrings) {
if (sql.trim().isEmpty()) continue;
DrillSqlWorker worker = new DrillSqlWorker(context);
PhysicalPlan p = worker.getPlan(sql);
}
}
protected String getFile(String resource) throws IOException {
URL url = Resources.getResource(resource);
if (url == null) {
throw new IOException(String.format("Unable to find path %s.", resource));
}
return Resources.toString(url, Charsets.UTF_8);
}
}
@BeforeClass
public static void setup() {
CONFIG = DrillConfig.create();
}
@Before
public void init() {
drillConfig = DrillConfig.create();
allocator = RootAllocatorFactory.newRoot(drillConfig);
}
@Test
public void testAllocators() throws Exception {
// Setup a drillbit (initializes a root allocator)
final DrillConfig config = DrillConfig.create(TEST_CONFIGURATIONS);
final RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
final Drillbit bit = new Drillbit(config, serviceSet);
bit.run();
final DrillbitContext bitContext = bit.getContext();
FunctionImplementationRegistry functionRegistry =
bitContext.getFunctionImplementationRegistry();
StoragePluginRegistry storageRegistry = new StoragePluginRegistry(bitContext);
// Create a few Fragment Contexts
BitControl.PlanFragment.Builder pfBuilder1 = BitControl.PlanFragment.newBuilder();
pfBuilder1.setMemInitial(1500000);
BitControl.PlanFragment pf1 = pfBuilder1.build();
BitControl.PlanFragment.Builder pfBuilder2 = BitControl.PlanFragment.newBuilder();
pfBuilder2.setMemInitial(500000);
BitControl.PlanFragment pf2 = pfBuilder1.build();
FragmentContext fragmentContext1 = new FragmentContext(bitContext, pf1, null, functionRegistry);
FragmentContext fragmentContext2 = new FragmentContext(bitContext, pf2, null, functionRegistry);
// Get a few physical operators. Easiest way is to read a physical plan.
PhysicalPlanReader planReader =
new PhysicalPlanReader(
config,
config.getMapper(),
CoordinationProtos.DrillbitEndpoint.getDefaultInstance(),
storageRegistry);
PhysicalPlan plan =
planReader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile(planFile), Charsets.UTF_8));
List<PhysicalOperator> physicalOperators = plan.getSortedOperators();
Iterator<PhysicalOperator> physicalOperatorIterator = physicalOperators.iterator();
PhysicalOperator physicalOperator1 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator2 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator3 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator4 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator5 = physicalOperatorIterator.next();
PhysicalOperator physicalOperator6 = physicalOperatorIterator.next();
// Create some bogus Operator profile defs and stats to create operator contexts
OpProfileDef def;
OperatorStats stats;
// Use some bogus operator type to create a new operator context.
def =
new OpProfileDef(
physicalOperator1.getOperatorId(),
UserBitShared.CoreOperatorType.MOCK_SUB_SCAN_VALUE,
OperatorContext.getChildCount(physicalOperator1));
stats = fragmentContext1.getStats().getOperatorStats(def, fragmentContext1.getAllocator());
// Add a couple of Operator Contexts
// Initial allocation = 1000000 bytes for all operators
OperatorContext oContext11 = fragmentContext1.newOperatorContext(physicalOperator1, true);
DrillBuf b11 = oContext11.getAllocator().buffer(1000000);
OperatorContext oContext12 =
fragmentContext1.newOperatorContext(physicalOperator2, stats, true);
DrillBuf b12 = oContext12.getAllocator().buffer(500000);
OperatorContext oContext21 = fragmentContext1.newOperatorContext(physicalOperator3, true);
def =
new OpProfileDef(
physicalOperator4.getOperatorId(),
UserBitShared.CoreOperatorType.TEXT_WRITER_VALUE,
OperatorContext.getChildCount(physicalOperator4));
stats = fragmentContext2.getStats().getOperatorStats(def, fragmentContext2.getAllocator());
OperatorContext oContext22 =
fragmentContext2.newOperatorContext(physicalOperator4, stats, true);
DrillBuf b22 = oContext22.getAllocator().buffer(2000000);
// New Fragment begins
BitControl.PlanFragment.Builder pfBuilder3 = BitControl.PlanFragment.newBuilder();
pfBuilder3.setMemInitial(1000000);
BitControl.PlanFragment pf3 = pfBuilder3.build();
FragmentContext fragmentContext3 = new FragmentContext(bitContext, pf3, null, functionRegistry);
// New fragment starts an operator that allocates an amount within the limit
def =
new OpProfileDef(
physicalOperator5.getOperatorId(),
UserBitShared.CoreOperatorType.UNION_VALUE,
OperatorContext.getChildCount(physicalOperator5));
stats = fragmentContext3.getStats().getOperatorStats(def, fragmentContext3.getAllocator());
OperatorContext oContext31 =
fragmentContext3.newOperatorContext(physicalOperator5, stats, true);
DrillBuf b31a = oContext31.getAllocator().buffer(200000);
// Previously running operator completes
b22.release();
((AutoCloseable) oContext22).close();
// Fragment 3 asks for more and fails
boolean outOfMem = false;
try {
DrillBuf b31b = oContext31.getAllocator().buffer(4400000);
if (b31b != null) {
b31b.release();
} else {
outOfMem = true;
}
} catch (Exception e) {
outOfMem = true;
}
assertEquals(true, (boolean) outOfMem);
// Operator is Exempt from Fragment limits. Fragment 3 asks for more and succeeds
outOfMem = false;
OperatorContext oContext32 = fragmentContext3.newOperatorContext(physicalOperator6, false);
DrillBuf b32 = null;
try {
b32 = oContext32.getAllocator().buffer(4400000);
} catch (Exception e) {
outOfMem = true;
} finally {
if (b32 != null) {
b32.release();
} else {
outOfMem = true;
}
closeOp(oContext32);
}
assertEquals(false, (boolean) outOfMem);
b11.release();
closeOp(oContext11);
b12.release();
closeOp(oContext12);
closeOp(oContext21);
b31a.release();
closeOp(oContext31);
fragmentContext1.close();
fragmentContext2.close();
fragmentContext3.close();
bit.close();
serviceSet.close();
}
@Test
public void sortOneKeyAscending(
@Injectable final DrillbitContext bitContext, @Injectable UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
bitContext.getCompiler();
result = CodeCompiler.getTestCompiler(c);
}
};
final PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
final PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(FileUtils.getResourceAsFile("/sort/one_key_sort.json"), Charsets.UTF_8));
final FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
final FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
final SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
int previousInt = Integer.MIN_VALUE;
int recordCount = 0;
int batchCount = 0;
while (exec.next()) {
batchCount++;
final IntVector c1 =
exec.getValueVectorById(
new SchemaPath("blue", ExpressionPosition.UNKNOWN), IntVector.class);
final IntVector c2 =
exec.getValueVectorById(
new SchemaPath("green", ExpressionPosition.UNKNOWN), IntVector.class);
final IntVector.Accessor a1 = c1.getAccessor();
final IntVector.Accessor a2 = c2.getAccessor();
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
recordCount++;
assertTrue(previousInt <= a1.get(i));
previousInt = a1.get(i);
assertEquals(previousInt, a2.get(i));
}
}
System.out.println(String.format("Sorted %,d records in %d batches.", recordCount, batchCount));
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}
public static Drillbit start(StartupOptions options) throws DrillbitStartupException {
return start(DrillConfig.create(options.getConfigLocation()));
}
@BeforeClass
public static void setUp() throws Exception {
config = DrillConfig.create();
zkUrl = config.getString(ExecConstants.ZK_CONNECTION);
setupTestDir();
}
public class TestSimpleFunctions extends ExecTest {
static final org.slf4j.Logger logger =
org.slf4j.LoggerFactory.getLogger(TestSimpleFunctions.class);
DrillConfig c = DrillConfig.create();
@Test
public void testHashFunctionResolution(@Injectable DrillConfig config)
throws JClassAlreadyExistsException, IOException {
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(config);
// test required vs nullable Int input
resolveHash(
config,
new TypedNullConstant(Types.optional(TypeProtos.MinorType.INT)),
Types.optional(TypeProtos.MinorType.INT),
Types.required(TypeProtos.MinorType.INT),
TypeProtos.DataMode.OPTIONAL,
registry);
resolveHash(
config,
new ValueExpressions.IntExpression(1, ExpressionPosition.UNKNOWN),
Types.required(TypeProtos.MinorType.INT),
Types.required(TypeProtos.MinorType.INT),
TypeProtos.DataMode.REQUIRED,
registry);
// test required vs nullable float input
resolveHash(
config,
new TypedNullConstant(Types.optional(TypeProtos.MinorType.FLOAT4)),
Types.optional(TypeProtos.MinorType.FLOAT4),
Types.required(TypeProtos.MinorType.FLOAT4),
TypeProtos.DataMode.OPTIONAL,
registry);
resolveHash(
config,
new ValueExpressions.FloatExpression(5.0f, ExpressionPosition.UNKNOWN),
Types.required(TypeProtos.MinorType.FLOAT4),
Types.required(TypeProtos.MinorType.FLOAT4),
TypeProtos.DataMode.REQUIRED,
registry);
// test required vs nullable long input
resolveHash(
config,
new TypedNullConstant(Types.optional(TypeProtos.MinorType.BIGINT)),
Types.optional(TypeProtos.MinorType.BIGINT),
Types.required(TypeProtos.MinorType.BIGINT),
TypeProtos.DataMode.OPTIONAL,
registry);
resolveHash(
config,
new ValueExpressions.LongExpression(100L, ExpressionPosition.UNKNOWN),
Types.required(TypeProtos.MinorType.BIGINT),
Types.required(TypeProtos.MinorType.BIGINT),
TypeProtos.DataMode.REQUIRED,
registry);
// test required vs nullable double input
resolveHash(
config,
new TypedNullConstant(Types.optional(TypeProtos.MinorType.FLOAT8)),
Types.optional(TypeProtos.MinorType.FLOAT8),
Types.required(TypeProtos.MinorType.FLOAT8),
TypeProtos.DataMode.OPTIONAL,
registry);
resolveHash(
config,
new ValueExpressions.DoubleExpression(100.0, ExpressionPosition.UNKNOWN),
Types.required(TypeProtos.MinorType.FLOAT8),
Types.required(TypeProtos.MinorType.FLOAT8),
TypeProtos.DataMode.REQUIRED,
registry);
}
public void resolveHash(
DrillConfig config,
LogicalExpression arg,
TypeProtos.MajorType expectedArg,
TypeProtos.MajorType expectedOut,
TypeProtos.DataMode expectedBestInputMode,
FunctionImplementationRegistry registry)
throws JClassAlreadyExistsException, IOException {
List<LogicalExpression> args = new ArrayList<>();
args.add(arg);
String[] registeredNames = {"hash"};
FunctionCall call = new FunctionCall("hash", args, ExpressionPosition.UNKNOWN);
FunctionResolver resolver = FunctionResolverFactory.getResolver(call);
DrillFuncHolder matchedFuncHolder = registry.findDrillFunction(resolver, call);
assertEquals(expectedBestInputMode, matchedFuncHolder.getParmMajorType(0).getMode());
}
@Test
public void testSubstring(
@Injectable final DrillbitContext bitContext,
@Injectable UserServer.UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(
FileUtils.getResourceAsFile("/functions/testSubstring.json"), Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
NullableVarCharVector c1 =
exec.getValueVectorById(
new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarCharVector.class);
NullableVarCharVector.Accessor a1;
a1 = c1.getAccessor();
int count = 0;
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
if (!a1.isNull(i)) {
NullableVarCharHolder holder = new NullableVarCharHolder();
a1.get(i, holder);
assertEquals("aaaa", holder.toString());
++count;
}
}
assertEquals(50, count);
}
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}
@Test
public void testSubstringNegative(
@Injectable final DrillbitContext bitContext,
@Injectable UserServer.UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(
FileUtils.getResourceAsFile("/functions/testSubstringNegative.json"),
Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
NullableVarCharVector c1 =
exec.getValueVectorById(
new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarCharVector.class);
NullableVarCharVector.Accessor a1;
a1 = c1.getAccessor();
int count = 0;
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
if (!a1.isNull(i)) {
NullableVarCharHolder holder = new NullableVarCharHolder();
a1.get(i, holder);
// when offset is negative, substring return empty string.
assertEquals("", holder.toString());
++count;
}
}
assertEquals(50, count);
}
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}
@Test
public void testByteSubstring(
@Injectable final DrillbitContext bitContext,
@Injectable UserServer.UserClientConnection connection)
throws Throwable {
new NonStrictExpectations() {
{
bitContext.getMetrics();
result = new MetricRegistry();
bitContext.getAllocator();
result = new TopLevelAllocator();
bitContext.getOperatorCreatorRegistry();
result = new OperatorCreatorRegistry(c);
bitContext.getConfig();
result = c;
}
};
PhysicalPlanReader reader =
new PhysicalPlanReader(
c, c.getMapper(), CoordinationProtos.DrillbitEndpoint.getDefaultInstance());
PhysicalPlan plan =
reader.readPhysicalPlan(
Files.toString(
FileUtils.getResourceAsFile("/functions/testByteSubstring.json"), Charsets.UTF_8));
FunctionImplementationRegistry registry = new FunctionImplementationRegistry(c);
FragmentContext context =
new FragmentContext(bitContext, PlanFragment.getDefaultInstance(), connection, registry);
SimpleRootExec exec =
new SimpleRootExec(
ImplCreator.getExec(
context, (FragmentRoot) plan.getSortedOperators(false).iterator().next()));
while (exec.next()) {
NullableVarBinaryVector c1 =
exec.getValueVectorById(
new SchemaPath("col3", ExpressionPosition.UNKNOWN), NullableVarBinaryVector.class);
NullableVarBinaryVector.Accessor a1;
a1 = c1.getAccessor();
int count = 0;
for (int i = 0; i < c1.getAccessor().getValueCount(); i++) {
if (!a1.isNull(i)) {
NullableVarBinaryHolder holder = new NullableVarBinaryHolder();
a1.get(i, holder);
assertEquals("aa", holder.toString());
++count;
}
}
assertEquals(50, count);
}
if (context.getFailureCause() != null) {
throw context.getFailureCause();
}
assertTrue(!context.isFailed());
}
}
private PhysicalPlan convert(LogicalPlan plan) throws OptimizerException {
if (logger.isDebugEnabled())
logger.debug("Converting logical plan {}.", plan.toJsonStringSafe(context.getConfig()));
return new BasicOptimizer(DrillConfig.create(), context)
.optimize(new BasicOptimizer.BasicOptimizationContext(), plan);
}