/** * Ensure the reader has been initialized and reset the VectorSchemaRoot row count to 0. * * @throws IOException on error */ protected void prepareLoadNextBatch() throws IOException { ensureInitialized(); root.setRowCount(0); }
public void execute(int recordCount, List<ValueVector> outVectors) throws Exception { root.setRowCount(recordCount); List<ArrowBuf> buffers = Lists.newArrayList(); for (FieldVector v : root.getFieldVectors()) { buffers.addAll(v.getFieldBuffers()); } projector.evaluate(recordCount, buffers, outVectors); }
/** * Loads the record batch in the vectors. * will not close the record batch * * @param recordBatch the batch to load */ public void load(ArrowRecordBatch recordBatch) { Iterator<ArrowBuf> buffers = recordBatch.getBuffers().iterator(); Iterator<ArrowFieldNode> nodes = recordBatch.getNodes().iterator(); for (FieldVector fieldVector : root.getFieldVectors()) { loadBuffers(fieldVector, fieldVector.getField(), buffers, nodes); } root.setRowCount(recordBatch.getLength()); if (nodes.hasNext() || buffers.hasNext()) { throw new IllegalArgumentException("not all nodes and buffers were consumed. nodes: " + Collections2.toList(nodes).toString() + " buffers: " + Collections2.toList(buffers).toString()); } }
public boolean read(VectorSchemaRoot root) throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { { int count = readNextField("count", Integer.class); root.setRowCount(count); nextFieldIs("columns"); readToken(START_ARRAY); { for (Field field : root.getSchema().getFields()) { FieldVector vector = root.getVector(field.getName()); readFromJsonIntoVector(field, vector); } } readToken(END_ARRAY); } readToken(END_OBJECT); return true; } else if (t == END_ARRAY) { root.setRowCount(0); return false; } else { throw new IllegalArgumentException("Invalid token: " + t); } }
/** * Filter a batch of records against the expression. * @param recordCount - number of records to consume * @return the number of records that passed the filter * @throws GandivaException on evaluation exception. */ public int filterBatch(int recordCount) throws GandivaException { if (recordCount == 0) { return 0; } root.setRowCount(recordCount); List<ArrowBuf> buffers = Lists.newArrayList(); for (FieldVector v : root.getFieldVectors()) { buffers.addAll(v.getFieldBuffers()); } selectionVector.allocateNew(recordCount); // do not take ownership of the buffer. ArrowBuf svBuffer = selectionVector.getBuffer(false); SelectionVector selectionVectorGandiva = new SelectionVectorInt16(svBuffer); filter.evaluate(recordCount, buffers, selectionVectorGandiva); selectionVector.setRecordCount(selectionVectorGandiva.getRecordCount()); return selectionVector.getCount(); }
public VectorSchemaRoot read() throws IOException { JsonToken t = parser.nextToken(); if (t == START_OBJECT) { VectorSchemaRoot recordBatch = VectorSchemaRoot.create(schema, allocator); { int count = readNextField("count", Integer.class); recordBatch.setRowCount(count); nextFieldIs("columns"); readToken(START_ARRAY); { for (Field field : schema.getFields()) { FieldVector vector = recordBatch.getVector(field.getName()); readFromJsonIntoVector(field, vector); } } readToken(END_ARRAY); } readToken(END_OBJECT); return recordBatch; } else if (t == END_ARRAY) { return null; } else { throw new IllegalArgumentException("Invalid token: " + t); } }