@Override public Object getObject(int index) throws InvalidAccessException { return v.getObject(index); }
private WritePartition getPartition(int offset){ String[] paths = new String[partitions.size()]; for(int i = 0; i < paths.length; i++){ paths[i] = fromObj(partitions.get(i).getObject(offset)); } final Integer bucketNumberValue = bucketNumber == null ? null : bucketNumber.get(offset); return new WritePartition(paths, bucketNumberValue); }
@Override public Object extractValue(String column, int index){ Pair<RecordBatchData, Integer> dataBatch = find(index); Integer columnIndex = nameToColumnIndex.get(column); return dataBatch.getKey().getVectors().get(columnIndex).getObject(dataBatch.getValue()); }
@Override public Object getObject(int index) { Map<String, Object> vv = new JsonStringHashMap<>(); for (String child : getChildFieldNames()) { ValueVector v = getChild(child); if (v != null && index < v.getValueCount()) { Object value = v.getObject(index); if (value != null) { vv.put(child, value); } } } return vv; }
private void unpivotHelper(PivotDef pivot, FixedBlockVector fbv, VariableBlockVector vbv, ValueVector[] in, ValueVector out[], int s, int e) throws Exception { Unpivots.unpivot(pivot, fbv, vbv, s, e); for (int i = 0; i < (e - s); i++) { for (int j = 0; j < in.length; j++) { assertEquals(in[j].getObject(s + i), out[j].getObject(i)); } } AutoCloseables.close(Arrays.asList(out)); }
/** * Get the element in the list vector at a particular index. * @param index position of the element * @return Object at given position */ @Override public Object getObject(int index) { if (isSet(index) == 0) { return null; } final List<Object> vals = new JsonStringArrayList<>(); final int start = offsetBuffer.getInt(index * OFFSET_WIDTH); final int end = offsetBuffer.getInt((index + 1) * OFFSET_WIDTH); final ValueVector vv = getDataVector(); for (int i = start; i < end; i++) { vals.add(vv.getObject(i)); } return vals; }
@Override public Object next() { if (currVec == null || indexInCurrentVector == currVec.getValueCount()) { currVec = hyperVector.getValueVectors()[indexInVectorList]; indexInVectorList++; indexInCurrentVector = 0; } Object obj = currVec.getObject(indexInCurrentVector); indexInCurrentVector++; totalValuesRead++; return obj; }
@SuppressWarnings("unchecked") private <T> void assertField(ValueVector valueVector, int index, TypeProtos.MinorType expectedMinorType, T value, String name, int parentFieldId) { if (expectedMinorType == TypeProtos.MinorType.STRUCT) { return; } final T val; try { val = (T) valueVector.getObject(index); } catch (Throwable ex) { throw ex; } if (val instanceof byte[]) { assertTrue(Arrays.equals((byte[]) value, (byte[]) val)); } else { assertEquals(value, val); } }
@Override public void apply(RecordBatchLoader loader) { assertEquals("invalid record count returned", count, loader.getRecordCount()); for (int r = 0; r < values.length; r++) { final Object[] row = values[r]; for (int c = 0; c<values[r].length; c++) { final Object expected = row[c]; final Object unconverted = loader.getValueAccessorById(ValueVector.class, c) .getValueVector().getObject(r); final Object actual = converter.convert(unconverted); assertEquals(String.format("row:%d - col:%d - expected:%s[%s] - actual:%s[%s]", r, c, expected, expected == null ? "null" : expected.getClass().getSimpleName(), actual, actual == null ? "null" : actual.getClass().getSimpleName()), actual, expected); } } } }
protected Object[] getRunResult(QueryType queryType, String planString) throws Exception { List<QueryDataBatch> resultList = testRunAndReturn(queryType, planString); List<Object> res = new ArrayList<>(); RecordBatchLoader loader = new RecordBatchLoader(getAllocator()); for(QueryDataBatch result : resultList) { if (result.getData() != null) { loader.load(result.getHeader().getDef(), result.getData()); ValueVector v = loader.iterator().next().getValueVector(); for (int j = 0; j < v.getValueCount(); j++) { if (v instanceof VarCharVector) { res.add(new String(((VarCharVector) v).get(j))); } else { res.add(v.getObject(j)); } } loader.clear(); result.release(); } } return res.toArray(); }
@Test public void testTimeStamp() throws Exception { try (ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { // run query. bit.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/record/vector/test_timestamp.json"), Charsets.UTF_8) .replace("#{TEST_FILE}", "/test_simple_date.json")); RecordBatchLoader batchLoader = new RecordBatchLoader(bit.getContext().getAllocator()); QueryDataBatch batch = results.get(0); assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData())); for (VectorWrapper<?> v : batchLoader) { ValueVector vv = v.getValueVector(); assertEquals(vv.getObject(0).toString() ,"1970-01-02 10:20:33.000"); assertEquals(vv.getObject(1).toString() ,"2008-12-28 11:34:00.129"); assertEquals(vv.getObject(2).toString(), "2000-02-27 14:24:00.000"); } batchLoader.clear(); for(QueryDataBatch b : results){ b.release(); } } }
@Test public void simpleEqualityJoin() throws Throwable { // Function checks hash join with single equality condition final String plan = Files.toString(FileUtils.getResourceAsFile("/join/hash_join.json"), Charsets.UTF_8) .replace("#{TEST_FILE_1}", FileUtils.getResourceAsFile("/build_side_input.json").toURI().toString()) .replace("#{TEST_FILE_2}", FileUtils.getResourceAsFile("/probe_side_input.json").toURI().toString()); List<QueryDataBatch> results = testRunAndReturn(QueryType.PHYSICAL, plan); try(RecordBatchLoader batchLoader = new RecordBatchLoader(nodes[0].getContext().getAllocator())){ QueryDataBatch batch = results.get(1); assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData())); Iterator<VectorWrapper<?>> itr = batchLoader.iterator(); // Just test the join key long colA[] = {1, 1, 2, 2, 1, 1}; // Check the output of decimal9 ValueVector intValueVector = itr.next().getValueVector(); for (int i = 0; i < intValueVector.getValueCount(); i++) { assertEquals(intValueVector.getObject(i), colA[i]); } assertEquals(6, intValueVector.getValueCount()); } for (QueryDataBatch result : results) { result.release(); } }
@SuppressWarnings("unchecked") public CellCompare compare(ValueVector vector, int index, boolean isValid) { V obj = isValid ? (V) vector.getObject(index) : null; if(obj == null && this.obj == null){ return new CellCompare(true, toString(obj)); } if(this.obj == null){ return new CellCompare(false, toString(obj) + " (null)"); } if(obj == null){ return new CellCompare(false, null + " ("+toString(this.obj)+")"); } if(!obj.getClass().equals(this.obj.getClass())){ return new CellCompare(false, obj.toString() + "(" + this.obj.toString() + ")"); } boolean isEqual = evaluateEquality(obj, this.obj); if(isEqual){ return new CellCompare(true, toString(this.obj)); }else{ return new CellCompare(false, toString(obj) + " ("+toString(this.obj)+")"); } }
protected Object[] getRunResult(QueryType queryType, String planString) throws Exception { List<QueryDataBatch> resultList = testRunAndReturn(queryType, planString); List<Object> res = new ArrayList<>(); try(RecordBatchLoader loader = new RecordBatchLoader(getAllocator())){; for(QueryDataBatch result : resultList) { if (result.getData() != null) { loader.load(result.getHeader().getDef(), result.getData()); ValueVector v = loader.iterator().next().getValueVector(); for (int j = 0; j < v.getValueCount(); j++) { res.add(v.getObject(j)); } loader.clear(); result.release(); } } } return res.toArray(); }
private void testCommon(String[] expectedResults, String physicalPlan, String resourceFile) throws Exception { try (ClusterCoordinator clusterCoordinator = LocalClusterCoordinator.newRunningCoordinator(); SabotNode bit = new SabotNode(DEFAULT_SABOT_CONFIG, clusterCoordinator, CLASSPATH_SCAN_RESULT, true); DremioClient client = new DremioClient(DEFAULT_SABOT_CONFIG, clusterCoordinator)) { // run query. bit.run(); client.connect(); List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile(physicalPlan), Charsets.UTF_8) .replace("#{TEST_FILE}", resourceFile)); try(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 vv = v.getValueVector(); System.out.println(vv.getObject(0)); assertEquals( expectedResults[i++], vv.getObject(0).toString()); } } for(QueryDataBatch b : results){ b.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 testMultiInputAdd() throws Exception { List<QueryDataBatch> results = client.runQuery(com.dremio.exec.proto.UserBitShared.QueryType.PHYSICAL, Files.toString(FileUtils.getResourceAsFile("/functions/multi_input_add_test.json"), Charsets.UTF_8)); try(RecordBatchLoader batchLoader = new RecordBatchLoader(nodes[0].getContext().getAllocator())){ QueryDataBatch batch = results.get(0); assertTrue(batchLoader.load(batch.getHeader().getDef(), batch.getData())); for (VectorWrapper<?> v : batchLoader) { ValueVector vv = v.getValueVector(); assertTrue((vv.getObject(0)).equals(10)); } batchLoader.clear(); for(QueryDataBatch b : results){ b.release(); } } } }
@Override public void dataArrived(QueryDataBatch result, ConnectionThrottle throttle) { QueryData queryHeader = result.getHeader(); int rows = queryHeader.getRowCount(); try { if ( result.hasData() ) { ArrowBuf data = result.getData(); loader.load(queryHeader.getDef(), data); for (int i = 0; i < rows; i++) { Map<String,String> record = Maps.newHashMap(); for (VectorWrapper<?> vw : loader) { final String field = vw.getValueVector().getField().getName(); final ValueVector vv = vw.getValueVector(); final Object value = i < vv.getValueCount() ? vv.getObject(i) : null; final String display = value == null ? null : value.toString(); record.put(field, display); } records.add(record); } loader.clear(); } result.release(); } catch (SchemaChangeException e) { fail(e.getMessage()); } }
@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(); }
@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(); }