Use ByteBuffer to receive data from Python

deephaven · Dec 11, 2023 · c649758 · c649758
1 parent 1d255d3
commit c649758
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 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
@@ -48,10 +48,9 @@ public class ArrowToTableConverter {
 
     private volatile boolean completed = false;
 
-    private static BarrageProtoUtil.MessageInfo parseArrowIpcMessage(final byte[] ipcMessage) throws IOException {
+    private BarrageProtoUtil.MessageInfo parseArrowIpcMessage(ByteBuffer bb) throws IOException {
         final BarrageProtoUtil.MessageInfo mi = new BarrageProtoUtil.MessageInfo();
 
-        final ByteBuffer bb = ByteBuffer.wrap(ipcMessage);
         bb.order(ByteOrder.LITTLE_ENDIAN);
         final int continuation = bb.getInt();
         final int metadata_size = bb.getInt();
@@ -70,7 +69,7 @@ private static BarrageProtoUtil.MessageInfo parseArrowIpcMessage(final byte[] ip
     }
 
     @ScriptApi
-    public synchronized void setSchema(final byte[] ipcMessage) {
+    public synchronized void setSchema(final ByteBuffer ipcMessage) {
         if (completed) {
             throw new IllegalStateException("Conversion is complete; cannot process additional messages");
         }
@@ -82,7 +81,7 @@ public synchronized void setSchema(final byte[] ipcMessage) {
     }
 
     @ScriptApi
-    public synchronized void addRecordBatch(final byte[] ipcMessage) {
+    public synchronized void addRecordBatch(final ByteBuffer ipcMessage) {
         if (completed) {
             throw new IllegalStateException("Conversion is complete; cannot process additional messages");
         }
@@ -192,7 +191,7 @@ protected BarrageMessage createBarrageMessage(BarrageProtoUtil.MessageInfo mi, i
         return msg;
     }
 
-    private BarrageProtoUtil.MessageInfo getMessageInfo(byte[] ipcMessage) {
+    private BarrageProtoUtil.MessageInfo getMessageInfo(ByteBuffer ipcMessage) {
         final BarrageProtoUtil.MessageInfo mi;
         try {
             mi = parseArrowIpcMessage(ipcMessage);
@@ -201,4 +200,6 @@ private BarrageProtoUtil.MessageInfo getMessageInfo(byte[] ipcMessage) {
         }
         return mi;
     }
+
+
 }
diff --git a/py/server/deephaven/arrow.py b/py/server/deephaven/arrow.py
@@ -101,12 +101,12 @@ def to_table(pa_table: pa.Table, cols: List[str] = None) -> Table:
 
     try:
         pa_buffer = dh_schema.serialize()
-        j_barrage_table_builder.setSchema(dtypes.array(dtypes.byte, pa_buffer))
+        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(dtypes.array(dtypes.byte, pa_buffer))
+            j_barrage_table_builder.addRecordBatch(pa_buffer)
         j_barrage_table_builder.onCompleted()
 
         return Table(j_table=j_barrage_table_builder.getResultTable())

diff --git a/py/server/tests/test_arrow.py b/py/server/tests/test_arrow.py
@@ -18,6 +18,7 @@
 from deephaven.table import Table
 from tests.testbase import BaseTestCase
 
+
 class ArrowTestCase(BaseTestCase):
     test_table: Table