Java Code Examples for org.apache.arrow.vector.ipc.message.ArrowRecordBatch#close()

/**
 * Attempts to deserialize the provided byte[] into an ArrowRecordBatch.
 *
 * @param in The byte[] that is expected to contain a serialized ArrowRecordBatch.
 * @return The resulting ArrowRecordBatch if the byte[] contains a valid ArrowRecordBatch.
 * @throws IOException
 */
public ArrowRecordBatch deserialize(byte[] in)
        throws IOException
{
    ArrowRecordBatch batch = null;
    try {
        return allocator.registerBatch((BufferAllocator root) ->
                (ArrowRecordBatch) MessageSerializer.deserializeMessageBatch(
                        new ReadChannel(Channels.newChannel(new ByteArrayInputStream(in))),
                        root)
        );
    }
    catch (Exception ex) {
        if (batch != null) {
            batch.close();
        }
        throw ex;
    }
}
 

public boolean drain() throws InterruptedException {
  if (isDone.get()) {
    logger.debug("already done, returning. {}", descriptor);
    return true;
  }
  logger.debug("it is not done. try to take from exchanger for {}", descriptor);
  ArrowRecordBatch batch = exchanger.poll(1, TimeUnit.SECONDS);
  if (batch == null) {
    logger.debug("timed out waiting, {}", descriptor);
    return false;
  }
  logger.debug("got batch, send to loader. exchanger size is {} for {}", exchanger.size(), descriptor);
  if (batch.getLength() == -1) {
    logger.debug("{} is done, stopping it", descriptor);
    return true;
  }
  loader.load(batch);
  logger.debug("tell listener to put next for {}", descriptor);
  listener.putNext();
  batch.close();
  logger.debug("got a batch and sent to listener for {}", descriptor);
  return false;
}
 

@Override
public void run(SourceContext<OUT> ctx) throws Exception {
	VectorLoader vectorLoader = new VectorLoader(root);
	while (running && !indexesToEmit.isEmpty()) {
		Tuple2<Integer, Integer> indexToEmit = indexesToEmit.peek();
		ArrowRecordBatch arrowRecordBatch = loadBatch(indexToEmit.f0);
		vectorLoader.load(arrowRecordBatch);
		arrowRecordBatch.close();

		ArrowReader<OUT> arrowReader = createArrowReader(root);
		int rowCount = root.getRowCount();
		int nextRowId = indexToEmit.f1;
		while (nextRowId < rowCount) {
			OUT element = arrowReader.read(nextRowId);
			synchronized (ctx.getCheckpointLock()) {
				ctx.collect(element);
				indexToEmit.setField(++nextRowId, 1);
			}
		}

		synchronized (ctx.getCheckpointLock()) {
			indexesToEmit.pop();
		}
	}
}
 

/**
 * Attempts to close all batches allocated by this BlockAllocator.
 */
@VisibleForTesting
protected synchronized void closeBatches()
{
    logger.debug("closeBatches: {}", recordBatches.size());
    for (ArrowRecordBatch next : recordBatches) {
        try {
            next.close();
        }
        catch (Exception ex) {
            logger.warn("closeBatches: Error closing batch", ex);
        }
    }
    recordBatches.clear();
}
 

/**
 * Serialized the provided ArrowRecordBatch to the provided OutputStream and closes the batch once
 * it is fully written to the OutputStream.
 *
 * @param batch The ArrowRecordBatch to serialize.
 * @param out The OutputStream to write to.
 * @throws IOException
 */
public void serialize(ArrowRecordBatch batch, OutputStream out)
        throws IOException
{
    try {
        MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), batch);
    }
    finally {
        batch.close();
    }
}
 

private byte[] serializeRecordBatch(ArrowRecordBatch recordBatch)
        throws IOException
{
    try {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), recordBatch);
        return out.toByteArray();
    }
    finally {
        recordBatch.close();
    }
}
 

@Override
protected void doSerialize(ArrowRecordBatch arrowRecordBatch, JsonGenerator jgen, SerializerProvider provider)
        throws IOException
{
    try {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowRecordBatch);
        jgen.writeBinary(out.toByteArray());
    }
    finally {
        arrowRecordBatch.close();
    }
}
 

/**
 * Used to load Apache Arrow data into this Block after it has been deserialized.
 *
 * @param batch An ArrowRecordBatch containing all row data you'd like to load into this Block.
 * @note The batch is closed after being loaded to avoid memory leaks or data corruption since the buffers
 * associated with the batch are now owned by this Block. Closing the batch essentially decrements the referrence
 * count in the Arrow Allocator.
 */
public void loadRecordBatch(ArrowRecordBatch batch)
{
    VectorLoader vectorLoader = new VectorLoader(vectorSchema);
    vectorLoader.load(batch);
    batch.close();
}