feat: Add arrow serialization helpers

Author: yevgenypats

cloudquery/plugin_v3/arrow.py

@@ -0,0 +1,51 @@
+from pyarrow import ipc
+import pyarrow as pa
+from typing import List
+
+def schema_to_bytes(sc) -> bytes:
+    """
+    Convert a schema to bytes
+    """
+    sink = pa.BufferOutputStream()
+    try:
+      with pa.ipc.new_stream(sink, sc) as writer:
+          pass
+      return sink.getvalue().to_pybytes()
+    finally:
+      sink.close()
+
+def new_schema_from_bytes(buf: bytes) -> pa.Schema:
+    """
+    Convert bytes to schema
+    """
+    schema = None
+    with pa.ipc.open_stream(buf) as reader:
+      schema = reader.schema
+    return schema
+
+def schemas_to_bytes(schemas: List[pa.Schema]):
+  res : List[bytes] = []
+  for schema in schemas:
+    res.append(schema_to_bytes(schema))
+  return res
+
+def new_schemas_from_bytes(bufs: List[bytes]) -> List[pa.Schema]:
+  res : List[pa.Schema] = []
+  for buf in bufs:
+    res.append(new_schema_from_bytes(buf))
+  return res
+
+def record_to_bytes(rec: pa.RecordBatch) -> bytes:
+  sink = pa.BufferOutputStream()
+  try:
+    with pa.ipc.new_stream(sink, rec.schema) as writer:
+        writer.write_batch(rec)
+    return sink.getvalue().to_pybytes()
+  finally:
+    sink.close()
+
+def new_record_from_bytes(buf :bytes) -> pa.RecordBatch:
+  rec = None
+  with pa.ipc.open_stream(buf) as reader:
+    rec = reader.read_next_batch()
+  return rec

setup.py

@@ -14,6 +14,7 @@
     "grpcio >= 1.56.0",
     "grpcio-tools >= 1.56.0",
     "protobuf >= 4.23.4",
+    "pyarrow >= 12.0.1"
 ]
 url = "https://github.com/cloudquery/plugin-pb-python"
 
@@ -36,7 +37,7 @@
 ]
 setuptools.setup(
     name=name,
-    version="0.0.12",
+    version="0.0.13",
     description=description,
     long_description=long_description,
     author="CloudQuery LTD",

tests/__init__.py

@@ -0,0 +1,16 @@
+import pyarrow as pa
+from cloudquery.plugin_v3.arrow import schemas_to_bytes, new_schemas_from_bytes, record_to_bytes, new_record_from_bytes
+
+def test_schema_round_trip():
+    sc = pa.schema(fields=[pa.field("a", pa.int64())], metadata={"foo":"bar", "baz":"quux"})
+    b = schemas_to_bytes([sc])
+    schemas = new_schemas_from_bytes(b)
+    assert len(schemas) == 1
+    assert schemas[0].equals(sc)
+
+def test_record_round_trip():
+  sc = pa.schema(fields=[pa.field("a", pa.int64())], metadata={"foo":"bar", "baz":"quux"})
+  rec = pa.RecordBatch.from_arrays([pa.array([1,2,3])], schema=sc)
+  b = record_to_bytes(rec)
+  rec2 = new_record_from_bytes(b)
+  assert rec.equals(rec2)