@Override public VectorAccessible next() { index++; if (index == dataBatches.size()) { throw new RuntimeException("Tried to call next when iterator had no more items."); } batchLoader.clear(); QueryDataBatch batch = dataBatches.get(index); try { batchLoader.load(batch.getHeader().getDef(), batch.getData()); } catch (SchemaChangeException e) { throw new RuntimeException(e); } return batchLoader; }
private static Set<String> extractSchema(QueryWritableBatch batch) { List<SerializedField> fields = batch.getHeader().getDef().getFieldList(); Set<String> schema = Sets.newHashSet(); for (SerializedField field : fields) { schema.add(field.getNamePart().getName()); } return schema; }
@Ignore("DX-3872") @Test public void testMergingReceiver() throws Exception { String plan = Files.toString(FileUtils.getResourceAsFile("/largefiles/merging_receiver_large_data.json"), Charsets.UTF_8).replace("#{TEST_FILE}", escapeJsonString(dataFile.getPath())); List<QueryDataBatch> results = testPhysicalWithResults(plan); int recordsInOutput = 0; for(QueryDataBatch batch : results) { recordsInOutput += batch.getHeader().getDef().getRecordCount(); batch.release(); } assertTrue(String.format("Number of records in output is wrong: expected=%d, actual=%s", NUM_RECORDS, recordsInOutput), NUM_RECORDS == recordsInOutput); }
@Test public void testColumnOrderingWithUnionVector() throws Exception { List<QueryDataBatch> results = null; try { results = testRunAndReturn(QueryType.SQL, "SELECT * FROM cp.\"type_changes.json\""); final RecordBatchDef def = results.get(0).getHeader().getDef(); assertEquals(2, def.getFieldCount()); assertEquals("a", def.getField(0).getNamePart().getName()); assertEquals("b", def.getField(1).getNamePart().getName()); } finally { if (results != null) { for(QueryDataBatch r : results) { r.release(); } } } }
public void runTest(Object[] values, String planPath, String dataPath) throws Throwable { List<QueryDataBatch> results = client.runQuery(QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile(planPath), Charsets.UTF_8).replace("#{TEST_FILE}", dataPath)); try (RecordBatchLoader batchLoader = new RecordBatchLoader(nodes[0].getContext().getAllocator())) { QueryDataBatch batch = results.get(1); assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData())); int i = 0; for (VectorWrapper<?> v : batchLoader) { ValueVector vv = v.getValueVector(); assertEquals(values[i++], (vv.getObject(0))); } } for (QueryDataBatch b : results) { b.release(); } }
@Test public void testNullableFilter() throws Exception { final List<QueryDataBatch> result = testSqlWithResults( "select count(wr_return_quantity) as row_count from dfs.\"/tmp/web_returns\" where wr_return_quantity = 1"); assertEquals("Only expected one batch with data, and then the empty finishing batch.", 2, result.size()); final RecordBatchLoader loader = new RecordBatchLoader(getSabotContext().getAllocator()); final QueryDataBatch b = result.get(0); loader.load(b.getHeader().getDef(), b.getData()); final VectorWrapper vw = loader.getValueAccessorById( BigIntVector.class, loader.getValueVectorId(SchemaPath.getCompoundPath("row_count")).getFieldIds() ); assertEquals(3573l, vw.getValueVector().getObject(0)); b.release(); loader.clear(); }
@Test public void testNullableAgg() throws Exception { final List<QueryDataBatch> result = testSqlWithResults( "select sum(a) as total_sum from dfs.\"/tmp/parquet_with_nulls_should_sum_100000_nulls_first.parquet\""); assertEquals("Only expected one batch with data, and then the empty finishing batch.", 2, result.size()); final RecordBatchLoader loader = new RecordBatchLoader(getSabotContext().getAllocator()); final QueryDataBatch b = result.get(0); loader.load(b.getHeader().getDef(), b.getData()); final VectorWrapper vw = loader.getValueAccessorById( BigIntVector.class, loader.getValueVectorId(SchemaPath.getCompoundPath("total_sum")).getFieldIds() ); assertEquals(4999950000l, vw.getValueVector().getObject(0)); b.release(); loader.clear(); }
protected int printResult(List<QueryDataBatch> results) throws SchemaChangeException { int rowCount = 0; final RecordBatchLoader loader = new RecordBatchLoader(getAllocator()); for(final QueryDataBatch result : results) { rowCount += result.getHeader().getRowCount(); loader.load(result.getHeader().getDef(), result.getData()); // TODO: Clean: DRILL-2933: That load(...) no longer throws // SchemaChangeException, so check/clean throw clause above. VectorUtil.showVectorAccessibleContent(loader, columnWidths); loader.clear(); result.release(); } System.out.println("Total record count: " + rowCount); return rowCount; }
protected void query(final String query, final Function<RecordBatchLoader> testBody) throws Exception { final List<QueryDataBatch> batches = testSqlWithResults(query); final RecordBatchLoader loader = new RecordBatchLoader(client.getRecordAllocator()); try { // first batch at index 0 is empty and used for fast schema return. Load the second one for the tests final QueryDataBatch batch = batches.get(0); loader.load(batch.getHeader().getDef(), batch.getData()); testBody.apply(loader); } finally { for (final QueryDataBatch batch:batches) { batch.release(); } loader.clear(); } }
@Test public void testRead() throws Exception { List<QueryDataBatch> results = testSqlWithResults( String.format("SELECT count(*) FROM dfs.\"%s\"", dataFile.getPath())); RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator()); for(QueryDataBatch batch : results) { batchLoader.load(batch.getHeader().getDef(), batch.getData()); if (batchLoader.getRecordCount() <= 0) { continue; } BigIntVector countV = batchLoader.getValueAccessorById(BigIntVector.class, 0).getValueVector(); assertTrue("Total of "+ NUM_RECORDS + " records expected in count", countV.get(0) == NUM_RECORDS); batchLoader.clear(); batch.release(); } }
public Builder mergeFrom(com.dremio.exec.proto.UserBitShared.QueryData other) { if (other == com.dremio.exec.proto.UserBitShared.QueryData.getDefaultInstance()) return this; if (other.hasQueryId()) { mergeQueryId(other.getQueryId()); } if (other.hasRowCount()) { setRowCount(other.getRowCount()); } if (other.hasDef()) { mergeDef(other.getDef()); } this.mergeUnknownFields(other.getUnknownFields()); return this; }
@Test public void test() throws Exception { ElasticsearchCluster.ColumnData[] data = new ElasticsearchCluster.ColumnData[]{ new ElasticsearchCluster.ColumnData("column", TEXT, new Object[][]{ {"value"} }) }; elastic.load(schema, table, data); List<QueryDataBatch> results = testSqlWithResults(String.format("select * from elasticsearch.%s.%s limit 0", schema, table)); MajorType t = results.get(0).getHeader().getDef().getField(0).getMajorType(); Assert.assertEquals("Mismatched type", MinorType.VARCHAR, t.getMinorType()); Assert.assertEquals("Mismatched datamode", DataMode.OPTIONAL, t.getMode()); } }
@Test public void readComplexWithStar() throws Exception { test("select * from cp.\"/store/json/test_complex_read_with_star.json\""); List<QueryDataBatch> results = testSqlWithResults("select * from cp.\"/store/json/test_complex_read_with_star.json\""); assertEquals(1, results.size()); RecordBatchLoader batchLoader = new RecordBatchLoader(getAllocator()); QueryDataBatch batch = results.get(0); assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData())); assertEquals(3, batchLoader.getSchema().getFieldCount()); testExistentColumns(batchLoader); batch.release(); batchLoader.clear(); }
private Map<SchemaPath, MajorType> getTypeMapFromBatch(QueryDataBatch batch) { Map<SchemaPath, MajorType> typeMap = new TreeMap<>(); for (int i = 0; i < batch.getHeader().getDef().getFieldCount(); i++) { typeMap.put(SchemaPath.getSimplePath(SerializedFieldHelper.create(batch.getHeader().getDef().getField(i)).getName()), batch.getHeader().getDef().getField(i).getMajorType()); } return typeMap; }
private void assertBatchCanBeLoaded(QueryWritableBatch batch, List<ValueVector> original) throws Exception { RpcOutcomeListener<GeneralRPCProtos.Ack> listener = Mockito.mock(RpcOutcomeListener.class); try (QueryDataBatch dataBatch = LocalUserUtil.acquireData(allocator, listener, batch); RecordBatchLoader loader = new RecordBatchLoader(allocator)) { loader.load(dataBatch.getHeader().getDef(), dataBatch.getData()); try (RecordBatchData loadedBatch = new RecordBatchData(loader, allocator)) { assertContainsFields(extractSchema(loadedBatch), original); } } }
private void checkResult(RecordBatchLoader loader, List<QueryDataBatch> results) { loader.load(results.get(0).getHeader().getDef(), results.get(0).getData()); RecordBatchDef def = results.get(0).getHeader().getDef(); // the entire row is returned as a single column assertEquals(1, def.getFieldCount()); assertTrue(def.getField(0).getMajorType().getMode() == DataMode.OPTIONAL); loader.clear(); for(QueryDataBatch result : results) { result.release(); } client.close(); }
@Test public void testDistributionBuckets() throws Exception { List<QueryDataBatch> result = testSqlWithResults("create table dfs_test.options_name DISTRIBUTE BY (name) STORE AS (type => 'TEXT', fieldDelimiter => ',') as select * from sys.options"); // Expect DISTRIBUTE BY to give multiple rows, i.e. multiple files int recordCount = 0; for (QueryDataBatch batch: result) { recordCount += batch.getHeader().getDef().getRecordCount(); batch.close(); } assertTrue(recordCount > 1); }
private void printSchema(String query) throws Exception { List<QueryDataBatch> results = testSqlWithResults(query); RecordBatchLoader loader = new RecordBatchLoader(getAllocator()); QueryDataBatch batch = results.get(0); loader.load(batch.getHeader().getDef(), batch.getData()); System.out.println(loader.getSchema()); VectorUtil.showVectorAccessibleContent(loader); } }
public void writeTo(io.protostuff.Output output, com.dremio.exec.proto.UserBitShared.QueryData message) throws java.io.IOException { if(message.hasQueryId()) output.writeObject(1, message.getQueryId(), com.dremio.exec.proto.SchemaUserBitShared.QueryId.WRITE, false); if(message.hasRowCount()) output.writeInt32(2, message.getRowCount(), false); if(message.hasDef()) output.writeObject(3, message.getDef(), com.dremio.exec.proto.SchemaUserBitShared.RecordBatchDef.WRITE, false); } public boolean isInitialized(com.dremio.exec.proto.UserBitShared.QueryData message)
@Override public QueryWritableBatch convertIfNecessary(QueryWritableBatch result) { if (result.getHeader().getDef().getRecordCount() > 0) { nonEmptyBatchSent = true; } return result; // does nothing by default } }