Replace the Arrow scan with a port of the neo branch's scan

_duckdb-stubs/__init__.pyi

@@ -443,6 +443,7 @@ class DuckDBPyConnection:
         compression: ParquetCompression | None = None,
     ) -> DuckDBPyRelation: ...
     def register(self, view_name: str, python_object: object) -> DuckDBPyConnection: ...
+    def _register_arrow_source(self, view_name: str, source: object) -> DuckDBPyConnection: ...
     def register_filesystem(self, filesystem: fsspec.AbstractFileSystem) -> None: ...
     def remove_function(self, name: str) -> DuckDBPyConnection: ...
     def rollback(self) -> DuckDBPyConnection: ...
@@ -1243,6 +1244,14 @@ def table_function(
     connection: DuckDBPyConnection | None = None,
 ) -> DuckDBPyRelation: ...
 def tf(*, connection: DuckDBPyConnection | None = None) -> dict[str, typing.Any]: ...
+def _set_arrow_sources(sources: object) -> None: ...
+def _translate_predicate(
+    connection: DuckDBPyConnection,
+    predicate: str,
+    columns: object = None,
+    schema: object = None,
+    select: str = "*",
+) -> lst[tuple[str, object | None]]: ...
 def tokenize(query: str) -> lst[tuple[int, token_type]]: ...
 def torch(*, connection: DuckDBPyConnection | None = None) -> dict[str, typing.Any]: ...
 def type(type_str: str, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...

duckdb/__init__.py

@@ -9,6 +9,7 @@
 - Future versions may move them into submodules with deprecation warnings.
 """
 
+import _duckdb
 from _duckdb import (
     Appender,
     BinderException,
@@ -158,6 +159,7 @@
     write_csv,
 )
 
+from duckdb import _sources
 from duckdb._dbapi_type_object import (
     BINARY,
     DATETIME,
@@ -208,6 +210,8 @@
     Value,
 )
 
+_duckdb._set_arrow_sources(_sources)
+
 __all__: list[str] = [
     "Appender",
     "BinaryValue",

duckdb/_sources/__init__.py

@@ -0,0 +1,59 @@
+"""The Python objects the Arrow scan reads, each behind a source that exports an Arrow C stream for it.
+
+The scan calls `adapt` for an Arrow shaped object and then speaks only to the source: `rows()`, `pull_under_gil`,
+`accepts(predicate)`, `stream(columns, filters)` and, when present, `__arrow_c_schema__()`; `one_shot` is
+informational, since the scan decides single use from the underlying object.
+Every module of the package is imported here, since importing one while a query runs could load the extension a
+second time; pyarrow, polars and pandas are imported only when an object's class, or a class it derives from, is
+defined in them, and their classes are then recognised by identity.
+"""
+
+from __future__ import annotations
+
+import datetime
+
+from duckdb._sources import nodes, to_arrow, to_polars
+from duckdb._sources.arrow import (
+    ArrowArraySource,
+    ArrowCapsuleSource,
+    ArrowScanSource,
+    ArrowStreamSource,
+    Source,
+    class_modules,
+)
+from duckdb._sources.polars import polars_source
+from duckdb._sources.pyarrow import DatasetSource, ScannerSource, pyarrow_source
+
+__all__ = ["ArrowScanSource", "Source", "adapt", "nodes", "to_arrow", "to_polars"]
+
+_CAPSULE = type(datetime.datetime_CAPI)
+
+
+def adapt(obj: object) -> Source:
+    """The source for an Arrow shaped object; an existing `Source` is returned unchanged."""
+    if isinstance(obj, Source):
+        return obj
+    if type(obj) is _CAPSULE:
+        return ArrowCapsuleSource(obj)
+    modules = class_modules(obj)
+    libraries = {module.partition(".")[0] for module in modules}
+    source = None
+    if "pyarrow" in libraries:
+        source = pyarrow_source(obj, modules)
+    if source is None and "polars" in libraries:
+        source = polars_source(obj)
+    if source is not None:
+        return source
+    if hasattr(obj, "__arrow_c_stream__"):
+        return ArrowStreamSource(obj)
+    if hasattr(obj, "__arrow_c_array__"):
+        return ArrowArraySource(obj)
+    if hasattr(obj, "to_reader") and hasattr(obj, "projected_schema"):
+        return ScannerSource(obj)
+    if hasattr(obj, "scanner") and hasattr(obj, "schema"):
+        return DatasetSource(obj)
+    message = (
+        "an Arrow object must export __arrow_c_stream__ or __arrow_c_array__, or be an Arrow stream capsule, a "
+        f"pyarrow Dataset or Scanner, or a polars DataFrame or LazyFrame; {type(obj).__name__} is none of these"
+    )
+    raise TypeError(message)

duckdb/_sources/arrow.py

@@ -0,0 +1,163 @@
+"""The source protocol, and the sources for any object exporting the Arrow PyCapsule interface or a bare capsule."""
+
+from __future__ import annotations
+
+import ctypes
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    import pyarrow as pa
+
+    from duckdb._sources.nodes import Node
+
+
+class Source:
+    """An object wrapped for a scan."""
+
+    #: Informational only: the scan decides from the underlying object whether it can be read once.
+    one_shot = False
+    #: Whether the schema the object declares cannot change while it is registered, so the scan may cache it.
+    schema_is_stable = False
+
+    def __init__(self, obj: object) -> None:
+        self.obj: Any = obj
+
+    def rows(self) -> int | None:
+        """How many rows a scan produces, when the object knows without reading itself, and None otherwise."""
+        return None
+
+
+class ArrowScanSource(Source):
+    """A source the Arrow scan reads through the stream `stream()` exports, optionally with `__arrow_c_schema__`."""
+
+    #: An opt in asking the scan to hold the GIL for pulls and release callbacks; no adapter here needs it, since a
+    #: producer that reaches Python takes the GIL itself, as pyarrow, ctypes, cffi and Cython callbacks do.
+    pull_under_gil = False
+
+    def accepts(self, predicate: Node) -> bool:
+        """Whether every stream will hold only rows where `predicate` is true by the engine's semantics."""
+        return False
+
+    def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]:
+        """A stream capsule or a (schema, array) capsule pair, and whether it holds exactly `columns` in order.
+
+        `columns` are declared column indexes, None for all; a source that cannot narrow returns every declared
+        column in declared order and False. `filters` are the accepted predicates, applied before narrowing.
+        """
+        return self.obj.__arrow_c_stream__(), False
+
+
+class _ArrowSchema(ctypes.Structure):
+    """The C struct an object's `_export_to_c` writes a schema into."""
+
+    _fields_ = (
+        ("format", ctypes.c_char_p),
+        ("name", ctypes.c_char_p),
+        ("metadata", ctypes.c_char_p),
+        ("flags", ctypes.c_int64),
+        ("n_children", ctypes.c_int64),
+        ("children", ctypes.c_void_p),
+        ("dictionary", ctypes.c_void_p),
+        ("release", ctypes.c_void_p),
+        ("private_data", ctypes.c_void_p),
+    )
+
+
+def class_modules(obj: object) -> set[str]:
+    """The modules defining the class of `obj` and the classes it derives from; a class may have no module."""
+    return {module for cls in type(obj).__mro__ if isinstance(module := getattr(cls, "__module__", None), str)}
+
+
+def _exported(exporter: object) -> pa.Schema | pa.Field:
+    """What `exporter._export_to_c` writes: a Schema for a struct, otherwise a Field keeping name and metadata."""
+    import pyarrow
+
+    exported = _ArrowSchema()
+    address = ctypes.addressof(exported)
+    exporter._export_to_c(address)  # type: ignore[attr-defined]
+    if exported.format == b"+s":
+        return pyarrow.Schema._import_from_c(address)
+    return pyarrow.Field._import_from_c(address)
+
+
+def arrow_schema(schema: object) -> pa.Schema:
+    """A pyarrow Schema for a schema object exporting through `__arrow_c_schema__` or the older `_export_to_c`."""
+    import pyarrow
+
+    if isinstance(schema, pyarrow.Schema):
+        return schema
+    hook = getattr(schema, "__arrow_c_schema__", None)
+    if callable(hook):
+        return pyarrow.Schema._import_from_c_capsule(hook())
+    exported = _exported(schema)
+    if not isinstance(exported, pyarrow.Schema):
+        message = f"a table schema must be a struct, not {exported.type}"
+        raise TypeError(message)
+    return exported
+
+
+def _schema_capsule(exporter: object) -> object:
+    """The schema capsule of a schema or type object exporting through `__arrow_c_schema__` or `_export_to_c`."""
+    hook = getattr(exporter, "__arrow_c_schema__", None)
+    if callable(hook):
+        return hook()
+    return _exported(exporter).__arrow_c_schema__()
+
+
+def _schema_hook(source: ArrowScanSource, attribute: str) -> None:
+    """Give `source` a `__arrow_c_schema__`, so bind reads the schema without opening a stream.
+
+    The hook reads the object's own `__arrow_c_schema__`, or else `attribute`, at every call, so an object whose
+    schema changes reports the new one.
+    """
+    obj = source.obj
+    if callable(getattr(obj, "__arrow_c_schema__", None)):
+        source.__dict__["__arrow_c_schema__"] = lambda: obj.__arrow_c_schema__()
+        return
+    exporter = getattr(obj, attribute, None)
+    if callable(getattr(exporter, "__arrow_c_schema__", None)) or callable(getattr(exporter, "_export_to_c", None)):
+        source.__dict__["__arrow_c_schema__"] = lambda: _schema_capsule(getattr(obj, attribute))
+
+
+class ArrowStreamSource(ArrowScanSource):
+    """Anything with `__arrow_c_stream__`: every column, its schema from `__arrow_c_schema__` or `schema`."""
+
+    #: The attribute that holds the object's schema when the object has no `__arrow_c_schema__` of its own.
+    schema_attribute = "schema"
+
+    def __init__(self, obj: object) -> None:
+        super().__init__(obj)
+        _schema_hook(self, self.schema_attribute)
+
+    def rows(self) -> int | None:
+        """The length of a pandas Series or DataFrame, which is its row count; unknown for anything else."""
+        if "pandas" not in {module.partition(".")[0] for module in class_modules(self.obj)}:
+            return None
+        import pandas
+
+        return len(self.obj) if isinstance(self.obj, pandas.Series | pandas.DataFrame) else None
+
+
+class ArrowArraySource(ArrowScanSource):
+    """Anything with `__arrow_c_array__` and no stream: one batch, read as often as asked."""
+
+    def __init__(self, obj: object) -> None:
+        super().__init__(obj)
+        _schema_hook(self, "type")
+
+    def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]:
+        """The schema and array capsules of the array, every column."""
+        return self.obj.__arrow_c_array__(), False
+
+
+class ArrowCapsuleSource(ArrowScanSource):
+    """A bare capsule; the scan checks that it is an unreleased `arrow_array_stream`."""
+
+    one_shot = True
+    schema_is_stable = True
+
+    def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]:
+        """The capsule itself, every column."""
+        return self.obj, False