Support empty Arrow files

Signed-off-by: Andrew Stein <[email protected]>
tarekquao · Nov 18, 2024 · 3e772ec · 3e772ec
1 parent f48e67b
commit 3e772ec
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/cpp/perspective/src/cpp/arrow_loader.cpp b/cpp/perspective/src/cpp/arrow_loader.cpp
@@ -77,10 +77,10 @@ load_file(
     } else {
         std::shared_ptr<arrow::ipc::RecordBatchFileReader> batch_reader =
             *status;
+
         std::vector<std::shared_ptr<arrow::RecordBatch>> batches;
         auto num_batches = batch_reader->num_record_batches();
         for (int i = 0; i < num_batches; ++i) {
-
             auto status2 = batch_reader->ReadRecordBatch(i);
             if (!status2.ok()) {
                 PSP_COMPLAIN_AND_ABORT(
@@ -91,7 +91,10 @@ load_file(
             std::shared_ptr<arrow::RecordBatch> chunk = *status2;
             batches.push_back(chunk);
         }
-        auto status3 = arrow::Table::FromRecordBatches(batches);
+
+        auto status3 =
+            arrow::Table::FromRecordBatches(batch_reader->schema(), batches);
+
         if (!status3.ok()) {
             std::stringstream ss;
             ss << "Failed to create Table from RecordBatches: "

diff --git a/rust/perspective-python/perspective/tests/table/test_table_arrow.py b/rust/perspective-python/perspective/tests/table/test_table_arrow.py
@@ -117,6 +117,27 @@ def test_table_arrow_loads_int_stream(self, util):
             "d": data[3],
         }
 
+    def test_empty_arrow(self, util):
+        table = pa.table(
+            {
+                "col1": [1, 2, 3],
+                "col2": ["abc", "foo", "bar"],
+            }
+        )
+
+        empty_table = table.schema.empty_table()
+        assert client.table(table, name="table").size() == 3
+        assert client.table(empty_table, name="table_empty_bad").size() == 0
+        assert client.table(table, name="table").schema() == {
+            "col1": "integer",
+            "col2": "string",
+        }
+
+        assert client.table(empty_table, name="table").schema() == {
+            "col1": "integer",
+            "col2": "string",
+        }
+
     def test_table_arrow_loads_float_stream(self, util):
         data = [[i for i in range(10)], [i * 1.5 for i in range(10)]]
         arrow_data = util.make_arrow(["a", "b"], data)