Write Arrow Table/RecordBatchReader to GDAL

pyogrio/_io.pyx

@@ -783,7 +783,7 @@ cdef process_fields(
             data[i] = bin_value[:ret_length]
 
         elif field_type == OFTDateTime or field_type == OFTDate:
-            
+
             if datetime_as_string:
                 # defer datetime parsing to user/ pandas layer
                 # Update to OGR_F_GetFieldAsISO8601DateTime when GDAL 3.7+ only
@@ -851,7 +851,7 @@ cdef get_features(
 
     field_data = [
         np.empty(shape=(num_features, ),
-        dtype = ("object" if datetime_as_string and 
+        dtype = ("object" if datetime_as_string and
                     fields[field_index,3].startswith("datetime") else fields[field_index,3])
         ) for field_index in range(n_fields)
     ]
@@ -950,8 +950,8 @@ cdef get_features_by_fid(
     field_ogr_types = fields[:,1]
     field_data = [
         np.empty(shape=(count, ),
-        dtype=("object" if datetime_as_string and fields[field_index,3].startswith("datetime") 
-            else fields[field_index,3])) 
+        dtype=("object" if datetime_as_string and fields[field_index,3].startswith("datetime")
+            else fields[field_index,3]))
         for field_index in range(n_fields)
     ]
 
@@ -1678,6 +1678,117 @@ cdef infer_field_types(list dtypes):
 
     return field_types
 
+def ogr_write_arrow(
+    obj,
+    str path
+):
+    IF CTE_GDAL_VERSION < (3, 8, 0):
+        raise RuntimeError("Need GDAL>=3.8 for Arrow write support")
+
+    stream_capsule = obj.__arrow_c_stream__()
+    return write_arrow_stream_capsule(stream_capsule)
+
+cdef OGRErr write_arrow_stream_capsule(OGRLayerH destLayer, object capsule):
+    cdef ArrowArrayStream stream
+    cdef ArrowSchema schema
+    cdef ArrowArray array
+
+    stream = PyCapsule_GetPointer(capsule, "arrow_array_stream")
+    if not stream:
+        raise RuntimeError("Accessing PyCapsule pointer named 'arrow_array_stream' failed.")
+
+    if stream.release == NULL:
+        raise RuntimeError("Arrow Array Stream was already released.")
+
+    errcode = stream.get_schema(stream, &schema)
+    if errcode != 0:
+        stream.release(stream)
+        raise RuntimeError("Error while accessing schema from stream.")
+
+    errcode = create_fields_from_arrow_schema(destLayer, &schema)
+    if errcode != 0:
+        schema.release(&schema)
+        stream.release(stream)
+        raise RuntimeError("Error creating Arrow Schema in OGR layer.")
+
+    while True:
+        errcode = stream.get_next(stream, &array)
+        if errcode != 1:
+            schema.release(&schema)
+            stream.release(stream)
+            raise RuntimeError("Error while accessing batch from stream.")
+
+        # We've reached the end of the stream
+        if array.release == NULL:
+            break
+
+        errcode = OGR_L_WriteArrowBatch(destLayer, &schema, &array, options)
+        if errcode:
+            if array.release != NULL:
+                array.release(&array)
+
+            schema.release(&schema)
+            stream.release(stream)
+            raise RuntimeError("Error while writing batch to OGR layer.")
+
+        if array.release != NULL:
+            array.release(&array)
+
+
+    schema.release(&schema)
+    stream.release(stream)
+
+# Create output fields using CreateFieldFromArrowSchema()
+static bool create_fields_from_arrow_schema(
+    OGRLayerH destLayer,
+    const struct ArrowSchema* schema,
+    char** options
+):
+    # The schema object is a struct type where each child is a column.
+    for child in schema.n_children:
+        # Access the metadata for this column
+        const char *metadata = child.metadata
+
+        # TODO: I don't know how to parse this metadata in C... I guess I can just use Python APIs for this in Cython?
+        # https://github.com/OSGeo/gdal/pull/9133/files#diff-37bedc92ae1d5e04706c7b9f8ea9e9fcccf984ca0c9997e2020ff85f1b958433R1159-R1185
+        # if metadata:
+
+        field_name = child.name
+
+
+# {
+#     for (int i = 0; i < schemaSrc->n_children; ++i)
+#     {
+#         const char *metadata =
+#             schemaSrc->children[i]->metadata;
+#         if( metadata )
+#         {
+#             char** keyValues = ParseArrowMetadata(metadata);
+#             const char *ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name";
+#             const char *EXTENSION_NAME_OGC_WKB = "ogc.wkb";
+#             const char *EXTENSION_NAME_GEOARROW_WKB = "geoarrow.wkb";
+#             const char* value = CSLFetchNameValue(keyValues, ARROW_EXTENSION_NAME_KEY);
+#             const bool bSkip = ( value && (EQUAL(value, EXTENSION_NAME_OGC_WKB) || EQUAL(value, EXTENSION_NAME_GEOARROW_WKB)) );
+#             CSLDestroy(keyValues);
+#             if( bSkip )
+#                 continue;
+#         }
+
+#         const char *pszFieldName =
+#             schemaSrc->children[i]->name;
+#         if (!EQUAL(pszFieldName, "OGC_FID") &&
+#             !EQUAL(pszFieldName, "wkb_geometry") &&
+#             !OGR_L_CreateFieldFromArrowSchema(
+#                 hDstLayer, schemaSrc->children[i], options))
+#         {
+#             CPLError(CE_Failure, CPLE_AppDefined,
+#                      "Cannot create field %s",
+#                      pszFieldName);
+#             return false;
+#         }
+#     }
+#     return true;
+# }
 
 # TODO: set geometry and field data as memory views?
 def ogr_write(
@@ -2077,4 +2188,4 @@ def ogr_write(
         # on close
         exc = exc_check()
         if exc:
-            raise DataSourceError(f"Failed to write features to dataset {path}; {exc}")
+            raise DataSourceError(f"Failed to write features to dataset {path}; {exc}")

pyogrio/_ogr.pxd

@@ -191,11 +191,20 @@ cdef extern from "ogr_srs_api.h":
 
 
 cdef extern from "arrow_bridge.h":
+    struct ArrowArray:
+        void (*release)(struct ArrowArray*)
+
     struct ArrowSchema:
+        const char* name
+        const char* metadata
+        struct ArrowSchema** children
         int64_t n_children
 
     struct ArrowArrayStream:
         int (*get_schema)(ArrowArrayStream* stream, ArrowSchema* out)
+        int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out);
+        const char* (*get_last_error)(struct ArrowArrayStream*);
+        void (*release)(struct ArrowArrayStream*);
 
 
 cdef extern from "ogr_api.h":
@@ -385,4 +394,4 @@ cdef extern from "gdal.h":
     const char*     GDALVersionInfo(const char *pszRequest)
 
 
-cdef get_string(const char *c_str, str encoding=*)
+cdef get_string(const char *c_str, str encoding=*)