Use zero-copy C data interface in Parquet adapter

Signed-off-by: Arham Chopra <[email protected]>

Author: arhamchopra

cpp/csp/python/adapters/parquetadapterimpl.cpp

@@ -17,8 +17,10 @@
 #include <csp/python/PyCppNode.h>
 #include <csp/python/PyNodeWrapper.h>
 #include <csp/python/NumpyConversions.h>
+#include <arrow/c/abi.h>
+#include <arrow/c/bridge.h>
 #include <arrow/io/memory.h>
-#include <arrow/ipc/reader.h>
+#include <arrow/table.h>
 #include <locale>
 #include <codecvt>
 
@@ -156,34 +158,30 @@ class ArrowTableGenerator : public csp::Generator<std::shared_ptr<arrow::Table>,
         {
             CSP_THROW( csp::python::PythonPassthrough, "" );
         }
-        if( nextVal == nullptr )
+        if( nextValPtr.get() == nullptr )
         {
             return false;
         }
 
-        if(!PyBytes_Check( nextVal ))
+        if( !PyCapsule_IsValid( nextValPtr.get(), "arrow_array_stream" ) )
         {
-            CSP_THROW( csp::TypeError, "Invalid arrow buffer type, expected bytes got " << Py_TYPE( nextVal ) -> tp_name );
+            CSP_THROW( csp::TypeError, "Invalid arrow data, expected PyCapsule got " << Py_TYPE( nextValPtr.get() ) -> tp_name );
         }
-        const char * data = PyBytes_AsString( nextVal );
-        if( !data )
-            CSP_THROW( csp::python::PythonPassthrough, "" );
-        auto size = PyBytes_Size(nextVal);
-        m_data = csp::python::PyObjectPtr::incref(nextVal);
-        std::shared_ptr<arrow::io::BufferReader> bufferReader = std::make_shared<arrow::io::BufferReader>(
-                reinterpret_cast<const uint8_t *>(data), size );
-        std::shared_ptr<arrow::ipc::RecordBatchStreamReader> reader = arrow::ipc::RecordBatchStreamReader::Open(bufferReader.get()).ValueOrDie();
-        auto result = reader->ToTable();
-        if (!(result.ok()))
-            CSP_THROW(csp::RuntimeException, "Failed read arrow table from buffer");
-        value = std::move(result.ValueUnsafe());
+        // Extract the record batch
+        struct ArrowArrayStream * c_stream = reinterpret_cast<struct ArrowArrayStream*>( PyCapsule_GetPointer( nextValPtr.get(), "arrow_array_stream" ) );
+        auto reader_result = arrow::ImportRecordBatchReader( c_stream );
+        if( !reader_result.ok() )
+            CSP_THROW( csp::ValueError, "Failed to load record batches through PyCapsule C Data interface: " << reader_result.status().ToString() );
+        auto reader = std::move( reader_result.ValueUnsafe() );
+        auto table_result = arrow::Table::FromRecordBatchReader( reader.get() );
+        if( !table_result.ok() )
+            CSP_THROW( csp::ValueError, "Failed to load table from record batches " << table_result.status().ToString() );
+        value = std::move( table_result.ValueUnsafe() );
         return true;
     }
 private:
     csp::python::PyObjectPtr m_wrappedGenerator;
     csp::python::PyObjectPtr m_iter;
-    // We need to keep the last buffer in memory since arrow doesn't copy it but can refer to strings in it
-    csp::python::PyObjectPtr m_data;
 };
 
 template< typename CspCType>

csp/adapters/parquet.py

@@ -88,9 +88,7 @@ def __init__(
             if not _CAN_READ_ARROW_BINARY:
                 raise TypeError("CSP Cannot load binary arrows derived from pyarrow versions less than 4.0.1")
             wrapped = self._filenames_gen
-            self._filenames_gen = lambda starttime, endtime: self._arrow_in_memory_table_to_buffers(
-                wrapped, starttime, endtime
-            )
+            self._filenames_gen = lambda starttime, endtime: self._arrow_c_data_interface(wrapped, starttime, endtime)
             binary_arrow = True
         self._properties = {"split_columns_to_files": split_columns_to_files}
         if symbol_column:
@@ -116,22 +114,12 @@ def __init__(
         self._properties["allow_missing_files"] = allow_missing_files
 
     @classmethod
-    def _arrow_in_memory_table_to_buffers(cls, gen, startime, endtime):
-        # This is a temporary solution until we implement PyArrow in-memory suport:
-        # Currently when we try to read in c++ pyarrow structures that were created in python it crashes. We believe that the reason is difference in compilers. c++ interface is generally not portable across compilers and compiler versions. As current workaround we create a temporary files in python from in-memory tables and consume those in c++ code, this is bad and ugly apporoach.
-        # Possible solutions:
-        # 1. Try using pyarrow c interface (introduced in pyarrow 5.0.0), it doesn't seem like it's sufficient for our needs but we could try using it.
-        # 2. Compile the arrow adapter in the user environment - generally it's hard to do and is a huge opening to a bunch of issues, we should probably avoid this approach as much as possible.
-        # 3. The most "sane" solution is probably to use numpy array as intermediate layer, i.e transform pyarrow tables to arrow arrays (and pass a bunch of metadata about those arrays - to properly resolve typing). We probably won't be able to support ALL arrow types, we will limit ourselves to a reasonable set of types for which we will have to implement conversions.
+    def _arrow_c_data_interface(cls, gen, startime, endtime):
         for v in gen(startime, endtime):
             if not isinstance(v, pyarrow.Table):
                 raise TypeError(f"Expected PyTable from generator, got {type(v).__name__}")
-
-            sink = io.BytesIO()
-            with pyarrow.ipc.new_stream(sink, v.schema) as writer:
-                writer.write_table(v)
-
-            yield sink.getvalue()
+            # Use the PyCapsule C data interface to pass data zero copy
+            yield v.__arrow_c_stream__()
 
     @node
     def _reconstruct_struct_array_fields(self, s: ts["T"], fields: {str: ts[object]}, struct_typ: "T") -> ts["T"]: