Refactor: add list of Arrow schema implementations to the data frame.

Author: anthony-tuininga

src/oracledb/arrow_impl.pxd

@@ -152,11 +152,12 @@ cdef class ArrowArrayImpl:
     cdef int get_uint(self, ArrowType arrow_type, int64_t index, bint* is_null,
                       uint64_t* value) except -1
     cdef object get_vector(self, int64_t index, bint* is_null)
-    cdef int populate_from_array(self, ArrowSchema* schema,
+    cdef int populate_from_array(self, ArrowSchemaImpl schema_impl,
                                  ArrowArray* array) except -1
     cdef int populate_from_schema(self, ArrowSchemaImpl schema_impl) except -1
 
 
 cdef class DataFrameImpl:
     cdef:
+        list schema_impls
         list arrays

src/oracledb/impl/arrow/array.pyx

@@ -484,27 +484,33 @@ cdef class ArrowArrayImpl:
         object implementing the PyCapsule Arrow array interface.
         """
         cdef:
+            ArrowSchemaImpl schema_impl
             ArrowArrayImpl array_impl
             ArrowSchema *arrow_schema
             ArrowArray *arrow_array
+
+        # convert schema
         schema_capsule, array_capsule = obj.__arrow_c_array__()
         arrow_schema = <ArrowSchema*> cpython.PyCapsule_GetPointer(
             schema_capsule, "arrow_schema"
         )
+        schema_impl = ArrowSchemaImpl.__new__(ArrowSchemaImpl)
+        schema_impl.populate_from_schema(arrow_schema)
+
+        # convert array
         arrow_array = <ArrowArray*> cpython.PyCapsule_GetPointer(
             array_capsule, "arrow_array"
         )
         array_impl = ArrowArrayImpl.__new__(ArrowArrayImpl)
-        array_impl.schema_impl = ArrowSchemaImpl.__new__(ArrowSchemaImpl)
-        array_impl.populate_from_array(arrow_schema, arrow_array)
+        array_impl.populate_from_array(schema_impl, arrow_array)
         return array_impl
 
-    cdef int populate_from_array(self, ArrowSchema* schema,
+    cdef int populate_from_array(self, ArrowSchemaImpl schema_impl,
                                  ArrowArray* array) except -1:
         """
         Populate the array from another array.
         """
-        self.schema_impl.populate_from_schema(schema)
+        self.schema_impl = schema_impl
         ArrowArrayMove(array, self.arrow_array)
 
     cdef int populate_from_schema(self, ArrowSchemaImpl schema_impl) except -1:

src/oracledb/impl/arrow/dataframe.pyx

@@ -38,24 +38,34 @@ cdef class DataFrameImpl:
         """
         cdef:
             ArrowArrayStream *arrow_stream
+            ArrowSchemaImpl schema_impl
+            ArrowArrayImpl array_impl
             ArrowSchema arrow_schema
             ArrowArray arrow_array
             DataFrameImpl df_impl
-            ArrowArrayImpl array_impl
-            ArrowSchemaImpl schema_impl
             ssize_t i
+
+        # initialization
         df_impl = DataFrameImpl.__new__(DataFrameImpl)
+        df_impl.schema_impls = []
         df_impl.arrays = []
         capsule = obj.__arrow_c_stream__()
         arrow_stream = <ArrowArrayStream*> cpython.PyCapsule_GetPointer(
             capsule, "arrow_array_stream"
         )
+
+        # populate list of schemas
         _check_nanoarrow(arrow_stream.get_schema(arrow_stream, &arrow_schema))
+        for i in range(arrow_schema.n_children):
+            schema_impl = ArrowSchemaImpl.__new__(ArrowSchemaImpl)
+            schema_impl.populate_from_schema(arrow_schema.children[i])
+            df_impl.schema_impls.append(schema_impl)
+
+        # populate list of arrays
         _check_nanoarrow(arrow_stream.get_next(arrow_stream, &arrow_array))
         for i in range(arrow_schema.n_children):
             array_impl = ArrowArrayImpl.__new__(ArrowArrayImpl)
-            array_impl.schema_impl = ArrowSchemaImpl.__new__(ArrowSchemaImpl)
-            array_impl.populate_from_array(arrow_schema.children[i],
+            array_impl.populate_from_array(df_impl.schema_impls[i],
                                            arrow_array.children[i])
             df_impl.arrays.append(array_impl)
         _check_nanoarrow(arrow_stream.get_next(arrow_stream, &arrow_array))
@@ -77,6 +87,7 @@ cdef class DataFrameImpl:
         encapsulates the arrays found in the data frame.
         """
         cdef:
+            ArrowSchemaImpl schema_impl
             ArrowArrayImpl array_impl
             ArrowArrayStream *stream
             int64_t i, num_arrays
@@ -100,14 +111,15 @@ cdef class DataFrameImpl:
                 ArrowArrayInitFromType(&array, NANOARROW_TYPE_STRUCT)
             )
             _check_nanoarrow(ArrowArrayAllocateChildren(&array, num_arrays))
-            for i, array_impl in enumerate(self.arrays):
+            for i, schema_impl in enumerate(self.schema_impls):
+                array_impl = self.arrays[i]
                 array.length = array_impl.arrow_array.length
                 copy_arrow_array(
                     array_impl, array_impl.arrow_array, array.children[i]
                 )
                 _check_nanoarrow(
                     ArrowSchemaDeepCopy(
-                        array_impl.schema_impl.arrow_schema, schema.children[i]
+                        schema_impl.arrow_schema, schema.children[i]
                     )
                 )
 

src/oracledb/impl/base/cursor.pyx

@@ -552,12 +552,16 @@ cdef class BaseCursorImpl:
         Flush all buffers and return an Oracle Data frame.
         """
         cdef:
+            ArrowArrayImpl array_impl
             DataFrameImpl df_impl
             BaseVarImpl var_impl
         df_impl = DataFrameImpl.__new__(DataFrameImpl)
+        df_impl.schema_impls = []
         df_impl.arrays = []
         for var_impl in self.fetch_var_impls:
-            df_impl.arrays.append(var_impl._finish_building_arrow_array())
+            array_impl = var_impl._finish_building_arrow_array()
+            df_impl.schema_impls.append(array_impl.schema_impl)
+            df_impl.arrays.append(array_impl)
         return PY_TYPE_DATAFRAME._from_impl(df_impl)
 
     def close(self, bint in_del=False):