@Test
  public void TestSingleBroadcastExchangeWithTwoScans() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();

    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        Drillbit bit2 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {

      bit1.run();
      bit2.run();
      client.connect();

      String physicalPlan =
          Files.toString(
                  FileUtils.getResourceAsFile("/sender/broadcast_exchange.json"), Charsets.UTF_8)
              .replace(
                  "#{LEFT_FILE}",
                  FileUtils.getResourceAsFile("/join/merge_single_batch.left.json")
                      .toURI()
                      .toString())
              .replace(
                  "#{RIGHT_FILE}",
                  FileUtils.getResourceAsFile("/join/merge_single_batch.right.json")
                      .toURI()
                      .toString());
      List<QueryResultBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan);
      int count = 0;
      for (QueryResultBatch b : results) {
        if (b.getHeader().getRowCount() != 0) count += b.getHeader().getRowCount();
        b.release();
      }
      assertEquals(25, count);
    }
  }
예제 #2
0
  @Test // DRILL-3580
  public void testExpressionInWindowFunction() throws Exception {
    String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
    String query =
        String.format(
            "select a1, b1, sum(b1) over (partition by a1) as c1, sum(a1 + b1) over (partition by a1) as c2\n"
                + "from dfs_test.`%s`",
            root);

    // Validate the plan
    final String[] expectedPlan = {
      "Window\\(window#0=\\[window\\(partition \\{0\\} order by \\[\\].*\\[SUM\\(\\$1\\), SUM\\(\\$2\\)\\]"
    };
    PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, new String[] {});

    testBuilder()
        .sqlQuery(query)
        .unOrdered()
        .baselineColumns("a1", "b1", "c1", "c2")
        .baselineValues(0l, 1l, 8l, 8l)
        .baselineValues(0l, 1l, 8l, 8l)
        .baselineValues(0l, 2l, 8l, 8l)
        .baselineValues(0l, 2l, 8l, 8l)
        .baselineValues(0l, 2l, 8l, 8l)
        .baselineValues(10l, 3l, 21l, 71l)
        .baselineValues(10l, 3l, 21l, 71l)
        .baselineValues(10l, 5l, 21l, 71l)
        .baselineValues(10l, 5l, 21l, 71l)
        .baselineValues(10l, 5l, 21l, 71l)
        .build()
        .run();
  }
  // specific tests should call this method, but it is not marked as a test itself intentionally
  public void testParquetFullEngineEventBased(
      boolean testValues,
      boolean generateNew,
      String plan,
      String readEntries,
      String filename,
      int numberOfTimesRead /* specified in json plan */,
      ParquetTestProperties props,
      QueryType queryType)
      throws Exception {
    if (generateNew) TestFileGenerator.generateParquetFile(filename, props);

    ParquetResultListener resultListener =
        new ParquetResultListener(getAllocator(), props, numberOfTimesRead, testValues);
    long C = System.nanoTime();
    String planText = Files.toString(FileUtils.getResourceAsFile(plan), Charsets.UTF_8);
    // substitute in the string for the read entries, allows reuse of the plan file for several
    // tests
    if (readEntries != null) {
      planText = planText.replaceFirst("&REPLACED_IN_PARQUET_TEST&", readEntries);
    }
    this.testWithListener(queryType, planText, resultListener);
    resultListener.getResults();
    long D = System.nanoTime();
    System.out.println(String.format("Took %f s to run query", (float) (D - C) / 1E9));
  }
예제 #4
0
  @Test
  public void twoBitTwoExchangeTwoEntryRun() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();

    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        Drillbit bit2 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator()); ) {

      bit1.run();
      bit2.run();
      client.connect();
      List<QueryResultBatch> results =
          client.runQuery(
              org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
              Files.toString(
                  FileUtils.getResourceAsFile("/sender/union_exchange.json"), Charsets.UTF_8));
      int count = 0;
      for (QueryResultBatch b : results) {
        if (b.getHeader().getRowCount() != 0) {
          count += b.getHeader().getRowCount();
        }
        b.release();
      }
      assertEquals(150, count);
    }
  }
  @Test
  public void testMultipleProvidersMixedSizes() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();

    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        Drillbit bit2 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator()); ) {

      bit1.run();
      bit2.run();
      client.connect();
      List<QueryResultBatch> results =
          client.runQuery(
              org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
              Files.toString(
                  FileUtils.getResourceAsFile("/mergerecv/multiple_providers.json"),
                  Charsets.UTF_8));
      int count = 0;
      RecordBatchLoader batchLoader = new RecordBatchLoader(client.getAllocator());
      // print the results
      Long lastBlueValue = null;
      for (QueryResultBatch b : results) {
        count += b.getHeader().getRowCount();
        for (int valueIdx = 0; valueIdx < b.getHeader().getRowCount(); valueIdx++) {
          List<Object> row = Lists.newArrayList();
          batchLoader.load(b.getHeader().getDef(), b.getData());
          for (VectorWrapper vw : batchLoader) {
            row.add(
                vw.getValueVector().getField().toExpr()
                    + ":"
                    + vw.getValueVector().getAccessor().getObject(valueIdx));
            if (vw.getValueVector()
                .getField()
                .getAsSchemaPath()
                .getRootSegment()
                .getPath()
                .equals("blue")) {
              // assert order is ascending
              if (((Long) vw.getValueVector().getAccessor().getObject(valueIdx)).longValue() == 0)
                continue; // ignore initial 0's from sort
              if (lastBlueValue != null)
                assertTrue(
                    ((Long) vw.getValueVector().getAccessor().getObject(valueIdx)).longValue()
                        >= ((Long) lastBlueValue).longValue());
              lastBlueValue = (Long) vw.getValueVector().getAccessor().getObject(valueIdx);
            }
          }
          for (Object cell : row) {
            int len = cell.toString().length();
            System.out.print(cell + " ");
            for (int i = 0; i < (30 - len); ++i) System.out.print(" ");
          }
          System.out.println();
        }
        b.release();
        batchLoader.clear();
      }
      assertEquals(400, count);
    }
  }
예제 #6
0
  public void testCommon(String[] expectedResults, String physicalPlan, String resourceFile)
      throws Exception {
    try (RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();
        Drillbit bit = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {

      // run query.
      bit.run();
      client.connect();
      List<QueryDataBatch> results =
          client.runQuery(
              org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
              Files.toString(FileUtils.getResourceAsFile(physicalPlan), Charsets.UTF_8)
                  .replace("#{TEST_FILE}", resourceFile));

      RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator());

      QueryDataBatch batch = results.get(0);
      assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData()));

      int i = 0;
      for (VectorWrapper<?> v : batchLoader) {

        ValueVector.Accessor accessor = v.getValueVector().getAccessor();
        System.out.println(accessor.getObject(0));
        assertEquals(expectedResults[i++], accessor.getObject(0).toString());
      }

      batchLoader.clear();
      for (QueryDataBatch b : results) {
        b.release();
      }
    }
  }
  @Test
  public void TestMultipleSendLocationBroadcastExchange() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();

    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        Drillbit bit2 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator())) {

      bit1.run();
      bit2.run();
      client.connect();

      String physicalPlan =
          Files.toString(
              FileUtils.getResourceAsFile("/sender/broadcast_exchange_long_run.json"),
              Charsets.UTF_8);
      List<QueryResultBatch> results = client.runQuery(QueryType.PHYSICAL, physicalPlan);
      int count = 0;
      for (QueryResultBatch b : results) {
        if (b.getHeader().getRowCount() != 0) count += b.getHeader().getRowCount();
        b.release();
      }
      System.out.println(count);
    }
  }
예제 #8
0
  @Test // see DRILL-3574
  public void testWithAndWithoutPartitions() throws Exception {
    String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
    String query =
        String.format(
            "select sum(a1) over(partition by b1, c1) as s1, sum(a1) over() as s2 \n"
                + "from dfs_test.`%s` \n"
                + "order by a1",
            root);
    test("alter session set `planner.slice_target` = 1");

    // Validate the plan
    final String[] expectedPlan = {
      "Window\\(window#0=\\[window\\(partition \\{\\}.*\n" + ".*UnionExchange"
    };
    PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, new String[] {});

    testBuilder()
        .sqlQuery(query)
        .unOrdered()
        .baselineColumns("s1", "s2")
        .baselineValues(0l, 50l)
        .baselineValues(0l, 50l)
        .baselineValues(0l, 50l)
        .baselineValues(0l, 50l)
        .baselineValues(0l, 50l)
        .baselineValues(10l, 50l)
        .baselineValues(10l, 50l)
        .baselineValues(10l, 50l)
        .baselineValues(20l, 50l)
        .baselineValues(20l, 50l)
        .build()
        .run();
  }
  @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());
  }
