Fix releasing arrow data

kuzudb · Sep 12, 2023 · 2dcc60c · 2dcc60c
1 parent 8520fcb
commit 2dcc60c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/common/arrow/arrow_converter.cpp b/src/common/arrow/arrow_converter.cpp
@@ -6,7 +6,7 @@ namespace kuzu {
 namespace common {
 
 static void releaseArrowSchema(ArrowSchema* schema) {
-    if (!schema || schema->release) {
+    if (!schema || !schema->release) {
         return;
     }
     schema->release = nullptr;

diff --git a/tools/python_api/src_cpp/py_query_result.cpp b/tools/python_api/src_cpp/py_query_result.cpp
@@ -214,8 +214,15 @@ kuzu::pyarrow::Table PyQueryResult::getAsArrow(std::int64_t chunkSize) {
 
     auto typesInfo = queryResult->getColumnTypesInfo();
     auto schema = ArrowConverter::toArrowSchema(typesInfo);
-    auto schemaObj = schemaImportFunc((std::uint64_t)schema.get());
+    // Prevent arrow from releasing the schema until it gets passed to the table
+    // It seems like you are expected to pass a new schema for each RecordBatch
+    auto release = schema->release;
+    schema->release = [](ArrowSchema*) {};
+
     py::list batches = getArrowChunks(*schema, chunkSize);
+    auto schemaObj = schemaImportFunc((std::uint64_t)schema.get());
+
+    schema->release = release;
     return py::cast<kuzu::pyarrow::Table>(fromBatchesFunc(batches, schemaObj));
 }