Use ByteBuffer to receive data from Python

deephaven · Dec 14, 2023 · bb34d3b · bb34d3b
1 parent 482cd8e
commit bb34d3b
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