예제 #10
0
  public Fragment getRootFragment(PhysicalPlanReader reader, String file)
      throws FragmentSetupException, IOException {
    MakeFragmentsVisitor f = new MakeFragmentsVisitor();

    PhysicalPlan plan =
        reader.readPhysicalPlan(Files.toString(FileUtils.getResourceAsFile(file), Charsets.UTF_8));
    PhysicalOperator o = plan.getSortedOperators(false).iterator().next();
    return o.accept(f, null);
  }
 public void testParquetFullEngineLocalPath(
     String planFileName,
     String filename,
     int numberOfTimesRead /* specified in json plan */,
     int numberOfRowGroups,
     int recordsPerRowGroup)
     throws Exception {
   testParquetFullEngineLocalText(
       Files.toString(FileUtils.getResourceAsFile(planFileName), Charsets.UTF_8),
       filename,
       numberOfTimesRead,
       numberOfRowGroups,
       recordsPerRowGroup,
       true);
 }
  @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);
    }
  }
  @Test
  public void twoBitTwoExchange() throws Exception {
    RemoteServiceSet serviceSet = RemoteServiceSet.getLocalServiceSet();

    try (Drillbit bit1 = new Drillbit(CONFIG, serviceSet);
        Drillbit bit2 = new Drillbit(CONFIG, serviceSet);
        DrillClient client = new DrillClient(CONFIG, serviceSet.getCoordinator()); ) {

      bit1.run();
      bit2.run();
      client.connect();
      List<QueryResultBatch> results =
          client.runQuery(
              org.apache.drill.exec.proto.UserBitShared.QueryType.PHYSICAL,
              Files.toString(
                  FileUtils.getResourceAsFile("/mergerecv/merging_receiver.json"), Charsets.UTF_8));
      int count = 0;
      RecordBatchLoader batchLoader = new RecordBatchLoader(client.getAllocator());
      // print the results
      for (QueryResultBatch b : results) {
        count += b.getHeader().getRowCount();
        for (int valueIdx = 0; valueIdx < b.getHeader().getRowCount(); valueIdx++) {
          List<Object> row = Lists.newArrayList();
          batchLoader.load(b.getHeader().getDef(), b.getData());
          for (VectorWrapper<?> vw : batchLoader)
            row.add(
                vw.getValueVector().getField().toExpr()
                    + ":"
                    + vw.getValueVector().getAccessor().getObject(valueIdx));
          for (Object cell : row) {
            if (cell == null) {
              System.out.print("<null>    ");
              continue;
            }
            int len = cell.toString().length();
            System.out.print(cell + " ");
            for (int i = 0; i < (30 - len); ++i) System.out.print(" ");
          }
          System.out.println();
        }
        b.release();
        batchLoader.clear();
      }
      assertEquals(200, count);
    }
  }
 // use this method to submit physical plan
 public void testParquetFullEngineLocalTextDistributed(
     String planName,
     String filename,
     int numberOfTimesRead /* specified in json plan */,
     int numberOfRowGroups,
     int recordsPerRowGroup)
     throws Exception {
   String planText = Files.toString(FileUtils.getResourceAsFile(planName), Charsets.UTF_8);
   testFull(
       QueryType.PHYSICAL,
       planText,
       filename,
       numberOfTimesRead,
       numberOfRowGroups,
       recordsPerRowGroup,
       true);
 }
 @Test
 public void testMultipleRowGroupsAndReads2() throws Exception {
   String readEntries;
   readEntries = "";
   // number of times to read the file
   int i = 3;
   for (int j = 0; j < i; j++) {
     readEntries += "\"" + fileName + "\"";
     if (j < i - 1) readEntries += ",";
   }
   String planText =
       Files.toString(
               FileUtils.getResourceAsFile("/parquet/parquet_scan_screen_read_entry_replace.json"),
               Charsets.UTF_8)
           .replaceFirst("&REPLACED_IN_PARQUET_TEST&", readEntries);
   testParquetFullEngineLocalText(
       planText, fileName, i, numberRowGroups, recordsPerRowGroup, true);
 }
 public void testParquetFullEngineRemote(
     String plan,
     String filename,
     int numberOfTimesRead /* specified in json plan */,
     int numberOfRowGroups,
     int recordsPerRowGroup)
     throws Exception {
   HashMap<String, FieldInfo> fields = new HashMap<>();
   ParquetTestProperties props =
       new ParquetTestProperties(
           numberRowGroups, recordsPerRowGroup, DEFAULT_BYTES_PER_PAGE, fields);
   TestFileGenerator.populateFieldInfoMap(props);
   ParquetResultListener resultListener =
       new ParquetResultListener(getAllocator(), props, numberOfTimesRead, true);
   testWithListener(
       QueryType.PHYSICAL,
       Files.toString(FileUtils.getResourceAsFile(plan), Charsets.UTF_8),
       resultListener);
   resultListener.getResults();
 }
  @Test
  @Ignore
  public void testDictionaryError() throws Exception {
    String readEntries;
    readEntries = "\"/tmp/lineitem_null_dict.parquet\"";

    String planText =
        Files.toString(
                FileUtils.getResourceAsFile("/parquet/parquet_scan_screen_read_entry_replace.json"),
                Charsets.UTF_8)
            .replaceFirst("&REPLACED_IN_PARQUET_TEST&", readEntries);
    // testParquetFullEngineLocalText(planText, fileName, 1, 1, 100000, false);

    testFull(
        QueryType.SQL,
        "select L_RECEIPTDATE from dfs.`/tmp/lineitem_null_dict.parquet`",
        "",
        1,
        1,
        100000,
        false);
  }
