From 6f79d5e943b3af5ba38e9e8554101489da5fb4f1 Mon Sep 17 00:00:00 2001 From: jianfengmao Date: Thu, 14 Dec 2023 09:50:03 -0700 Subject: [PATCH] Batch load arrow record batches --- .../extensions/barrage/util/ArrowToTableConverter.java | 7 +++++++ py/server/deephaven/arrow.py | 4 +--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/ArrowToTableConverter.java b/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/ArrowToTableConverter.java index 12063e81b6e..8d23a4dc7ef 100644 --- a/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/ArrowToTableConverter.java +++ b/extensions/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 index 4b2cb6b90b7..b61e62debb1 100644 --- 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())