feat(python/sedonadb): Implement GDAL/OGR formats

python/sedonadb/Cargo.toml

@@ -43,6 +43,7 @@ pyo3 = { version = "0.25.1" }
 sedona = { path = "../../rust/sedona" }
 sedona-adbc = { path = "../../rust/sedona-adbc" }
 sedona-expr = { path = "../../rust/sedona-expr" }
+sedona-datasource = { path = "../../rust/sedona-datasource" }
 sedona-geoparquet = { path = "../../rust/sedona-geoparquet" }
 sedona-schema = { path = "../../rust/sedona-schema" }
 sedona-proj = { path = "../../c/sedona-proj", default-features = false }

python/sedonadb/python/sedonadb/context.py

@@ -152,6 +152,26 @@ def read_parquet(
             self.options,
         )
 
+    def read_ogr(
+        self,
+        table_paths: Union[str, Path, Iterable[str]],
+        options: Optional[Dict[str, Any]] = None,
+    ) -> DataFrame:
+        from sedonadb.datasource import PyogrioFormatSpec
+
+        if isinstance(table_paths, (str, Path)):
+            table_paths = [table_paths]
+
+        spec = PyogrioFormatSpec()
+        if options is not None:
+            spec = spec.with_options(options)
+
+        return DataFrame(
+            self._impl,
+            self._impl.read_external_format(spec, [str(path) for path in table_paths], False),
+            self.options,
+        )
+
     def sql(self, sql: str) -> DataFrame:
         """Create a [DataFrame][sedonadb.dataframe.DataFrame] by executing SQL
 

python/sedonadb/python/sedonadb/datasource.py

@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from sedonadb._lib import PyExternalFormat
+
+
+class ExternalFormatSpec:
+    def clone(self):
+        raise NotImplementedError()
+
+    @property
+    def extension(self):
+        return ""
+
+    def with_options(self, options):
+        raise NotImplementedError(
+            f"key/value options not supported by {type(self).__name__}"
+        )
+
+    def open_reader(self, args):
+        raise NotImplementedError()
+
+    def infer_schema(self, object):
+        raise NotImplementedError()
+
+    def __sedona_external_format__(self):
+        return PyExternalFormat(self)
+
+
+class PyogrioFormatSpec(ExternalFormatSpec):
+    def __init__(self, extension=""):
+        import pyogrio.raw
+
+        self._raw = pyogrio.raw
+        self._extension = extension
+        self._options = {}
+
+    def clone(self):
+        cloned = type(self)(self.extension)
+        cloned._options.update(self._options)
+        return cloned
+
+    def with_options(self, options):
+        cloned = self.clone()
+        cloned._options.update(options)
+        return cloned
+
+    @property
+    def extension(self) -> str:
+        return self._extension
+
+    def open_reader(self, args):
+        url = args.src.to_url()
+        if url is None:
+            raise ValueError(f"Can't convert {args.src} to OGR-openable object")
+
+        if url.startswith("http://") or url.startswith("https://"):
+            ogr_src = f"/vsicurl/{url}"
+        elif url.startswith("file://"):
+            ogr_src = url.removeprefix("file://")
+        else:
+            raise ValueError(f"Can't open {url} with OGR")
+
+        if args.is_projected():
+            file_names = args.file_schema.names
+            columns = [file_names[i] for i in args.file_projection]
+        else:
+            columns = None
+
+        # TODO: Column order is not respected here, so we still need a utility to
+        # ensure match with the projected file schema
+        PyogrioReaderShelter(self._raw.ogr_open_arrow(ogr_src, {}, columns=columns))
+
+
+class PyogrioReaderShelter:
+    def __init__(self, inner):
+        self._inner = inner
+        self._meta, self._reader = self._inner.__enter__()
+
+    def __del__(self):
+        self._inner.__exit__(None, None, None)
+
+    def __arrow_c_stream__(self, requested_schema=None):
+        return self._reader.__arrow_c_stream__()

python/sedonadb/src/context.rs

@@ -23,6 +23,7 @@ use tokio::runtime::Runtime;
 
 use crate::{
     dataframe::InternalDataFrame,
+    datasource::PyExternalFormat,
     error::PySedonaError,
     import_from::{import_ffi_scalar_udf, import_table_provider_from_any},
     runtime::wait_for_future,
@@ -107,6 +108,26 @@ impl InternalContext {
         Ok(InternalDataFrame::new(df, self.runtime.clone()))
     }
 
+    pub fn read_external_format<'py>(
+        &self,
+        py: Python<'py>,
+        format_spec: Bound<PyAny>,
+        table_paths: Vec<String>,
+        check_extension: bool,
+    ) -> Result<InternalDataFrame, PySedonaError> {
+        let spec = format_spec
+            .call_method0("__sedona_external_format__")?
+            .extract::<PyExternalFormat>()?;
+        let df = wait_for_future(
+            py,
+            &self.runtime,
+            self.inner
+                .read_external_format(Arc::new(spec), table_paths, None, check_extension),
+        )??;
+
+        Ok(InternalDataFrame::new(df, self.runtime.clone()))
+    }
+
     pub fn sql<'py>(
         &self,
         py: Python<'py>,