예제 #18
0
  @Test
  public void testCompoundIdentifierInWindowDefinition() throws Exception {
    String root =
        FileUtils.getResourceAsFile("/multilevel/csv/1994/Q1/orders_94_q1.csv").toURI().toString();
    String query =
        String.format(
            "SELECT count(*) OVER w as col1, count(*) OVER w as col2 \n"
                + "FROM dfs_test.`%s` \n"
                + "WINDOW w AS (PARTITION BY columns[1] ORDER BY columns[0] DESC)",
            root);

    // Validate the plan
    final String[] expectedPlan = {
      "Window.*partition \\{1\\} order by \\[0 DESC\\].*COUNT\\(\\)",
      "Scan.*columns=\\[`columns`\\[0\\], `columns`\\[1\\]\\]"
    };
    final String[] excludedPatterns = {"Scan.*columns=\\[`\\*`\\]"};
    PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, excludedPatterns);

    // Validate the result
    testBuilder()
        .sqlQuery(query)
        .ordered()
        .baselineColumns("col1", "col2")
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .baselineValues((long) 1, (long) 1)
        .build()
        .run();
  }
예제 #19
0
  @Test // DRILL-3567
  public void testMultiplePartitions2() throws Exception {
    String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
    String query =
        String.format(
            "select count(*) over(partition by b1 order by c1) as count1, \n"
                + "count(*) over(partition by a1 order by c1) as count2, \n"
                + "sum(a1)  over(partition by b1 order by c1) as sum1 \n"
                + "from dfs_test.`%s`",
            root);

    // Validate the plan
    final String[] expectedPlan = {
      "Window.*partition \\{2\\} order by \\[1\\].*COUNT\\(\\)",
      "Window.*partition \\{0\\} order by \\[1\\].*COUNT\\(\\), SUM\\(\\$2\\)",
      "Scan.*columns=\\[`b1`, `c1`, `a1`\\]"
    };
    final String[] excludedPatterns = {"Scan.*columns=\\[`\\*`\\]"};
    PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, excludedPatterns);

    testBuilder()
        .sqlQuery(query)
        .unOrdered()
        .baselineColumns("count1", "count2", "sum1")
        .baselineValues(1l, 2l, 0l)
        .baselineValues(1l, 2l, 0l)
        .baselineValues(2l, 5l, 0l)
        .baselineValues(3l, 5l, 0l)
        .baselineValues(3l, 5l, 0l)
        .baselineValues(1l, 2l, 10l)
        .baselineValues(1l, 2l, 10l)
        .baselineValues(2l, 5l, 20l)
        .baselineValues(3l, 5l, 30l)
        .baselineValues(3l, 5l, 30l)
        .build()
        .run();
  }
예제 #20
0
  @Test // see DRILL-3657
  public void testConstantsInMultiplePartitions() throws Exception {
    String root = FileUtils.getResourceAsFile("/store/text/data/t.json").toURI().toString();
    String query =
        String.format(
            "select sum(1) over(partition by b1 order by a1) as sum1, sum(1) over(partition by a1) as sum2, rank() over(order by b1) as rank1, rank() over(order by 1) as rank2 \n"
                + "from dfs_test.`%s` \n"
                + "order by 1, 2, 3, 4",
            root);

    // Validate the plan
    final String[] expectedPlan = {
      "Window.*SUM\\(\\$3\\).*\n"
          + ".*SelectionVectorRemover.*\n"
          + ".*Sort.*\n"
          + ".*Window.*SUM\\(\\$2\\).*"
    };
    PlanTestBase.testPlanMatchingPatterns(query, expectedPlan, new String[] {});

    testBuilder()
        .sqlQuery(query)
        .unOrdered()
        .baselineColumns("sum1", "sum2", "rank1", "rank2")
        .baselineValues(2l, 5l, 1l, 1l)
        .baselineValues(2l, 5l, 1l, 1l)
        .baselineValues(2l, 5l, 6l, 1l)
        .baselineValues(2l, 5l, 6l, 1l)
        .baselineValues(3l, 5l, 3l, 1l)
        .baselineValues(3l, 5l, 3l, 1l)
        .baselineValues(3l, 5l, 3l, 1l)
        .baselineValues(3l, 5l, 8l, 1l)
        .baselineValues(3l, 5l, 8l, 1l)
        .baselineValues(3l, 5l, 8l, 1l)
        .build()
        .run();
  }
 public String getPlanForFile(String pathFileName, String parquetFileName) throws IOException {
   return Files.toString(FileUtils.getResourceAsFile(pathFileName), Charsets.UTF_8)
       .replaceFirst("&REPLACED_IN_PARQUET_TEST&", parquetFileName);
 }
예제 #22
0
  @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();
  }
예제 #23
0
  @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());
  }