minor streamwriter changes to allow compression (#55)

Author: mlathara

src/genomicsdb.pyx

@@ -199,14 +199,15 @@ cdef class _GenomicsDB:
                             flatten_intervals=False,
                             json_output=None,
                             arrow_output=None,
-                            # batching only used with arrow_output
-                            batching=False):
+                            # batching/compress only used with arrow_output
+                            batching=False,
+                            compress=None):
         """ Query for variant calls from the GenomicsDB workspace using array, column_ranges and row_ranges for subsetting """
 
         if json_output is not None:
             return self.query_variant_calls_json(array, column_ranges, row_ranges, query_protobuf, json_output);
         elif arrow_output is not None:
-            return self.query_variant_calls_arrow(array, column_ranges, row_ranges, query_protobuf, batching);
+            return self.query_variant_calls_arrow(array, column_ranges, row_ranges, query_protobuf, batching, compress);
         elif flatten_intervals is True:
             return self.query_variant_calls_columnar(array, column_ranges, row_ranges, query_protobuf)
         else:
@@ -318,7 +319,8 @@ cdef class _GenomicsDB:
                                   column_ranges=None,
                                   row_ranges=None,
                                   query_protobuf: query_pb.QueryConfiguration=None,
-                                  batching=False): 
+                                  batching=False,
+                                  compress=None): 
       """ Query for variant calls from the GenomicsDB workspace using array, column_ranges and row_ranges for subsetting """
 
       cdef ArrowVariantCallProcessor processor
@@ -366,19 +368,24 @@ cdef class _GenomicsDB:
         schema_capsule = pycapsule_get_arrow_schema(arrow_schema)
         schema_obj = _ArrowSchemaWrapper._import_from_c_capsule(schema_capsule)
         schema = pa.schema(schema_obj.children_schema)
-        yield schema.serialize().to_pybytes()
       else:
         raise GenomicsDBException("Failed to retrieve arrow schema for query_variant_calls()")
 
       cdef void* arrow_array = NULL
+      w_opts = pa.ipc.IpcWriteOptions(allow_64bit=True, compression=compress)
       while True:
         try:
           arrow_array = processor.arrow_array()
           if arrow_array:
             array_capsule = pycapsule_get_arrow_array(arrow_array)
             array_obj = _ArrowArrayWrapper._import_from_c_capsule(schema_capsule, array_capsule)
             arrays = [pa.array(array_obj.child(i)) for i in range(array_obj.n_children)]
-            yield pa.RecordBatch.from_arrays(arrays, schema=schema).serialize().to_pybytes()
+            batch = pa.RecordBatch.from_arrays(arrays, schema=schema)
+            sink = pa.BufferOutputStream()
+            writer = pa.RecordBatchStreamWriter(sink, schema, options=w_opts)
+            writer.write_batch(batch)
+            writer.close()
+            yield sink.getvalue().to_pybytes()
           else:
             break
         except Exception as e:

test/test_genomicsdb.py

@@ -78,29 +78,21 @@ def test_connect_and_query_with_protobuf(setup):
                                             json_output=9999)
 
   # test with query protobuf and arrow output
-  first = True
   schema = pa.schema([("null_field", pa.string())])
   for output in gdb.query_variant_calls(row_ranges=[(0, 3)], array="t0_1_2", arrow_output=True):
-    if first:
-      schema = pa.ipc.read_schema(pa.py_buffer(output))
-      first = False
-    else:
-      batch = pa.ipc.read_record_batch(pa.py_buffer(output), schema)
-      assert batch.num_columns == 6
-      assert batch.num_rows == 5
+    reader = pa.ipc.open_stream(output)
+    batch = reader.read_next_batch()
+    assert batch.num_columns == 6
+    assert batch.num_rows == 5
 
-  first = True
   batch = None
   for output in gdb.query_variant_calls(
       row_ranges=[(0, 3)], array="t0_1_2", arrow_output=True, batching=True
   ):
-    if first:
-      schema = pa.ipc.read_schema(pa.py_buffer(output))
-      first = False
-    else:
-      batch = pa.ipc.read_record_batch(pa.py_buffer(output), schema)
-      assert batch.num_columns == 6
-      assert batch.num_rows == 1 or batch.num_rows == 3
+    reader = pa.ipc.open_stream(output)
+    batch = reader.read_next_batch()
+    assert batch.num_columns == 6
+    assert batch.num_rows == 1 or batch.num_rows == 3
 
   # test with query contig interval and no results
   interval = query_coords.ContigInterval()

test/test_genomicsdb_demo.py

@@ -111,14 +111,10 @@ def test_genomicsdb_demo_with_arrow_output(self):
       start = time.time()
       gdb = genomicsdb.connect_with_protobuf(self.query_config)
       print("\nSummary for batching mode=" + str(batching_mode) + ":")
-      first = True
       for output in gdb.query_variant_calls(arrow_output=True, batching=batching_mode):
-        if first:
-          schema = pa.ipc.read_schema(pa.py_buffer(output))
-          first = False
-        else:
-          batch = pa.ipc.read_record_batch(pa.py_buffer(output), schema)
-          print("batch num_rows="+str(batch.num_rows)+" num_columns="+str(batch.num_columns))
+        reader = pa.ipc.open_stream(output)
+        batch = reader.read_next_batch()
+        print("batch num_rows="+str(batch.num_rows)+" num_columns="+str(batch.num_columns))
       print("\tElapsed time: " + str(time.time() - start))
 
 if __name__ == '__main__':