Batch load arrow record batches

deephaven · Dec 14, 2023 · 6f79d5e · 6f79d5e
1 parent bb34d3b
commit 6f79d5e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/...ons/barrage/src/main/java/io/deephaven/extensions/barrage/util/ArrowToTableConverter.java b/...ons/barrage/src/main/java/io/deephaven/extensions/barrage/util/ArrowToTableConverter.java
@@ -80,6 +80,13 @@ public synchronized void setSchema(final ByteBuffer ipcMessage) {
  parseSchema((Schema) mi.header.header(new Schema()));
  }
 
+ @ScriptApi
+ public synchronized void addRecordBatches(final ByteBuffer... ipcMessages) {
+ for (final ByteBuffer ipcMessage : ipcMessages) {
+ addRecordBatch(ipcMessage);
+ }
+ }
+
  @ScriptApi
  public synchronized void addRecordBatch(final ByteBuffer ipcMessage) {
  if (completed) {

diff --git a/py/server/deephaven/arrow.py b/py/server/deephaven/arrow.py
@@ -104,9 +104,7 @@ def to_table(pa_table: pa.Table, cols: List[str] = None) -> Table:
  j_barrage_table_builder.setSchema(pa_buffer)
 
  record_batches = pa_table.to_batches()
- for rb in record_batches:
- pa_buffer = rb.serialize()
- j_barrage_table_builder.addRecordBatch(pa_buffer)
+ j_barrage_table_builder.addRecordBatches([jpy.byte_buffer(rb.serialize()) for rb in record_batches])
  j_barrage_table_builder.onCompleted()
 
  return Table(j_table=j_barrage_table_builder.getResultTable())