From eb412b15ed4520c1fa09bf88f13cd16f416a7454 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Sat, 17 Oct 2026 20:50:09 +0000 Subject: [PATCH 1/2] Replace the Arrow scan with a port of the neo branch's scan --- _duckdb-stubs/__init__.pyi | 9 + duckdb/__init__.py | 4 + duckdb/_sources/__init__.py | 59 + duckdb/_sources/arrow.py | 163 ++ duckdb/_sources/nodes.py | 258 ++ duckdb/_sources/polars.py | 104 + duckdb/_sources/pyarrow.py | 157 ++ duckdb/_sources/to_arrow.py | 261 ++ duckdb/_sources/to_polars.py | 211 ++ src/arrow/CMakeLists.txt | 4 +- src/arrow/arrow_array_stream.cpp | 417 --- src/arrow/arrow_export_utils.cpp | 23 +- src/arrow/arrow_importer.cpp | 335 +++ src/arrow/filter_pushdown_visitor.cpp | 260 -- src/arrow/polars_filter_pushdown.cpp | 157 -- src/arrow/predicate_translation.cpp | 609 +++++ src/arrow/pyarrow_filter_pushdown.cpp | 336 --- src/arrow/python_arrow_scan.cpp | 822 ++++++ src/arrow/single_use.cpp | 108 + src/duckdb_python.cpp | 6 + .../arrow/arrow_array_stream.hpp | 114 - .../arrow/arrow_export_utils.hpp | 28 + .../duckdb_python/arrow/arrow_importer.hpp | 40 + .../arrow/filter_pushdown_visitor.hpp | 100 - .../arrow/polars_filter_pushdown.hpp | 23 - .../arrow/predicate_translation.hpp | 63 + .../arrow/pyarrow_filter_pushdown.hpp | 24 - .../duckdb_python/arrow/python_arrow_scan.hpp | 81 + .../duckdb_python/arrow/single_use.hpp | 31 + src/include/duckdb_python/exceptions.hpp | 5 + .../pyconnection/pyconnection.hpp | 10 +- src/include/duckdb_python/pyrelation.hpp | 2 +- .../duckdb_python/python_replacement_scan.hpp | 14 +- src/include/duckdb_python/pyutil.hpp | 14 + src/pyconnection.cpp | 115 +- src/pyresult.cpp | 1 - src/python_replacement_scan.cpp | 196 +- src/python_udf.cpp | 127 +- tests/fast/arrow/_pushdown_helpers.py | 156 +- tests/fast/arrow/scan/__init__.py | 0 .../fast/arrow/scan/test_arrow_scan_rules.py | 1145 ++++++++ .../arrow/scan/test_arrow_scan_sources.py | 1126 ++++++++ .../arrow/scan/test_arrow_scan_streams.py | 1420 ++++++++++ .../arrow/scan/test_predicate_translation.py | 795 ++++++ .../arrow/scan/test_python_arrow_scan_core.py | 2423 +++++++++++++++++ tests/fast/arrow/scan/test_sources_adapt.py | 849 ++++++ tests/fast/arrow/scan/test_translators.py | 1018 +++++++ .../fast/arrow/test_arrow_replacement_scan.py | 50 + .../arrow/test_arrow_single_use_stream.py | 2 +- tests/fast/arrow/test_arrow_stream_scan.py | 17 +- tests/fast/arrow/test_arrow_string_view.py | 32 + .../arrow/test_dictionary_filter_pushdown.py | 200 +- tests/fast/arrow/test_filter_pushdown.py | 359 +-- .../arrow/test_nan_null_filter_pushdown.py | 194 +- .../fast/arrow/test_polars_filter_pushdown.py | 311 ++- tests/fast/arrow/test_reader_readahead.py | 195 ++ .../arrow/test_timestamp_filter_pushdown.py | 158 +- tests/fast/udf/test_scalar_arrow.py | 56 + tests/slow/test_interpreter_shutdown.py | 46 + 59 files changed, 13531 insertions(+), 2312 deletions(-) create mode 100644 duckdb/_sources/__init__.py create mode 100644 duckdb/_sources/arrow.py create mode 100644 duckdb/_sources/nodes.py create mode 100644 duckdb/_sources/polars.py create mode 100644 duckdb/_sources/pyarrow.py create mode 100644 duckdb/_sources/to_arrow.py create mode 100644 duckdb/_sources/to_polars.py delete mode 100644 src/arrow/arrow_array_stream.cpp create mode 100644 src/arrow/arrow_importer.cpp delete mode 100644 src/arrow/filter_pushdown_visitor.cpp delete mode 100644 src/arrow/polars_filter_pushdown.cpp create mode 100644 src/arrow/predicate_translation.cpp delete mode 100644 src/arrow/pyarrow_filter_pushdown.cpp create mode 100644 src/arrow/python_arrow_scan.cpp create mode 100644 src/arrow/single_use.cpp delete mode 100644 src/include/duckdb_python/arrow/arrow_array_stream.hpp create mode 100644 src/include/duckdb_python/arrow/arrow_importer.hpp delete mode 100644 src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp delete mode 100644 src/include/duckdb_python/arrow/polars_filter_pushdown.hpp create mode 100644 src/include/duckdb_python/arrow/predicate_translation.hpp delete mode 100644 src/include/duckdb_python/arrow/pyarrow_filter_pushdown.hpp create mode 100644 src/include/duckdb_python/arrow/python_arrow_scan.hpp create mode 100644 src/include/duckdb_python/arrow/single_use.hpp create mode 100644 tests/fast/arrow/scan/__init__.py create mode 100644 tests/fast/arrow/scan/test_arrow_scan_rules.py create mode 100644 tests/fast/arrow/scan/test_arrow_scan_sources.py create mode 100644 tests/fast/arrow/scan/test_arrow_scan_streams.py create mode 100644 tests/fast/arrow/scan/test_predicate_translation.py create mode 100644 tests/fast/arrow/scan/test_python_arrow_scan_core.py create mode 100644 tests/fast/arrow/scan/test_sources_adapt.py create mode 100644 tests/fast/arrow/scan/test_translators.py create mode 100644 tests/fast/arrow/test_reader_readahead.py diff --git a/_duckdb-stubs/__init__.pyi b/_duckdb-stubs/__init__.pyi index a4ca53f4..6003c94c 100644 --- a/_duckdb-stubs/__init__.pyi +++ b/_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: ... diff --git a/duckdb/__init__.py b/duckdb/__init__.py index dfbf93b0..2d9b3759 100644 --- a/duckdb/__init__.py +++ b/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", diff --git a/duckdb/_sources/__init__.py b/duckdb/_sources/__init__.py new file mode 100644 index 00000000..f8937542 --- /dev/null +++ b/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) diff --git a/duckdb/_sources/arrow.py b/duckdb/_sources/arrow.py new file mode 100644 index 00000000..e9cb4851 --- /dev/null +++ b/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 diff --git a/duckdb/_sources/nodes.py b/duckdb/_sources/nodes.py new file mode 100644 index 00000000..fec4c265 --- /dev/null +++ b/duckdb/_sources/nodes.py @@ -0,0 +1,258 @@ +"""The predicate tree the Arrow scan offers a source, built in C++ from a filter the engine would apply. + +A tree holds `Binary` for AND, OR and the comparisons = != < <= > >= of a `Col` (always on the left) with a +`Lit`, `Unary("NOT", predicate)`, `Postfix("IS NULL" or "IS NOT NULL", Col)` and `In(Col, [Lit, ...])`; BETWEEN +arrives as AND of two comparisons. `Col.parts` is the declared column name followed by the struct field names that +lead into it. `Lit.type` is the engine's name for the constant's type, which is exactly the column's type, such as +"INTEGER" or "DECIMAL(10,2)"; a "TIMESTAMP WITH TIME ZONE" value is the instant as an int count of microseconds +since the epoch, and every other value is the plain Python value. +""" + +from __future__ import annotations + +import datetime +import decimal +import math +import re +import struct +from dataclasses import dataclass + + +class Untranslatable(ValueError): + """Raised by a translator for a node it cannot evaluate exactly as the engine would.""" + + +class Node: + """A node of the predicate tree.""" + + def fragment(self) -> str: + """The node as parenthesised SQL, the text EXPLAIN shows under Filters.""" + return fragment(self) + + def _children(self) -> tuple[Node, ...]: + return () + + def _render(self, children: list[str]) -> str: + raise NotImplementedError + + +@dataclass(frozen=True, eq=False) +class Col(Node): + """A column, or a struct field inside one.""" + + parts: tuple[str, ...] + + def _render(self, children: list[str]) -> str: + return ".".join(_quote(part) for part in self.parts) + + +@dataclass(frozen=True, eq=False) +class Lit(Node): + """A constant of the column's type.""" + + value: object + type: str + + def _render(self, children: list[str]) -> str: + return _literal(self.value, self.type) + + +@dataclass(frozen=True, eq=False) +class Binary(Node): + """A comparison of a column with a constant, or AND or OR of two predicates.""" + + op: str + left: Node + right: Node + + def _children(self) -> tuple[Node, ...]: + return (self.left, self.right) + + def _render(self, children: list[str]) -> str: + return f"({children[0]} {self.op} {children[1]})" + + +@dataclass(frozen=True, eq=False) +class Unary(Node): + """NOT of a predicate.""" + + op: str + operand: Node + + def _children(self) -> tuple[Node, ...]: + return (self.operand,) + + def _render(self, children: list[str]) -> str: + return f"({self.op} {children[0]})" + + +@dataclass(frozen=True, eq=False) +class Postfix(Node): + """IS NULL or IS NOT NULL of a column.""" + + op: str + operand: Node + + def _children(self) -> tuple[Node, ...]: + return (self.operand,) + + def _render(self, children: list[str]) -> str: + return f"({children[0]} {self.op})" + + +@dataclass(frozen=True, eq=False) +class In(Node): + """Membership of a column in a list of at least one constant.""" + + operand: Node + values: list[Node] + + def _children(self) -> tuple[Node, ...]: + return (self.operand, *self.values) + + def _render(self, children: list[str]) -> str: + if not self.values: + message = "an IN needs at least one candidate" + raise Untranslatable(message) + return f"({children[0]} IN ({', '.join(children[1:])}))" + + +def fragment(node: Node) -> str: + """The SQL text of `node`, built without recursion so that any depth of AND and OR renders.""" + texts: list[str] = [] + pending: list[tuple[Node, bool]] = [(node, False)] + while pending: + current, expanded = pending.pop() + children = current._children() + if children and not expanded: + pending.append((current, True)) + pending.extend((child, False) for child in reversed(children)) + continue + start = len(texts) - len(children) + rendered = current._render(texts[start:]) + del texts[start:] + texts.append(rendered) + return texts[0] + + +_INTEGER_RANGES = { + "TINYINT": (-(1 << 7), (1 << 7) - 1), + "SMALLINT": (-(1 << 15), (1 << 15) - 1), + "INTEGER": (-(1 << 31), (1 << 31) - 1), + "BIGINT": (-(1 << 63), (1 << 63) - 1), + "UTINYINT": (0, (1 << 8) - 1), + "USMALLINT": (0, (1 << 16) - 1), + "UINTEGER": (0, (1 << 32) - 1), + "UBIGINT": (0, (1 << 64) - 1), +} + +#: The step of each naive timestamp type, in microseconds. +_TIMESTAMP_STEPS = {"TIMESTAMP_S": 1_000_000, "TIMESTAMP_MS": 1_000, "TIMESTAMP": 1} + +_DECIMAL = re.compile(r"DECIMAL\((\d+),(\d+)\)") + + +def _decimal_type(type_text: str) -> tuple[int, int] | None: + """The precision and scale of a DECIMAL type text, None for any other type.""" + match = _DECIMAL.fullmatch(type_text) + return (int(match[1]), int(match[2])) if match else None + + +def exact_value(constant: Lit) -> bool: + """Whether the constant's value is exactly a value of its type, as the translators require.""" + value, type_text = constant.value, constant.type + if type_text == "BOOLEAN": + return isinstance(value, bool) + if type_text in _INTEGER_RANGES or type_text == "TIMESTAMP WITH TIME ZONE": + low, high = _INTEGER_RANGES.get(type_text, (-(1 << 63), (1 << 63) - 1)) + return isinstance(value, int) and not isinstance(value, bool) and low <= value <= high + if type_text in ("FLOAT", "DOUBLE"): + if not isinstance(value, float) or not math.isfinite(value): + return False + try: + return type_text == "DOUBLE" or struct.unpack("f", struct.pack("f", value))[0] == value + except OverflowError: + return False + if type_text == "VARCHAR": + return isinstance(value, str) + if type_text == "BLOB": + return isinstance(value, bytes) + if type_text == "DATE": + return isinstance(value, datetime.date) and not isinstance(value, datetime.datetime) + if type_text == "TIME": + return isinstance(value, datetime.time) and value.tzinfo is None + if type_text in _TIMESTAMP_STEPS: + return ( + isinstance(value, datetime.datetime) + and value.tzinfo is None + and value.microsecond % _TIMESTAMP_STEPS[type_text] == 0 + ) + shape = _decimal_type(type_text) + if shape is not None and isinstance(value, decimal.Decimal): + return _fits(value, *shape) + return False + + +def _fits(value: decimal.Decimal, precision: int, scale: int) -> bool: + if not value.is_finite(): + return False + _, digits, exponent = value.as_tuple() + unscaled = int("".join(map(str, digits)) or "0") + shift = int(exponent) + scale + if shift >= 0: + unscaled *= 10**shift + elif unscaled % 10**-shift == 0: + unscaled //= 10**-shift + else: + return False + return len(str(unscaled)) <= precision + + +def _quote(name: str) -> str: + escaped = name.replace('"', '""') + return f'"{escaped}"' + + +def _text(value: str) -> str: + escaped = value.replace("'", "''") + return f"'{escaped}'" + + +def _fraction(microseconds: int) -> str: + return f".{microseconds:06d}".rstrip("0") if microseconds else "" + + +def _clock(value: datetime.time | datetime.datetime) -> str: + return f"{value.hour:02d}:{value.minute:02d}:{value.second:02d}{_fraction(value.microsecond)}" + + +def _literal(value: object, type_text: str) -> str: + if isinstance(value, bool): + return "TRUE" if value else "FALSE" + if type_text == "TIMESTAMP WITH TIME ZONE" and isinstance(value, int): + moment = datetime.datetime(1970, 1, 1) + datetime.timedelta(microseconds=value) + return f"TIMESTAMPTZ '{moment.date().isoformat()} {_clock(moment)}+00'" + if isinstance(value, int): + return str(value) + if isinstance(value, float): + return f"{value!r}::DOUBLE" + if isinstance(value, str): + return _text(value) + if isinstance(value, bytes): + # One escape per byte, since the engine reads exactly two hex digits after each. + escaped = "".join(f"\\x{byte:02x}" for byte in value) + return f"'{escaped}'::BLOB" + if isinstance(value, decimal.Decimal) and value.is_finite(): + shape = _decimal_type(type_text) + if shape is None: + _, digits, exponent = value.as_tuple() + scale = max(0, -int(exponent)) + shape = (max(len(digits) + max(int(exponent), 0), scale + 1), scale) + return f"CAST('{value:f}' AS DECIMAL({shape[0]}, {shape[1]}))" + if isinstance(value, datetime.datetime): + return f"TIMESTAMP '{value.date().isoformat()} {_clock(value)}'" + if isinstance(value, datetime.date): + return f"DATE '{value.isoformat()}'" + if isinstance(value, datetime.time): + return f"TIME '{_clock(value)}'" + return f"{_text(str(value))}::{type_text}" diff --git a/duckdb/_sources/polars.py b/duckdb/_sources/polars.py new file mode 100644 index 00000000..59488bc5 --- /dev/null +++ b/duckdb/_sources/polars.py @@ -0,0 +1,104 @@ +"""The polars sources: a DataFrame read as the pyarrow Table of one `to_arrow()`, a LazyFrame collected per scan. + +Polars exports a column of Python objects (`pl.Object`) as the objects' addresses, so a frame holding one is refused. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from duckdb._sources.arrow import ArrowScanSource, ArrowStreamSource +from duckdb._sources.nodes import Untranslatable +from duckdb._sources.pyarrow import PyArrowTableSource +from duckdb._sources.to_arrow import nan_may_satisfy +from duckdb._sources.to_polars import to_polars + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from duckdb._sources.nodes import Node + + +def polars_source(obj: object) -> ArrowScanSource | None: + """The source for a polars object, or None when polars has none for it.""" + import polars + + if isinstance(obj, polars.DataFrame): + _refuse_object_columns(obj.schema) + return PyArrowTableSource(obj.to_arrow()) + if isinstance(obj, polars.LazyFrame): + return LazyFrameSource(obj) + if isinstance(obj, polars.Series): + _refuse_object_columns({obj.name: obj.dtype}) + return PolarsStreamSource(obj) + return None + + +class PolarsStreamSource(ArrowStreamSource): + """A polars Series.""" + + schema_is_stable = True + + def rows(self) -> int | None: + """The length of the series.""" + return len(self.obj) + + +class LazyFrameSource(ArrowScanSource): + """A polars LazyFrame, filtered, narrowed and collected once per scan, declaring what an empty collect exports.""" + + schema_is_stable = True + + def __init__(self, obj: object) -> None: + super().__init__(obj) + self._empty: Any = None + self._declared: Any = None + + def _empty_frame(self) -> Any: # noqa: ANN401 + if self._empty is None: + empty = self.obj.head(0).collect() + _refuse_object_columns(empty.schema) + self._empty = empty + return self._empty + + def _declared_schema(self) -> Any: # noqa: ANN401 + if self._declared is None: + self._declared = self._empty_frame().to_arrow().schema + return self._declared + + def __arrow_c_schema__(self) -> object: + """The schema of the frame's first zero rows as Arrow.""" + return self._declared_schema().__arrow_c_schema__() + + def accepts(self, predicate: Node) -> bool: + """Whether `to_polars` translates `predicate` exactly and statistics pruning cannot drop a row it keeps.""" + try: + to_polars(predicate, self._empty_frame().schema) + # A scan over files prunes row groups by minimum and maximum statistics, which leave NaN out. + return not nan_may_satisfy(predicate, self._declared_schema()) + except Untranslatable: + return False + + def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]: + """The collected frame's stream after `filters`, narrowed to `columns` unless none are asked.""" + schema = self._empty_frame().schema + plan = self.obj + for predicate in filters: + plan = plan.filter(to_polars(predicate, schema)) + # A polars frame without columns has no rows, so a request for no columns reads them all. + if not columns: + return plan.collect().to_arrow().__arrow_c_stream__(), False + import polars + + # Polars reads a name like ^x$ or * as a selector, so columns are taken by position. + plan = plan.select([polars.nth(i) for i in columns]) + return plan.collect().to_arrow().__arrow_c_stream__(), True + + +def _refuse_object_columns(schema: Mapping[str, object]) -> None: + import polars + + for name, dtype in schema.items(): + if dtype == polars.Object: + message = f"the polars column '{name}' holds Python objects, which polars cannot export as Arrow" + raise TypeError(message) diff --git a/duckdb/_sources/pyarrow.py b/duckdb/_sources/pyarrow.py new file mode 100644 index 00000000..6b805c5f --- /dev/null +++ b/duckdb/_sources/pyarrow.py @@ -0,0 +1,157 @@ +"""The pyarrow sources: a Table or RecordBatch, a RecordBatchReader, a Dataset and a Scanner. + +An object that only looks like a Dataset or Scanner reaches `adapt` only through the private registration hook, +since the scan's entry points accept the pyarrow classes alone; it is never offered a filter. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from duckdb._sources.arrow import ArrowArraySource, ArrowScanSource, ArrowStreamSource, arrow_schema +from duckdb._sources.nodes import Untranslatable +from duckdb._sources.to_arrow import nan_may_satisfy, to_arrow + +if TYPE_CHECKING: + from collections.abc import Sequence + + import pyarrow as pa + + from duckdb._sources.nodes import Node + + +def pyarrow_source(obj: object, modules: set[str]) -> ArrowScanSource | None: + """The source for a pyarrow object, given the modules of its classes, or None when pyarrow has none for it.""" + import pyarrow + + if isinstance(obj, pyarrow.RecordBatchReader): + return PyArrowReaderSource(obj) + if isinstance(obj, pyarrow.Table | pyarrow.RecordBatch): + return PyArrowTableSource(obj) + if "pyarrow._dataset" in modules: + import pyarrow.dataset + + if isinstance(obj, pyarrow.dataset.Scanner): + return PyArrowScannerSource(obj) + if isinstance(obj, pyarrow.dataset.Dataset): + return PyArrowDatasetSource(obj) + if isinstance(obj, pyarrow.ChunkedArray): + return PyArrowStreamSource(obj) + if isinstance(obj, pyarrow.Array): + return PyArrowArraySource(obj) + return None + + +class PyArrowTableSource(ArrowStreamSource): + """A pyarrow Table or RecordBatch, narrowed by position with `select`, which copies nothing.""" + + schema_is_stable = True + + def rows(self) -> int | None: + """The number of rows.""" + return len(self.obj) + + def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]: + """The object's stream, narrowed to `columns` when asked.""" + if columns is None: + return self.obj.__arrow_c_stream__(), False + return self.obj.select(list(columns)).__arrow_c_stream__(), True + + +class PyArrowReaderSource(ArrowStreamSource): + """A pyarrow RecordBatchReader, every column.""" + + one_shot = True + schema_is_stable = True + + def rows(self) -> int | None: + """Unknown, since only reading the stream counts it.""" + return None + + +class PyArrowStreamSource(ArrowStreamSource): + """A pyarrow ChunkedArray, its schema from its `type`.""" + + schema_attribute = "type" + + schema_is_stable = True + + def rows(self) -> int | None: + """The length of the array.""" + return len(self.obj) + + +class PyArrowArraySource(ArrowArraySource): + """A pyarrow Array or StructArray, one batch.""" + + schema_is_stable = True + + def rows(self) -> int | None: + """The length of the array.""" + return len(self.obj) + + +class DatasetSource(ArrowScanSource): + """An object with a pyarrow Dataset's `schema` and `scanner(columns=, filter=)`, scanned afresh per stream.""" + + def __init__(self, obj: object) -> None: + super().__init__(obj) + self._declared: Any = None + + def declared_schema(self) -> pa.Schema: + """The schema the dataset declares, kept for the life of a stable source and read anew for any other.""" + if not self.schema_is_stable: + return arrow_schema(self.obj.schema) + if self._declared is None: + self._declared = arrow_schema(self.obj.schema) + return self._declared + + def __arrow_c_schema__(self) -> object: + """The schema the dataset declares.""" + return self.declared_schema().__arrow_c_schema__() + + def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]: + """A reader over a scanner that applies `filters` and narrows to `columns` by name.""" + schema = self.declared_schema() + mask = None + for predicate in filters: + term = to_arrow(predicate, schema) + mask = term if mask is None else mask & term + names = schema.names + chosen = None if columns is None else [names[i] for i in columns] + scanner = type(self.obj).scanner(self.obj, columns=chosen, filter=mask) + return scanner.to_reader().__arrow_c_stream__(), chosen is not None + + +class PyArrowDatasetSource(DatasetSource): + """A pyarrow Dataset, which applies every predicate `to_arrow` translates and prunes by it.""" + + schema_is_stable = True + + def accepts(self, predicate: Node) -> bool: + """Whether `to_arrow` translates `predicate` exactly and statistics pruning cannot drop a row it keeps.""" + schema = self.declared_schema() + try: + to_arrow(predicate, schema) + except Untranslatable: + return False + # Row groups are pruned by minimum and maximum statistics, which leave NaN out. + return not nan_may_satisfy(predicate, schema) + + +class ScannerSource(ArrowScanSource): + """An object with a pyarrow Scanner's `projected_schema` and `to_reader()`; it keeps its projection and filter.""" + + def __arrow_c_schema__(self) -> object: + """The schema the scanner projects.""" + return arrow_schema(self.obj.projected_schema).__arrow_c_schema__() + + def stream(self, columns: Sequence[int] | None, filters: Sequence[Node]) -> tuple[object, bool]: + """A fresh reader over the scanner, every column.""" + return self.obj.to_reader().__arrow_c_stream__(), False + + +class PyArrowScannerSource(ScannerSource): + """A pyarrow Scanner.""" + + schema_is_stable = True diff --git a/duckdb/_sources/to_arrow.py b/duckdb/_sources/to_arrow.py new file mode 100644 index 00000000..c1c6815b --- /dev/null +++ b/duckdb/_sources/to_arrow.py @@ -0,0 +1,261 @@ +"""A predicate tree as a pyarrow compute expression that a Dataset scanner evaluates exactly as the engine would. + +A constant becomes a scalar of its own engine type, never of the column's Arrow type, so a value between two points +of a coarser column keeps its meaning; a zoned scalar is labelled with the column's zone, since pyarrow compares +across units but not across zones. +""" + +from __future__ import annotations + +import operator +from typing import TYPE_CHECKING, Any + +from duckdb._sources.nodes import Binary, Col, In, Lit, Postfix, Unary, Untranslatable, exact_value + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + import pyarrow as pa + import pyarrow.compute as pc + + from duckdb._sources.nodes import Node + +_COMPARISONS: dict[str, Callable[[Any, Any], Any]] = { + "=": operator.eq, + "!=": operator.ne, + "<": operator.lt, + "<=": operator.le, + ">": operator.gt, + ">=": operator.ge, +} + +_INTEGERS = { + "int8": "TINYINT", + "int16": "SMALLINT", + "int32": "INTEGER", + "int64": "BIGINT", + "uint8": "UTINYINT", + "uint16": "USMALLINT", + "uint32": "UINTEGER", + "uint64": "UBIGINT", +} + +_NAIVE_TIMESTAMPS = {"s": "TIMESTAMP_S", "ms": "TIMESTAMP_MS", "us": "TIMESTAMP"} + +_UNITS = {name: unit for unit, name in _NAIVE_TIMESTAMPS.items()} + +_PER_SECOND = {"s": 1, "ms": 1_000, "us": 1_000_000} + + +def to_arrow(node: Node, schema: pa.Schema) -> pc.Expression: + """`node` over `schema`, the schema the source declares, or `Untranslatable` where pyarrow would differ.""" + _refuse_unfilterable(schema) + try: + return _translate(node, schema) + except RecursionError as error: + raise Untranslatable(node) from error + + +def _translate(node: Node, schema: pa.Schema) -> pc.Expression: + import pyarrow as pa + import pyarrow.compute as pc + + match node: + case Binary(op="AND", left=left, right=right): + return _translate(left, schema) & _translate(right, schema) + case Binary(op="OR", left=left, right=right): + return _translate(left, schema) | _translate(right, schema) + case Unary(op="NOT", operand=operand): + return ~_translate(operand, schema) + case Postfix(op="IS NULL", operand=Col(parts=parts)): + return _column(schema, parts)[0].is_null() + case Postfix(op="IS NOT NULL", operand=Col(parts=parts)): + return _column(schema, parts)[0].is_valid() + case Binary(op=op, left=Col(parts=parts), right=Lit() as constant) if op in _COMPARISONS: + field, kind = _column(schema, parts) + compared = _COMPARISONS[op](field, _scalar(constant, kind)) + if op in (">", ">=") and pa.types.is_floating(kind): + # The engine orders NaN above every number; NaN alone is unequal to itself, also in a dictionary, + # where pyarrow has no is_nan kernel. + return compared | (field != field) + return compared + case In(operand=Col(parts=parts), values=values) if values: + field, kind = _column(schema, parts) + member = _membership(field, kind, values) + # pyarrow's is_in answers false for a NULL row where the engine answers NULL, which NOT turns into true. + return pc.if_else(field.is_valid(), member, pa.scalar(None, pa.bool_())) + raise Untranslatable(node) + + +def nan_may_satisfy(node: Node, schema: pa.Schema) -> bool: + """Whether a NaN in a floating point column could make `node` true, as > and >= do and anything under NOT may.""" + import pyarrow as pa + + pending: list[tuple[Node, bool]] = [(node, False)] + while pending: + current, negated = pending.pop() + match current: + case Binary(op="AND" | "OR", left=left, right=right): + pending.extend(((left, negated), (right, negated))) + case Unary(op="NOT", operand=operand): + pending.append((operand, True)) + case Binary(left=Col(parts=parts)) | In(operand=Col(parts=parts)) | Postfix(operand=Col(parts=parts)) if ( + pa.types.is_floating(_column(schema, parts)[1]) + ): + if negated or (isinstance(current, Binary) and current.op in (">", ">=", "!=")): + return True + return False + + +def _refuse_unfilterable(schema: pa.Schema) -> None: + """Refuse a view layout or run end encoding at any depth, or a dictionary nested inside another column.""" + import pyarrow as pa + + pending = [(field.type, True) for field in schema] + while pending: + kind, column = pending.pop() + if ( + pa.types.is_string_view(kind) + or pa.types.is_binary_view(kind) + or pa.types.is_list_view(kind) + or pa.types.is_large_list_view(kind) + or pa.types.is_run_end_encoded(kind) + or (pa.types.is_dictionary(kind) and not column) + ): + raise Untranslatable(kind) + if pa.types.is_dictionary(kind): + pending.append((kind.value_type, False)) + elif isinstance(kind, pa.BaseExtensionType): + pending.append((kind.storage_type, False)) + else: + pending.extend((kind.field(i).type, False) for i in range(kind.num_fields)) + + +def _engine_type(kind: pa.DataType) -> str | None: + """The engine's name for the type it reads an Arrow value type as, for the types translated here.""" + import pyarrow as pa + + if pa.types.is_boolean(kind): + return "BOOLEAN" + if pa.types.is_integer(kind): + return _INTEGERS[str(kind)] + if pa.types.is_float32(kind): + return "FLOAT" + if pa.types.is_float64(kind): + return "DOUBLE" + if pa.types.is_string(kind) or pa.types.is_large_string(kind): + return "VARCHAR" + if pa.types.is_binary(kind) or pa.types.is_large_binary(kind) or pa.types.is_fixed_size_binary(kind): + return "BLOB" + # A date64 may hold part of a day, which pyarrow compares and the engine truncates. + if pa.types.is_date32(kind): + return "DATE" + if pa.types.is_time(kind) and kind.unit in _PER_SECOND: + return "TIME" + if pa.types.is_timestamp(kind) and kind.tz is None: + return _NAIVE_TIMESTAMPS.get(kind.unit) + if pa.types.is_timestamp(kind) and kind.unit in _PER_SECOND: + return "TIMESTAMP WITH TIME ZONE" + if pa.types.is_decimal128(kind): + return f"DECIMAL({kind.precision},{kind.scale})" + return None + + +def _column(schema: pa.Schema, parts: tuple[str, ...]) -> tuple[pc.Expression, Any]: + """The field expression for `parts` and its value type; a constant is checked against the type where it is used.""" + import pyarrow as pa + import pyarrow.compute as pc + + if not parts: + raise Untranslatable(parts) + # A missing or repeated name has index -1, so a name is never resolved to one of two columns. + index = schema.get_field_index(parts[0]) + if index < 0: + raise Untranslatable(parts) + kind = schema.field(index).type + for name in parts[1:]: + child = kind.get_field_index(name) if pa.types.is_struct(kind) else -1 + if child < 0: + raise Untranslatable(parts) + kind = kind.field(child).type + if pa.types.is_dictionary(kind): + kind = kind.value_type + # A union has no validity of its own, so whether its row is NULL is not plainly the engine's answer. + if pa.types.is_union(kind): + raise Untranslatable(kind) + return pc.field(*parts), kind + + +def _scalar(constant: Lit, kind: pa.DataType) -> pa.Scalar: + """A scalar of the constant's own engine type, refusing a constant that is not a value of the column's type.""" + import pyarrow as pa + + if constant.type != _engine_type(kind) or not exact_value(constant): + raise Untranslatable(constant) + match constant.type: + case "BOOLEAN": + own = pa.bool_() + case "FLOAT": + own = pa.float32() + case "DOUBLE": + own = pa.float64() + case "VARCHAR": + own = pa.string() + case "BLOB": + own = pa.binary() + case "DATE": + own = pa.date32() + case "TIME": + own = pa.time64("us") + case "TIMESTAMP WITH TIME ZONE": + own = pa.timestamp("us", tz=kind.tz) + case "TIMESTAMP_S" | "TIMESTAMP_MS" | "TIMESTAMP": + own = pa.timestamp(_UNITS[constant.type]) + case _: + # The integer and decimal types match the column's Arrow type one to one. + own = kind + return pa.scalar(constant.value, own) + + +def _membership(field: pc.Expression, kind: pa.DataType, values: Sequence[Node]) -> pc.Expression: + """Membership of `field` in the constants, each converted exactly to the column's own type.""" + import pyarrow as pa + import pyarrow.compute as pc + + held = [] + zero = False + for constant in values: + if not isinstance(constant, Lit): + raise Untranslatable(constant) + value = _held(_scalar(constant, kind), kind) + if value is None: + continue + if pa.types.is_floating(kind) and value == 0.0: + zero = True + else: + held.append(value) + member = pc.scalar(False) if not held else field.isin(pa.array(held, type=kind)) + if zero: + # is_in compares the bits of a float, which tells -0.0 from 0.0 where the engine does not. + member = (field == pa.scalar(0.0, kind)) | member + return member + + +def _held(scalar: pa.Scalar, kind: pa.DataType) -> object: + """The scalar's value as a value of the column's type, or None when no row of the column can hold it.""" + import pyarrow as pa + + if pa.types.is_timestamp(kind) and kind.tz is not None: + return _on_grid(scalar.value, kind.unit) + if pa.types.is_time(kind): + return _on_grid(scalar.value, kind.unit) + value = scalar.as_py() + if pa.types.is_fixed_size_binary(kind) and len(value) != kind.byte_width: + return None + return value + + +def _on_grid(count: int, unit: str) -> int | None: + """A count of microseconds as a count of `unit`, or None when it lies between two of them.""" + step = 1_000_000 // _PER_SECOND[unit] + return count // step if count % step == 0 else None diff --git a/duckdb/_sources/to_polars.py b/duckdb/_sources/to_polars.py new file mode 100644 index 00000000..c794b1d9 --- /dev/null +++ b/duckdb/_sources/to_polars.py @@ -0,0 +1,211 @@ +"""A predicate tree as a polars expression that a LazyFrame filter evaluates exactly as the engine would. + +Polars orders NaN above every number, equates -0.0 with 0.0 and combines NULL with three valued logic, as the +engine does. A constant becomes a literal of its own engine type, never of the column's dtype; a zoned literal is +labelled with the column's zone, since polars refuses to compare across zones. +""" + +from __future__ import annotations + +import operator +from typing import TYPE_CHECKING, Any + +from duckdb._sources.nodes import Binary, Col, In, Lit, Postfix, Unary, Untranslatable, exact_value + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + import polars as pl + + from duckdb._sources.nodes import Node + +_COMPARISONS: dict[str, Callable[[Any, Any], Any]] = { + "=": operator.eq, + "!=": operator.ne, + "<": operator.lt, + "<=": operator.le, + ">": operator.gt, + ">=": operator.ge, +} + +_INTEGERS = { + "Int8": "TINYINT", + "Int16": "SMALLINT", + "Int32": "INTEGER", + "Int64": "BIGINT", + "UInt8": "UTINYINT", + "UInt16": "USMALLINT", + "UInt32": "UINTEGER", + "UInt64": "UBIGINT", +} + +_MICROS_PER = {"ms": 1_000, "us": 1} + +_NAIVE_TIMESTAMPS = {"ms": "TIMESTAMP_MS", "us": "TIMESTAMP"} + + +def _on_grid(count: int, unit: str) -> int | None: + """A count of microseconds as a count of `unit`, or None when it lies between two of them.""" + step = _MICROS_PER[unit] + return count // step if count % step == 0 else None + + +def to_polars(node: Node, schema: pl.Schema) -> pl.Expr: + """`node` over `schema`, the schema of the frame, or `Untranslatable` where polars would differ.""" + try: + return _translate(node, schema) + except RecursionError as error: + raise Untranslatable(node) from error + + +def _translate(node: Node, schema: pl.Schema) -> pl.Expr: + import polars as pl + + match node: + case Binary(op="AND", left=left, right=right): + return _translate(left, schema) & _translate(right, schema) + case Binary(op="OR", left=left, right=right): + return _translate(left, schema) | _translate(right, schema) + case Unary(op="NOT", operand=operand): + return ~_translate(operand, schema) + case Postfix(op="IS NULL", operand=Col(parts=parts)): + return _column(schema, parts)[0].is_null() + case Postfix(op="IS NOT NULL", operand=Col(parts=parts)): + return _column(schema, parts)[0].is_not_null() + case Binary(op=op, left=Col(parts=parts), right=Lit() as constant) if op in _COMPARISONS: + column, kind = _column(schema, parts) + if _categories(kind): + if op not in ("=", "!="): + # Polars orders a Categorical or Enum by its categories, where the engine orders the strings. + raise Untranslatable(node) + column, kind = column.cast(pl.String), pl.String + literal = _literal(constant, kind) + if constant.type == "TIMESTAMP WITH TIME ZONE": + # Polars casts a literal down to a coarser column's unit, so the column goes up to microseconds instead. + column = column.dt.cast_time_unit("us") + compared: pl.Expr = _COMPARISONS[op](column, literal) + return compared + case In(operand=Col(parts=parts), values=values) if values: + column, kind = _column(schema, parts) + if _categories(kind): + column, kind = column.cast(pl.String), pl.String + member = _membership(column, kind, values) + # Whether is_in answers NULL for a NULL row depends on the polars version; the engine answers NULL. + return pl.when(column.is_null()).then(pl.lit(None, dtype=pl.Boolean)).otherwise(member) + raise Untranslatable(node) + + +def _categories(kind: object) -> bool: + """Whether `kind` is a Categorical or Enum, which the engine reads as VARCHAR, so its values compare as strings.""" + import polars as pl + + return isinstance(kind, pl.Categorical | pl.Enum) + + +def _engine_type(kind: Any) -> str | None: # noqa: ANN401 + """The engine's name for the type it reads a polars dtype as, for the dtypes translated here.""" + import polars as pl + + if kind == pl.Boolean: + return "BOOLEAN" + if kind.is_integer(): + return _INTEGERS.get(str(kind)) + if kind == pl.Float32: + return "FLOAT" + if kind == pl.Float64: + return "DOUBLE" + if kind == pl.String: + return "VARCHAR" + if kind == pl.Binary: + return "BLOB" + if kind == pl.Date: + return "DATE" + if isinstance(kind, pl.Datetime) and kind.time_zone is None: + return _NAIVE_TIMESTAMPS.get(kind.time_unit) + if isinstance(kind, pl.Datetime) and kind.time_unit in _NAIVE_TIMESTAMPS: + return "TIMESTAMP WITH TIME ZONE" + if isinstance(kind, pl.Decimal): + return f"DECIMAL({_precision(kind)},{kind.scale})" + return None + + +def _precision(kind: pl.Decimal) -> int: + """The precision a Decimal dtype exports with, which is the widest when polars leaves it open.""" + return kind.precision if kind.precision is not None else 38 + + +def _column(schema: pl.Schema, parts: tuple[str, ...]) -> tuple[pl.Expr, Any]: + """The column expression for `parts` and its dtype; a constant is checked against the dtype where it is used.""" + import polars as pl + + if not parts or parts[0] not in schema: + raise Untranslatable(parts) + # Polars reads a name like ^x$ or * as a selector, so the column is taken by position. + column = pl.nth(schema.names().index(parts[0])) + kind: Any = schema[parts[0]] + for name in parts[1:]: + fields = [field for field in kind.fields if field.name == name] if isinstance(kind, pl.Struct) else [] + # A struct field has no positional access polars can plan, so a field named like a selector is refused. + if len(fields) != 1 or name == "*" or (name.startswith("^") and name.endswith("$")): + raise Untranslatable(parts) + column = column.struct.field(name) + kind = fields[0].dtype + return column, kind + + +def _literal(constant: Lit, kind: Any) -> pl.Expr: # noqa: ANN401 + """A literal of the constant's own engine type, refusing a constant that is not a value of the column's type.""" + import polars as pl + + if constant.type != _engine_type(kind) or not exact_value(constant): + raise Untranslatable(constant) + if constant.type == "TIMESTAMP WITH TIME ZONE": + return pl.lit(constant.value, dtype=pl.Int64).cast(pl.Datetime("us", kind.time_zone)) + return pl.lit(constant.value, dtype=_own_dtype(constant.type, kind)) + + +def _own_dtype(type_text: str, kind: Any) -> Any: # noqa: ANN401 + import polars as pl + + match type_text: + case "BOOLEAN": + return pl.Boolean + case "FLOAT": + return pl.Float32 + case "DOUBLE": + return pl.Float64 + case "VARCHAR": + return pl.String + case "BLOB": + return pl.Binary + case "DATE": + return pl.Date + case "TIMESTAMP_MS": + return pl.Datetime("ms") + case "TIMESTAMP": + return pl.Datetime("us") + # The integer and decimal types match the column's dtype one to one. + return kind + + +def _membership(column: pl.Expr, kind: Any, values: Sequence[Node]) -> pl.Expr: # noqa: ANN401 + """Membership of `column` in the constants, each converted exactly to the column's own dtype.""" + import polars as pl + + held = [] + for constant in values: + if not isinstance(constant, Lit): + raise Untranslatable(constant) + _literal(constant, kind) + value = constant.value + if isinstance(kind, pl.Datetime) and kind.time_zone is not None: + value = _on_grid(value, kind.time_unit) # type: ignore[arg-type] + if value is not None: + held.append(value) + if not held: + return pl.lit(False) + if isinstance(kind, pl.Datetime) and kind.time_zone is not None: + candidates = pl.Series(held, dtype=pl.Int64).cast(kind) + else: + candidates = pl.Series(held, dtype=kind) + return column.is_in(candidates.implode()) diff --git a/src/arrow/CMakeLists.txt b/src/arrow/CMakeLists.txt index da8185a0..4631afcd 100644 --- a/src/arrow/CMakeLists.txt +++ b/src/arrow/CMakeLists.txt @@ -1,7 +1,7 @@ # this is used for clang-tidy checks add_library( python_arrow OBJECT - arrow_array_stream.cpp arrow_export_utils.cpp filter_pushdown_visitor.cpp - polars_filter_pushdown.cpp pyarrow_filter_pushdown.cpp) + arrow_export_utils.cpp arrow_importer.cpp predicate_translation.cpp + python_arrow_scan.cpp single_use.cpp) target_link_libraries(python_arrow PRIVATE _duckdb_dependencies) diff --git a/src/arrow/arrow_array_stream.cpp b/src/arrow/arrow_array_stream.cpp deleted file mode 100644 index 04299941..00000000 --- a/src/arrow/arrow_array_stream.cpp +++ /dev/null @@ -1,417 +0,0 @@ -#include "duckdb_python/arrow/arrow_array_stream.hpp" -#include "duckdb_python/arrow/polars_filter_pushdown.hpp" -#include "duckdb_python/arrow/pyarrow_filter_pushdown.hpp" - -#include "duckdb_python/pyconnection/pyconnection.hpp" -#include "duckdb_python/pyrelation.hpp" -#include "duckdb_python/pyresult.hpp" -#include "duckdb/function/table/arrow.hpp" - -#include "duckdb/common/assert.hpp" -#include "duckdb/common/common.hpp" -#include "duckdb/common/limits.hpp" -#include "duckdb/common/unordered_set.hpp" -#include "duckdb/common/vector.hpp" -#include "duckdb/main/client_config.hpp" - -namespace duckdb { - -static constexpr const char *CONSUMED_READER_ERROR = - "This RecordBatchReader has already been consumed and cannot be scanned again. A RecordBatchReader can be read " - "only once; to scan its data more than once, read it into a table first, for example with read_all()."; - -namespace { - -string TypeText(nb::object value) { - if (!value.is_valid() || !nb::str::check_(value)) { - PyErr_Clear(); - return string(); - } - return nb::cast(value); -} - -//! The fallback for a reader whose pyarrow left sys.modules: an extension type's module and name come from tp_name, -//! while a Python class can forge both through __module__ and __qualname__, so it is never the reader. -bool IsPyarrowReaderType(PyTypeObject *type) { - if (PyType_GetFlags(type) & Py_TPFLAGS_HEAPTYPE) { - return false; - } - auto name = TypeText(nb::steal(PyType_GetName(type))); - auto module = TypeText(nb::steal(PyObject_GetAttrString(reinterpret_cast(type), "__module__"))); - return name == "RecordBatchReader" && module == "pyarrow.lib"; -} - -//! Walks the bases the type was created with, not __mro__, which a metaclass can replace; each ancestor once, since -//! a diamond hierarchy repeats them -bool DerivesFromPyarrowReader(PyTypeObject *type) { - vector pending {type}; - unordered_set seen; - while (!pending.empty()) { - auto current = pending.back(); - pending.pop_back(); - if (!seen.insert(current).second) { - continue; - } - if (IsPyarrowReaderType(current)) { - return true; - } - auto bases = static_cast(PyType_GetSlot(current, Py_tp_bases)); - if (!bases || !PyTuple_Check(bases)) { - PyErr_Clear(); - continue; - } - auto count = PyTuple_Size(bases); - for (Py_ssize_t i = 0; i < count; i++) { - pending.push_back(reinterpret_cast(PyTuple_GetItem(bases, i))); - } - } - return false; -} - -} // namespace - -bool PythonTableArrowArrayStreamFactory::IsSingleUse(nb::handle arrow_object, PyArrowObjectType type) { - if (type != PyArrowObjectType::PyCapsuleInterface) { - return false; - } - if (ModuleIsLoaded()) { - auto reader_type = DuckDBPyConnection::ImportCache()->pyarrow.RecordBatchReader(); - if (reader_type.ptr() && PyType_Check(reader_type.ptr())) { - return PyType_IsSubtype(Py_TYPE(arrow_object.ptr()), reinterpret_cast(reader_type.ptr())) != - 0; - } - } - return DerivesFromPyarrowReader(Py_TYPE(arrow_object.ptr())); -} - -namespace { - -//! Records a single-use reader as consumed, by identity, and takes the record back unless Keep() is called -class ConsumedReaderMark { -public: - ConsumedReaderMark(nb::handle reader, bool single_use) { - if (!single_use) { - return; - } - auto readers = DuckDBPyConnection::ConsumedArrowReaders(); - key = nb::steal(PyLong_FromVoidPtr(reader.ptr())); - if (!key) { - throw nb::python_error(); - } - auto forget = nb::module_::import_("functools").attr("partial")(readers.attr("pop"), key); - auto entry = nb::steal(PyWeakref_NewRef(reader.ptr(), forget.ptr())); - if (!entry) { - throw nb::python_error(); - } - // One setdefault tests and records at once; the entry leaves with the reader, before its id can be reused. - auto recorded = readers.attr("setdefault")(key, entry); - if (!recorded.is(entry)) { - throw InvalidInputException(CONSUMED_READER_ERROR); - } - record = std::move(readers); - } - ~ConsumedReaderMark() { - if (record && !kept && PyDict_DelItem(record.ptr(), key.ptr()) < 0) { - PyErr_Clear(); - } - } - ConsumedReaderMark(const ConsumedReaderMark &) = delete; - ConsumedReaderMark &operator=(const ConsumedReaderMark &) = delete; - - void Keep() { - kept = true; - } - -private: - //! Set only once the reader is marked - nb::object record; - nb::object key; - bool kept = false; -}; - -} // namespace - -void TransformDuckToArrowChunk(nb::object pyarrow_schema, ArrowArray &data, nb::list &batches) { - DUCKDB_PY_ASSERT_GIL_HELD(); - auto pyarrow_lib_module = nb::module_::import_("pyarrow").attr("lib"); - auto batch_import_func = pyarrow_lib_module.attr("RecordBatch").attr("_import_from_c"); - batches.append(batch_import_func(reinterpret_cast(&data), pyarrow_schema)); -} - -void VerifyArrowDatasetLoaded() { - auto &import_cache = *DuckDBPyConnection::ImportCache(); - if (!import_cache.pyarrow.dataset() || !ModuleIsLoaded()) { - throw InvalidInputException("Optional module 'pyarrow.dataset' is required to perform this action"); - } -} - -nb::object PythonTableArrowArrayStreamFactory::ProduceScanner(nb::object &arrow_scanner, nb::handle &arrow_obj_handle, - ArrowStreamParameters ¶meters, - const ClientProperties &client_properties) { - D_ASSERT(!nb::isinstance(arrow_obj_handle)); - ArrowSchemaWrapper schema; - PythonTableArrowArrayStreamFactory::GetSchemaInternal(arrow_obj_handle, schema.arrow_schema); - ArrowTableSchema arrow_table; - ArrowTableFunction::PopulateArrowTableSchema(*client_properties.client_context.get_mutable(), arrow_table, - schema.arrow_schema); - - auto filters = parameters.filters; - auto &column_list = parameters.projected_columns.columns; - auto &filter_to_col = parameters.projected_columns.filter_to_col; - nb::list projection_list(nb::cast(column_list)); - - bool has_filter = filters && filters->HasFilters(); - nb::dict kwargs; - if (!column_list.empty()) { - kwargs["columns"] = projection_list; - } - - if (has_filter) { - auto filter = - PyArrowFilterPushdown::TransformFilter(*filters, parameters.projected_columns.projection_map, filter_to_col, - client_properties, arrow_table, schema.arrow_schema); - if (!filter.is(nb::none())) { - kwargs["filter"] = filter; - } - } - return arrow_scanner(arrow_obj_handle, **kwargs); -} - -unique_ptr -PythonTableArrowArrayStreamFactory::ProduceStream(ArrowStreamParameters ¶meters) { - nb::gil_scoped_acquire acquire; - D_ASSERT(arrow_object.obj.ptr()); - nb::handle arrow_obj_handle(arrow_object.obj); - auto arrow_object_type = cached_arrow_type; - - if (arrow_object_type == PyArrowObjectType::PolarsLazyFrame) { - nb::object lf = nb::borrow(arrow_obj_handle); - - auto filters = parameters.filters; - bool filters_pushed = false; - - // Translate DuckDB filters to Polars expressions and push into the lazy plan. - // The walker only fails (throws / returns nb::none()) for filters that are not - // required for correctness — optional/runtime wrappers it skips, or shapes the - // optimizer keeps above the scan. A throw here would mean the optimizer fully - // pushed something we can't translate (a correctness bug), so we let it surface - // rather than silently returning unfiltered rows — the arrow scan does not - // re-apply pushed filters. Mirrors the pyarrow ProduceScanner path. - if (filters && filters->HasFilters()) { - auto filter_expr = - PolarsFilterPushdown::TransformFilter(*filters, parameters.projected_columns.projection_map, - parameters.projected_columns.filter_to_col, client_properties); - if (!filter_expr.is(nb::none())) { - lf = lf.attr("filter")(filter_expr); - filters_pushed = true; - } - } - - // If no filters were pushed and we have a cached Arrow table, reuse it. This avoids re-reading from source and - // re-converting on repeated unfiltered scans. - nb::object arrow_table; - if (!filters_pushed && cached_arrow_table.obj.ptr() != nullptr) { - arrow_table = cached_arrow_table.obj; - } else { - arrow_table = lf.attr("collect")().attr("to_arrow")(); - // Cache only unfiltered results (filtered results are partial) - if (!filters_pushed) { - cached_arrow_table.obj = arrow_table; - } - } - - // Apply column projection - auto &column_list = parameters.projected_columns.columns; - if (!column_list.empty()) { - arrow_table = arrow_table.attr("select")(nb::cast(column_list)); - } - - auto capsule_obj = arrow_table.attr("__arrow_c_stream__")(); - auto capsule = nb::borrow(capsule_obj); - auto stream = reinterpret_cast(capsule.data()); - auto res = make_uniq(); - res->arrow_array_stream = *stream; - stream->release = nullptr; - return res; - } - - if (arrow_object_type == PyArrowObjectType::PyCapsuleInterface || arrow_object_type == PyArrowObjectType::Table) { - ConsumedReaderMark mark(arrow_obj_handle, single_use); - nb::object capsule_obj = arrow_obj_handle.attr("__arrow_c_stream__")(); - mark.Keep(); - auto capsule = nb::borrow(capsule_obj); - auto stream = reinterpret_cast(capsule.data()); - if (!stream->release) { - throw InvalidInputException( - "The __arrow_c_stream__() method returned a released stream. " - "If this object is single-use, implement __arrow_c_schema__() or expose a .schema attribute " - "with _export_to_c() so that DuckDB can extract the schema without consuming the stream."); - } - - auto &import_cache_check = *DuckDBPyConnection::ImportCache(); - if (import_cache_check.pyarrow.dataset()) { - // Tier A: full pushdown via pyarrow.dataset - // Import as RecordBatchReader, feed through Scanner.from_batches for projection/filter pushdown. - auto pyarrow_lib_module = nb::module_::import_("pyarrow").attr("lib"); - auto import_func = pyarrow_lib_module.attr("RecordBatchReader").attr("_import_from_c"); - nb::object reader = import_func(reinterpret_cast(stream)); - // _import_from_c takes ownership of the stream; null out to prevent capsule double-free - stream->release = nullptr; - auto &import_cache = *DuckDBPyConnection::ImportCache(); - nb::object arrow_batch_scanner = import_cache.pyarrow.dataset.Scanner().attr("from_batches"); - nb::handle reader_handle = reader; - auto scanner = ProduceScanner(arrow_batch_scanner, reader_handle, parameters, client_properties); - auto record_batches = scanner.attr("to_reader")(); - auto res = make_uniq(); - auto export_to_c = record_batches.attr("_export_to_c"); - export_to_c(reinterpret_cast(&res->arrow_array_stream)); - return res; - } else { - // Tier B: no pyarrow.dataset, return raw stream (no pushdown) - // DuckDB applies projection/filter post-scan via arrow_scan_dumb - auto res = make_uniq(); - res->arrow_array_stream = *stream; - stream->release = nullptr; - return res; - } - } - - if (arrow_object_type == PyArrowObjectType::PyCapsule) { - auto res = make_uniq(); - auto capsule = nb::borrow(arrow_obj_handle); - auto stream = reinterpret_cast(capsule.data()); - if (!stream->release) { - throw InvalidInputException("This ArrowArrayStream has already been consumed and cannot be scanned again."); - } - res->arrow_array_stream = *stream; - stream->release = nullptr; - return res; - } - - // Scanner and Dataset: require pyarrow.dataset for pushdown - VerifyArrowDatasetLoaded(); - auto &import_cache = *DuckDBPyConnection::ImportCache(); - nb::object scanner; - nb::object arrow_batch_scanner = import_cache.pyarrow.dataset.Scanner().attr("from_batches"); - switch (arrow_object_type) { - case PyArrowObjectType::Scanner: { - // If it's a scanner we have to turn it to a record batch reader, and then a scanner again since we can't stack - // scanners on arrow Otherwise pushed-down projections and filters will disappear like tears in the rain - auto record_batches = arrow_obj_handle.attr("to_reader")(); - scanner = ProduceScanner(arrow_batch_scanner, record_batches, parameters, client_properties); - break; - } - case PyArrowObjectType::Dataset: { - nb::object arrow_scanner = arrow_obj_handle.attr("__class__").attr("scanner"); - scanner = ProduceScanner(arrow_scanner, arrow_obj_handle, parameters, client_properties); - break; - } - default: { - // nb::object wrap: nb::str() of a bare .attr() accessor is an ambiguous overload on MSVC. - auto py_object_type = nb::cast(nb::str(nb::object((arrow_obj_handle).type().attr("__name__")))); - throw InvalidInputException("Object of type '%s' is not a recognized Arrow object", py_object_type); - } - } - - auto record_batches = scanner.attr("to_reader")(); - auto res = make_uniq(); - auto export_to_c = record_batches.attr("_export_to_c"); - export_to_c(reinterpret_cast(&res->arrow_array_stream)); - return res; -} - -void PythonTableArrowArrayStreamFactory::GetSchemaInternal(nb::handle arrow_obj_handle, ArrowSchema &schema) { - // PyCapsule (from bare capsule Produce path) - if (nb::isinstance(arrow_obj_handle)) { - auto capsule = nb::borrow(arrow_obj_handle); - auto stream = reinterpret_cast(capsule.data()); - if (!stream->release) { - throw InvalidInputException("This ArrowArrayStream has already been consumed and cannot be scanned again."); - } - if (stream->get_schema(stream, &schema)) { - throw InvalidInputException("Failed to get Arrow schema from stream: %s", - stream->get_last_error ? stream->get_last_error(stream) : "unknown error"); - } - return; - } - - // Scanner: use projected_schema; everything else (RecordBatchReader, Dataset): use .schema - VerifyArrowDatasetLoaded(); - auto &import_cache = *DuckDBPyConnection::ImportCache(); - if (duckdb::PyUtil::IsInstance(arrow_obj_handle, import_cache.pyarrow.dataset.Scanner())) { - auto obj_schema = arrow_obj_handle.attr("projected_schema"); - obj_schema.attr("_export_to_c")(reinterpret_cast(&schema)); - } else { - auto obj_schema = arrow_obj_handle.attr("schema"); - obj_schema.attr("_export_to_c")(reinterpret_cast(&schema)); - } -} - -void PythonTableArrowArrayStreamFactory::GetSchema(ArrowSchema &schema) { - if (schema_cached.load(std::memory_order_acquire)) { - schema = cached_schema; - schema.release = nullptr; - return; - } - - nb::gil_scoped_acquire acquire; - D_ASSERT(arrow_object.obj.ptr()); - nb::handle arrow_obj_handle(arrow_object.obj); - - auto type = cached_arrow_type; - if (type == PyArrowObjectType::PolarsLazyFrame) { - // head(0).collect().to_arrow() gives the Arrow-exported schema (e.g. large_string) without materializing data. - // collect_schema() would give Polars-native types (e.g. string_view) that don't match the actual export. - const auto empty_arrow = arrow_obj_handle.attr("head")(0).attr("collect")().attr("to_arrow")(); - const auto schema_capsule = empty_arrow.attr("schema").attr("__arrow_c_schema__")(); - const auto capsule = nb::borrow(schema_capsule); - const auto arrow_schema = reinterpret_cast(capsule.data()); - cached_schema = *arrow_schema; - arrow_schema->release = nullptr; - schema_cached.store(true, std::memory_order_release); - schema = cached_schema; - schema.release = nullptr; - return; - } - if (type == PyArrowObjectType::PyCapsuleInterface || type == PyArrowObjectType::Table) { - // Get __arrow_c_schema__ if it exists - if (nb::hasattr(arrow_obj_handle, "__arrow_c_schema__")) { - auto schema_capsule = arrow_obj_handle.attr("__arrow_c_schema__")(); - auto capsule = nb::borrow(schema_capsule); - auto arrow_schema = reinterpret_cast(capsule.data()); - cached_schema = *arrow_schema; - arrow_schema->release = nullptr; - schema_cached.store(true, std::memory_order_release); - schema = cached_schema; - schema.release = nullptr; - return; - } - // Otherwise try to use .schema with _export_to_c - if (nb::hasattr(arrow_obj_handle, "schema")) { - auto obj_schema = arrow_obj_handle.attr("schema"); - if (nb::hasattr(obj_schema, "_export_to_c")) { - obj_schema.attr("_export_to_c")(reinterpret_cast(&schema)); - return; - } - } - // Fallback: create a temporary stream just for the schema (consumes single-use streams!) - auto stream_capsule = arrow_obj_handle.attr("__arrow_c_stream__")(); - auto capsule = nb::borrow(stream_capsule); - auto stream = reinterpret_cast(capsule.data()); - if (stream->get_schema(stream, &schema)) { - throw InvalidInputException("Failed to get Arrow schema from stream: %s", - stream->get_last_error ? stream->get_last_error(stream) : "unknown error"); - } - return; // stream_capsule goes out of scope, stream released by capsule destructor - } - GetSchemaInternal(arrow_obj_handle, schema); - - // Cache for Table and Dataset (immutable schema) - if (type == PyArrowObjectType::Table || type == PyArrowObjectType::Dataset) { - cached_schema = schema; - schema.release = nullptr; - schema_cached.store(true, std::memory_order_release); - } -} - -} // namespace duckdb diff --git a/src/arrow/arrow_export_utils.cpp b/src/arrow/arrow_export_utils.cpp index 55ef48e7..d2d3da12 100644 --- a/src/arrow/arrow_export_utils.cpp +++ b/src/arrow/arrow_export_utils.cpp @@ -1,17 +1,7 @@ -#include "duckdb_python/arrow/arrow_array_stream.hpp" - -#include "duckdb/common/assert.hpp" -#include "duckdb/common/common.hpp" -#include "duckdb/common/limits.hpp" -#include "duckdb/main/client_config.hpp" -#include "duckdb/planner/filter/conjunction_filter.hpp" -#include "duckdb/planner/filter/constant_filter.hpp" -#include "duckdb/planner/table_filter.hpp" -#include "duckdb/common/arrow/arrow_converter.hpp" +#include "duckdb_python/arrow/arrow_export_utils.hpp" -#include "duckdb_python/pyconnection/pyconnection.hpp" -#include "duckdb_python/pyrelation.hpp" -#include "duckdb_python/pyresult.hpp" +#include "duckdb/common/arrow/arrow_converter.hpp" +#include "duckdb_python/pyutil.hpp" namespace duckdb { @@ -43,4 +33,11 @@ nb::object ToArrowTable(const vector &types, const vector & } // namespace pyarrow +void TransformDuckToArrowChunk(nb::object pyarrow_schema, ArrowArray &data, nb::list &batches) { + DUCKDB_PY_ASSERT_GIL_HELD(); + auto pyarrow_lib_module = nb::module_::import_("pyarrow").attr("lib"); + auto batch_import_func = pyarrow_lib_module.attr("RecordBatch").attr("_import_from_c"); + batches.append(batch_import_func(reinterpret_cast(&data), pyarrow_schema)); +} + } // namespace duckdb diff --git a/src/arrow/arrow_importer.cpp b/src/arrow/arrow_importer.cpp new file mode 100644 index 00000000..555ea342 --- /dev/null +++ b/src/arrow/arrow_importer.cpp @@ -0,0 +1,335 @@ +#include "duckdb_python/arrow/arrow_importer.hpp" + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/limits.hpp" +#include "duckdb/common/numeric_utils.hpp" +#include "duckdb/common/vector/flat_vector.hpp" +#include "duckdb/function/table/arrow/arrow_type_info.hpp" + +namespace duckdb { + +namespace { + +//! The rows the batch's own validity bitmap calls NULL, from `from` on; null when it has no bitmap or no NULL row. +unique_ptr BatchValidity(const ArrowArray &batch, idx_t from, idx_t rows) { + if (batch.null_count == 0 || batch.n_buffers == 0 || !batch.buffers || !batch.buffers[0]) { + return nullptr; + } + auto bits = static_cast(batch.buffers[0]); + auto mask = make_uniq(rows); + for (idx_t row = 0; row < rows; row++) { + auto index = NumericCast(batch.offset) + from + row; + if (!(bits[index / 8] & (1 << (index % 8)))) { + mask->SetInvalid(row); + } + } + if (!mask->CanHaveNull()) { + return nullptr; + } + return mask; +} + +bool ValidRange(int64_t offset, int64_t length) { + return offset >= 0 && length >= 0 && offset <= NumericLimits::Maximum() - length; +} + +//! Checks the arrays below a requested column as the batch is checked, and refuses an offset the engine would drop: its +//! struct conversion hands its children only the struct's own offset and its list conversion none, so an offset met +//! further up would be lost there and the rows below read from the wrong places. +//! `received` is the offset the engine passes the array as its parent's, and `in_list` whether a list above it makes +//! the engine ignore parent offsets. +void CheckNested(const ArrowArray &array, const ArrowType &type, idx_t column, int64_t received, bool in_list) { + if (!ValidRange(array.offset, array.length)) { + throw InvalidInputException("an Arrow batch carried column %d with a nested array of offset %d and length %d, " + "which do not describe a range of rows", + column, array.offset, array.length); + } + // The engine reads validity only through the bitmap, so NULL rows would read as values. + if (array.null_count > 0 && array.n_buffers > 0 && (!array.buffers || !array.buffers[0])) { + throw InvalidInputException( + "an Arrow batch carried column %d with a non-zero null count but no validity bitmap", column); + } + // The engine reads every child its format names, so a missing one would be read through a null pointer. + auto child_of = [&](idx_t index) -> const ArrowArray & { + if (!array.children || array.n_children <= NumericCast(index) || !array.children[index]) { + throw InvalidInputException("an Arrow batch carried column %d with a nested array missing its child %d", + column, index); + } + if (!array.children[index]->release) { + throw InvalidInputException( + "an Arrow batch carried column %d with a nested array whose child %d is released already", column, + index); + } + return *array.children[index]; + }; + // Like the batch: the children of a struct or a sparse union cover its offset plus its length. + auto check_covers = [&](const ArrowArray &child, idx_t factor) { + const auto rows = array.offset + array.length; + const auto width = NumericCast(factor); + if (width > 0 && rows > NumericLimits::Maximum() / width) { + throw InvalidInputException("an Arrow batch carried column %d with a nested array too long to address", + column); + } + const auto needed = rows * width; + if (child.length < needed) { + throw InvalidInputException("Arrow array child length does not match the array length: column %d holds a " + "nested array of %d rows where %d were expected", + column, child.length, needed); + } + }; + switch (type.GetPhysicalType()) { + case ArrowArrayPhysicalType::DICTIONARY_ENCODED: + if (!array.dictionary) { + throw InvalidInputException("an Arrow batch carried column %d with a dictionary-encoded array that has no " + "dictionary", + column); + } + if (!array.dictionary->release) { + throw InvalidInputException( + "an Arrow batch carried column %d with a dictionary-encoded array whose dictionary is released already", + column); + } + // The dictionary is converted on its own, from its start. + CheckNested(*array.dictionary, type.GetDictionary(), column, 0, false); + return; + case ArrowArrayPhysicalType::RUN_END_ENCODED: { + auto &info = type.GetTypeInfo(); + auto &run_ends = child_of(0); + auto &values = child_of(1); + if (run_ends.length != values.length) { + throw InvalidInputException("an Arrow batch carried column %d with a run end encoded array of %d run ends " + "and %d values", + column, run_ends.length, values.length); + } + // The run ends and the values are converted on their own, from their start. + CheckNested(run_ends, info.GetChild(0), column, 0, false); + CheckNested(values, info.GetChild(1), column, 0, false); + return; + } + default: + break; + } + switch (type.GetDuckType().id()) { + // A VARIANT is converted as the struct it is stored as. + case LogicalTypeId::VARIANT: + case LogicalTypeId::STRUCT: + case LogicalTypeId::TUPLE: + case LogicalTypeId::UNION: { + const bool is_union = type.GetDuckType().id() == LogicalTypeId::UNION; + // A union adds its own offset to the one it received, a struct passes on only its own. + if ((in_list && array.offset != 0) || (!is_union && !in_list && received != 0)) { + throw InvalidInputException( + "an Arrow batch carried column %d with a struct or union below an offset the conversion would drop, " + "which would read the wrong rows", + column); + } + const auto passed = is_union ? received + array.offset : array.offset; + auto &info = type.GetTypeInfo(); + // The engine converts every child the array has into vectors sized by the type. + if (array.n_children != NumericCast(info.ChildCount())) { + throw InvalidInputException("an Arrow batch carried column %d with a struct or union of %d child arrays " + "where its type has %d", + column, array.n_children, info.ChildCount()); + } + for (idx_t i = 0; i < info.ChildCount(); i++) { + auto &child = child_of(i); + check_covers(child, 1); + CheckNested(child, info.GetChild(i), column, passed, in_list); + } + return; + } + case LogicalTypeId::LIST: + case LogicalTypeId::MAP: { + auto &info = type.GetTypeInfo(); + auto &child = child_of(0); + auto value_at = [&](idx_t buffer, int64_t index) -> int64_t { + if (info.GetSizeType() == ArrowVariableSizeType::SUPER_SIZE) { + return static_cast(array.buffers[buffer])[index]; + } + return static_cast(array.buffers[buffer])[index]; + }; + // When the lists address no child values and start at offset 0, the engine reads the child as plain values, + // which a dictionary-encoded child is not; a list view counts as starting at 0 whenever its lists are empty. + const bool dictionary_child = info.GetChild().GetPhysicalType() == ArrowArrayPhysicalType::DICTIONARY_ENCODED; + bool only_empty_from_zero = false; + if (!info.IsView() && array.length > 0) { + if (array.n_buffers < 2 || !array.buffers || !array.buffers[1]) { + throw InvalidInputException("an Arrow batch carried column %d with a list that has no offsets", column); + } + // The engine takes the child rows to read from the first and the last offset. + const auto first = value_at(1, array.offset); + const auto last = value_at(1, array.offset + array.length); + if (first < 0 || first > last || last > child.length) { + throw InvalidInputException( + "an Arrow batch carried column %d with a list whose offsets %d to %d do not " + "fit its child of %d rows", + column, first, last, child.length); + } + only_empty_from_zero = last == 0; + } else if (info.IsView() && array.length > 0 && dictionary_child) { + if (array.n_buffers < 3 || !array.buffers || !array.buffers[2]) { + throw InvalidInputException("an Arrow batch carried column %d with a list view that has no sizes", + column); + } + only_empty_from_zero = true; + for (int64_t row = array.offset; row < array.offset + array.length && only_empty_from_zero; row++) { + only_empty_from_zero = value_at(2, row) == 0; + } + } + if (dictionary_child && only_empty_from_zero) { + throw InvalidInputException("an Arrow batch carried column %d with only empty lists over a " + "dictionary-encoded child, which the conversion cannot read", + column); + } + CheckNested(child, info.GetChild(), column, 0, true); + return; + } + case LogicalTypeId::ARRAY: { + auto &child = child_of(0); + auto &info = type.GetTypeInfo(); + // The engine's fixed size list conversion reads a run end encoded child as plain values. + if (array.length > 0 && info.GetChild().GetPhysicalType() == ArrowArrayPhysicalType::RUN_END_ENCODED) { + throw InvalidInputException("an Arrow batch carried column %d with a run end encoded array inside a fixed " + "size list, which the conversion cannot read", + column); + } + check_covers(child, info.FixedSize()); + CheckNested(child, info.GetChild(), column, 0, true); + return; + } + default: + return; + } +} + +void ApplyBatchValidity(Vector &vector, const ValidityMask &batch_validity, idx_t rows) { + auto &mask = FlatVector::ValidityMutable(vector); + for (idx_t row = 0; row < rows; row++) { + if (!batch_validity.RowIsValid(row)) { + mask.SetInvalid(row); + } + } +} + +} // namespace + +ArrowBatchImporter::ArrowBatchImporter(ClientContext &context, shared_ptr schema_p, + vector columns_p) + : context(context), schema(std::move(schema_p)), arrow_columns(schema->GetColumns()), + columns(std::move(columns_p)) { +} + +void ArrowBatchImporter::Append(ArrowArray &incoming) { + const auto expected = arrow_columns.size(); + if (incoming.n_children < 0 || NumericCast(incoming.n_children) != expected) { + throw InvalidInputException("an Arrow batch carried %d columns where %d were expected", incoming.n_children, + expected); + } + if (!ValidRange(incoming.offset, incoming.length)) { + throw InvalidInputException("an Arrow batch has offset %d and length %d, which do not describe a range of rows", + incoming.offset, incoming.length); + } + if (expected > 0 && !incoming.children) { + throw InvalidInputException("an Arrow batch has no column arrays"); + } + // Columns of a sliced batch keep their full length; the batch offset applies to every one of them. + const auto rows_needed = incoming.offset + incoming.length; + for (idx_t i = 0; i < expected; i++) { + auto child = incoming.children[i]; + if (!child || !child->release) { + throw InvalidInputException("an Arrow batch carried column %d missing or released already", i); + } + if (!ValidRange(child->offset, child->length)) { + throw InvalidInputException("an Arrow batch carried column %d with offset %d and length %d, which do not " + "describe a range of rows", + i, child->offset, child->length); + } + if (child->length < rows_needed) { + throw InvalidInputException("Arrow array child length does not match the array length: column %d holds %d " + "rows where %d were expected", + i, child->length, rows_needed); + } + } + if (incoming.null_count > 0 && (incoming.n_buffers == 0 || !incoming.buffers || !incoming.buffers[0])) { + throw InvalidInputException("an Arrow batch has a non-zero null count but no validity bitmap"); + } + for (auto column : columns) { + CheckNested(*incoming.children[column], *arrow_columns.at(column), column, 0, false); + } + states.clear(); + auto owned = make_shared_ptr(); + owned->arrow_array = incoming; + incoming.release = nullptr; + batch = std::move(owned); + offset = 0; + length = NumericCast(batch->arrow_array.length); +} + +ArrowArrayScanState &ArrowBatchImporter::StateOf(idx_t column) { + auto entry = states.find(column); + if (entry != states.end()) { + return *entry->second; + } + auto state = make_uniq(context); + // Zero-copy vectors keep the batch alive through this. + state->owned_data = batch; + auto &result = *state; + states.emplace(column, std::move(state)); + return result; +} + +bool ArrowBatchImporter::Next(DataChunk &output) { + if (!batch || offset >= length) { + states.clear(); + batch.reset(); + return false; + } + D_ASSERT(output.ColumnCount() == columns.size()); + auto &array = batch->arrow_array; + const auto rows = MinValue(STANDARD_VECTOR_SIZE, length - offset); + // The batch offset is folded into the row position, which the engine hands on to nested children, rather than + // passed as the parent offset, which its struct conversion does not hand on. + const auto position = NumericCast(array.offset) + offset; + output.SetChildCardinality(rows); + auto batch_validity = BatchValidity(array, offset, rows); + for (idx_t i = 0; i < columns.size(); i++) { + const auto column = columns[i]; + auto &child = *array.children[column]; + auto &arrow_type = *arrow_columns.at(column); + auto &state = StateOf(column); + auto &vector = output.data[i]; + switch (arrow_type.GetPhysicalType()) { + case ArrowArrayPhysicalType::DICTIONARY_ENCODED: + if (!child.dictionary) { + throw InvalidInputException("a dictionary-encoded Arrow column has no dictionary"); + } + ArrowToDuckDBConversion::ColumnArrowToDuckDBDictionary(vector, child, position, state, rows, arrow_type, -1, + batch_validity.get()); + break; + case ArrowArrayPhysicalType::RUN_END_ENCODED: + ArrowToDuckDBConversion::ColumnArrowToDuckDBRunEndEncoded(vector, child, position, state, rows, arrow_type, + -1, batch_validity.get()); + break; + case ArrowArrayPhysicalType::DEFAULT: + ArrowToDuckDBConversion::SetValidityMask(vector, child, position, rows, 0, -1); + // Merged before the conversion so a struct or union column passes the NULL rows on to its children. + if (batch_validity) { + ApplyBatchValidity(vector, *batch_validity, rows); + } + ArrowToDuckDBConversion::ColumnArrowToDuckDB(vector, child, position, state, rows, arrow_type, -1, + batch_validity.get()); + break; + default: + throw NotImplementedException("Only default, dictionary and run-end encoded Arrow layouts are supported"); + } + FlatVector::SetSize(vector, rows); + // A conversion may rebuild the validity mask from the column's own buffers. + if (batch_validity && vector.GetVectorType() == VectorType::FLAT_VECTOR) { + ApplyBatchValidity(vector, *batch_validity, rows); + } + } + offset += rows; + return true; +} + +} // namespace duckdb diff --git a/src/arrow/filter_pushdown_visitor.cpp b/src/arrow/filter_pushdown_visitor.cpp deleted file mode 100644 index 385a8540..00000000 --- a/src/arrow/filter_pushdown_visitor.cpp +++ /dev/null @@ -1,260 +0,0 @@ -#include "duckdb_python/arrow/filter_pushdown_visitor.hpp" - -#include "duckdb/function/scalar/struct_utils.hpp" -#include "duckdb/planner/expression/bound_comparison_expression.hpp" -#include "duckdb/planner/expression/bound_conjunction_expression.hpp" -#include "duckdb/planner/expression/bound_constant_expression.hpp" -#include "duckdb/planner/expression/bound_function_expression.hpp" -#include "duckdb/planner/expression/bound_operator_expression.hpp" -#include "duckdb/planner/filter/expression_filter.hpp" -#include "duckdb/planner/filter/table_filter_functions.hpp" - -namespace duckdb { - -namespace { - -bool ValueIsNan(const Value &value) { - if (value.type().id() == LogicalTypeId::FLOAT) { - return Value::IsNan(value.GetValue()); - } - if (value.type().id() == LogicalTypeId::DOUBLE) { - return Value::IsNan(value.GetValue()); - } - return false; -} - -// ResolveColumn accepts a bare BoundReferenceExpression and (nested) -// `struct_extract` chains. Anything else throws NotImplementedException — -// that gives the OPTIONAL_FILTER catch point a chance to swallow it. -struct ResolvedColumn { - vector path; - const ArrowType *leaf_type; - const ArrowSchema *leaf_schema; -}; - -// MakeScalar needs the type of the values, which a dictionary column's ArrowType and ArrowSchema keep apart from its -// indices, and pyarrow has no kernel that compares a run end encoded column. -const ArrowType *ValueType(const ArrowType *type, const vector &path) { - if (type && type->RunEndEncoded()) { - vector names; - for (auto &segment : path) { - names.push_back(segment.GetIdentifierName()); - } - throw NotImplementedException("Cannot push down a filter on run end encoded column \"%s\"", - StringUtil::Join(names, ".")); - } - return type && type->HasDictionary() ? &type->GetDictionary() : type; -} - -const ArrowSchema *ValueSchema(const ArrowSchema *schema) { - return schema && schema->dictionary ? schema->dictionary : schema; -} - -ResolvedColumn ResolveColumn(const Expression &expr, const vector &root_path, const ArrowType *root_type, - const ArrowSchema *root_schema) { - if (expr.GetExpressionClass() == ExpressionClass::BOUND_REF) { - return {root_path, ValueType(root_type, root_path), ValueSchema(root_schema)}; - } - if (expr.GetExpressionClass() != ExpressionClass::BOUND_FUNCTION) { - throw NotImplementedException("Cannot push down arrow scan filter on column-side expression: %s", - ExpressionClassToString(expr.GetExpressionClass())); - } - auto &func = expr.Cast(); - idx_t child_idx; - if (!TryGetStructExtractChildIndex(func, child_idx)) { - throw NotImplementedException("Cannot push down arrow scan filter on column-side function: %s", - ExpressionTypeToString(expr.GetExpressionType())); - } - // Recurse innermost-first so names accumulate root → leaf. - auto inner = ResolveColumn(*func.GetChildren()[0], root_path, root_type, root_schema); - inner.path.push_back(StructType::GetChildName(func.GetChildren()[0]->GetReturnType(), child_idx)); - if (inner.leaf_type) { - inner.leaf_type = ValueType(&inner.leaf_type->GetTypeInfo().GetChild(child_idx), inner.path); - } - if (inner.leaf_schema) { - D_ASSERT(child_idx < NumericCast(inner.leaf_schema->n_children)); - inner.leaf_schema = ValueSchema(inner.leaf_schema->children[child_idx]); - } - return inner; -} - -nb::object EmitCompare(FilterBackend &backend, ExpressionType op, nb::object col, const Value &constant, - const ResolvedColumn &resolved, const string &timezone_config) { - if (ValueIsNan(constant)) { - return backend.NaNCompare(op, std::move(col)); - } - auto scalar = backend.MakeScalar(constant, resolved.leaf_type, resolved.leaf_schema, timezone_config); - // DuckDB orders NaN as the greatest float value, so `nan > finite` and `nan >= finite` are TRUE, while - // IEEE (pyarrow) makes them FALSE. For a finite FLOAT/DOUBLE constant with `>` / `>=`, the plain - // comparison would silently drop NaN column rows the engine keeps (the arrow scan never re-applies - // pushed filters). OR the NaN rows back in so the pushed filter matches DuckDB semantics. `<`, `<=`, - // `=`, `<>` already agree with IEEE, so they are left unchanged. (Idempotent for the polars backend, - // which already treats NaN as greatest, and only reached for float constants so is_nan is always valid.) - // N3: keying is_nan on the CONSTANT's float-ness is safe -- a float constant here implies a float column - // (int/float comparisons are constant-folded to int bounds or wrapped in a non-pushed CAST upstream), so - // col.is_nan() always resolves to a valid pyarrow kernel. - const auto constant_type_id = constant.type().id(); - const bool constant_is_float = - constant_type_id == LogicalTypeId::FLOAT || constant_type_id == LogicalTypeId::DOUBLE; - if (constant_is_float && - (op == ExpressionType::COMPARE_GREATERTHAN || op == ExpressionType::COMPARE_GREATERTHANOREQUALTO)) { - auto compare = backend.Compare(op, col, std::move(scalar)); - auto is_nan = backend.IsNaN(std::move(col)); - return backend.Or(std::move(compare), std::move(is_nan)); - } - return backend.Compare(op, std::move(col), std::move(scalar)); -} - -} // anonymous namespace - -nb::object TransformExpression(const Expression &expression, const vector &column_path, - FilterBackend &backend, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, - const string &timezone_config) { - auto expression_class = expression.GetExpressionClass(); - auto expression_type = expression.GetExpressionType(); - - if (expression_class == ExpressionClass::BOUND_FUNCTION) { - auto &bound_function_expression = expression.Cast(); - if (BoundComparisonExpression::IsComparison(expression_type)) { - auto &left = BoundComparisonExpression::Left(bound_function_expression); - auto &right = BoundComparisonExpression::Right(bound_function_expression); - - optional_ptr column_side; - optional_ptr constant_side; - - if (right.GetExpressionType() == ExpressionType::VALUE_CONSTANT) { - column_side = &left; - constant_side = &right.Cast(); - } else if (left.GetExpressionType() == ExpressionType::VALUE_CONSTANT) { - column_side = &right; - constant_side = &left.Cast(); - expression_type = FlipComparisonExpression(expression_type); - } else { - throw NotImplementedException("Can only push down constant comparisons."); - } - - auto resolved = ResolveColumn(*column_side, column_path, arrow_type, arrow_schema); - auto col = backend.MakeColumnRef(resolved.path); - return EmitCompare(backend, expression_type, std::move(col), constant_side->GetValue(), resolved, - timezone_config); - } - - // Internal table-filter functions. Since the table-filter -> expression-filter - // migration in core, optional / dynamic / bloom / perfect-hash-join / prefix-range - // filters no longer have dedicated TableFilter subtypes. They arrive as scalar - // function wrappers inside the ExpressionFilter expression tree (see - // table_filter_functions.hpp). - const auto &func_name = bound_function_expression.Function().GetName(); - - // OPTIONAL / SELECTIVITY_OPTIONAL wrap a child predicate that lives in `bind_info` - // (their `children` hold only a placeholder column ref). An optional filter is never - // required for correctness, so if its child can't be translated we push nothing for - // it rather than failing the whole scan. - if (func_name == OptionalFilterScalarFun::NAME || func_name == SelectivityOptionalFilterScalarFun::NAME) { - optional_ptr child; - if (bound_function_expression.BindInfo()) { - if (func_name == OptionalFilterScalarFun::NAME) { - child = bound_function_expression.BindInfo() - ->Cast() - .child_filter_expr.get(); - } else { - child = bound_function_expression.BindInfo() - ->Cast() - .child_filter_expr.get(); - } - } - if (!child) { - return nb::none(); - } - try { - return TransformExpression(*child, column_path, backend, arrow_type, arrow_schema, timezone_config); - } catch (const NotImplementedException &) { - return nb::none(); - } - } - - // DYNAMIC / BLOOM / PERFECT_HASH_JOIN / PREFIX_RANGE are runtime filters with no - // static pyarrow/polars equivalent. They are not required for correctness (the - // engine applies them above the scan), so skip them. - if (TableFilterFunctions::IsTableFilterFunction(func_name)) { - return nb::none(); - } - } - - if (expression_class == ExpressionClass::BOUND_OPERATOR) { - auto &op_expr = expression.Cast(); - if (expression_type == ExpressionType::OPERATOR_IS_NULL) { - auto resolved = ResolveColumn(*op_expr.GetChildren()[0], column_path, arrow_type, arrow_schema); - auto col = backend.MakeColumnRef(resolved.path); - return backend.IsNull(std::move(col)); - } - if (expression_type == ExpressionType::OPERATOR_IS_NOT_NULL) { - auto resolved = ResolveColumn(*op_expr.GetChildren()[0], column_path, arrow_type, arrow_schema); - auto col = backend.MakeColumnRef(resolved.path); - return backend.IsNotNull(std::move(col)); - } - if (expression_type == ExpressionType::COMPARE_IN) { - auto resolved = ResolveColumn(*op_expr.GetChildren()[0], column_path, arrow_type, arrow_schema); - auto col = backend.MakeColumnRef(resolved.path); - vector values; - for (idx_t i = 1; i < op_expr.GetChildren().size(); i++) { - auto &const_expr = op_expr.GetChildren()[i]->Cast(); - values.push_back(const_expr.GetValue()); - } - auto col_type = op_expr.GetChildren()[0]->GetReturnType(); - return backend.IsIn(std::move(col), values, col_type, timezone_config); - } - } - - if (expression_class == ExpressionClass::BOUND_CONJUNCTION) { - if (expression_type == ExpressionType::CONJUNCTION_OR || expression_type == ExpressionType::CONJUNCTION_AND) { - const bool is_and = expression_type == ExpressionType::CONJUNCTION_AND; - auto &conj_expr = expression.Cast(); - nb::object result = nb::none(); - for (idx_t i = 0; i < conj_expr.GetChildren().size(); i++) { - nb::object child_expression = TransformExpression(*conj_expr.GetChildren()[i], column_path, backend, - arrow_type, arrow_schema, timezone_config); - if (child_expression.is(nb::none())) { - if (is_and) { - // A conjunct we can't push can simply be dropped: the remaining AND - // terms still form a correct (if weaker) filter, and the engine - // re-applies the rest above the scan. - continue; - } - // An OR branch that can't be translated (e.g. a dynamic filter) would - // make the pushed-down predicate stricter than the engine intends — - // fall back to no pushdown for the whole disjunction. - return nb::none(); - } - if (result.is(nb::none())) { - result = std::move(child_expression); - } else if (is_and) { - result = backend.And(std::move(result), std::move(child_expression)); - } else { - result = backend.Or(std::move(result), std::move(child_expression)); - } - } - return result; - } - } - - throw NotImplementedException("Pushdown Filter Type %s is not currently supported in arrow scans", - ExpressionClassToString(expression_class)); -} - -nb::object TransformFilter(const TableFilter &filter, const vector &column_path, FilterBackend &backend, - const ArrowType *arrow_type, const ArrowSchema *arrow_schema, - const string &timezone_config) { - switch (filter.filter_type) { - case TableFilterType::EXPRESSION_FILTER: { - auto &expression_filter = filter.Cast(); - return TransformExpression(*expression_filter.expr, column_path, backend, arrow_type, arrow_schema, - timezone_config); - } - default: - throw NotImplementedException("Pushdown Filter Type %s is not currently supported in arrow scans", - EnumUtil::ToString(filter.filter_type)); - } -} - -} // namespace duckdb diff --git a/src/arrow/polars_filter_pushdown.cpp b/src/arrow/polars_filter_pushdown.cpp deleted file mode 100644 index f23a5449..00000000 --- a/src/arrow/polars_filter_pushdown.cpp +++ /dev/null @@ -1,157 +0,0 @@ -#include "duckdb_python/arrow/polars_filter_pushdown.hpp" - -#include "duckdb_python/arrow/filter_pushdown_visitor.hpp" -#include "duckdb_python/import_cache/python_import_cache.hpp" -#include "duckdb_python/pyconnection/pyconnection.hpp" -#include "duckdb_python/python_objects.hpp" - -namespace duckdb { - -namespace { - -struct PolarsBackend : public FilterBackend { - explicit PolarsBackend(const ClientProperties &client_properties_p) - : client_properties(client_properties_p), import_cache(*DuckDBPyConnection::ImportCache()) { - } - - nb::object MakeColumnRef(const vector &path) override { - // pl.col(path[0]).struct.field(path[1]).struct.field(...) — polars supports arbitrary - // chaining for nested struct access, verified empirically up to 3 levels. - nb::object col = import_cache.polars.col()(path[0]); - for (idx_t i = 1; i < path.size(); i++) { - col = col.attr("struct").attr("field")(path[i].GetIdentifierName()); - } - return col; - } - - nb::object MakeScalar(const Value &v, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, - const string &timezone_config) override { - // Polars handles type coercion for primitives; no ArrowType lookup is needed. - (void)arrow_type; - (void)arrow_schema; - (void)timezone_config; - return PythonObject::FromValue(v, v.type(), client_properties); - } - - nb::object Compare(ExpressionType op, nb::object col, nb::object scalar) override { - switch (op) { - case ExpressionType::COMPARE_EQUAL: - return col.attr("__eq__")(scalar); - case ExpressionType::COMPARE_NOTEQUAL: - return col.attr("__ne__")(scalar); - case ExpressionType::COMPARE_LESSTHAN: - return col.attr("__lt__")(scalar); - case ExpressionType::COMPARE_GREATERTHAN: - return col.attr("__gt__")(scalar); - case ExpressionType::COMPARE_LESSTHANOREQUALTO: - return col.attr("__le__")(scalar); - case ExpressionType::COMPARE_GREATERTHANOREQUALTO: - return col.attr("__ge__")(scalar); - default: - throw NotImplementedException("Comparison Type %s can't be a polars pushdown filter", - ExpressionTypeToString(op)); - } - } - - nb::object NaNCompare(ExpressionType op, nb::object col) override { - switch (op) { - case ExpressionType::COMPARE_EQUAL: - case ExpressionType::COMPARE_GREATERTHANOREQUALTO: - return col.attr("is_nan")(); - case ExpressionType::COMPARE_LESSTHAN: - case ExpressionType::COMPARE_NOTEQUAL: - return col.attr("is_nan")().attr("__invert__")(); - case ExpressionType::COMPARE_GREATERTHAN: - // Nothing is greater than NaN. - return import_cache.polars.lit()(false); - case ExpressionType::COMPARE_LESSTHANOREQUALTO: - // NaN sorts above every other value, so only NULL rows fail. - return IsNotNull(std::move(col)); - default: - throw NotImplementedException("Unsupported comparison type (%s) for NaN values", - ExpressionTypeToString(op)); - } - } - - nb::object IsNaN(nb::object col) override { - return col.attr("is_nan")(); - } - - nb::object IsNull(nb::object col) override { - return col.attr("is_null")(); - } - - nb::object IsNotNull(nb::object col) override { - return col.attr("is_not_null")(); - } - - nb::object IsIn(nb::object col, const vector &values, const LogicalType &col_logical_type, - const string &timezone_config) override { - (void)timezone_config; - nb::list py_values; - for (auto &val : values) { - py_values.append(PythonObject::FromValue(val, val.type(), client_properties)); - } - if (col_logical_type.id() == LogicalTypeId::DECIMAL) { - // Polars infers Decimal(38, scale) for a plain list of Python Decimal values, - // which doesn't match the column's declared Decimal(precision, scale) — the call - // then fails with `'is_in' cannot check for List(Decimal(38, _)) values in - // Decimal(p, s) data`. Build a typed Series matching the column to side-step - // that, and wrap it with `.implode()` to silence the - // `is_in`-with-same-dtype-Series deprecation (issue 22149). - uint8_t width; - uint8_t scale; - col_logical_type.GetDecimalProperties(width, scale); - nb::object dtype = import_cache.polars.Decimal()(nb::arg("precision") = width, nb::arg("scale") = scale); - nb::object typed_series = - import_cache.polars.Series()(nb::arg("values") = py_values, nb::arg("dtype") = dtype); - return col.attr("is_in")(typed_series.attr("implode")()); - } - return col.attr("is_in")(py_values); - } - - nb::object And(nb::object a, nb::object b) override { - return a.attr("__and__")(b); - } - - nb::object Or(nb::object a, nb::object b) override { - return a.attr("__or__")(b); - } - -private: - const ClientProperties &client_properties; - PythonImportCache &import_cache; -}; - -} // anonymous namespace - -nb::object PolarsFilterPushdown::TransformFilter(const TableFilterSet &filter_collection, - unordered_map &columns, - const unordered_map &filter_to_col, - const ClientProperties &client_properties) { - (void)filter_to_col; - PolarsBackend backend(client_properties); - nb::object expression = nb::none(); - for (auto &entry : filter_collection) { - auto column_idx = entry.GetIndex(); - auto &column_name = columns[column_idx]; - D_ASSERT(columns.find(column_idx) != columns.end()); - - vector column_path = {Identifier(column_name)}; - // Polars does not need ArrowType information — `nullptr` here propagates through the - // shared walker; the PolarsBackend ignores both the ArrowType and the ArrowSchema node in MakeScalar. - nb::object child_expression = duckdb::TransformFilter(entry.Filter(), std::move(column_path), backend, nullptr, - nullptr, client_properties.time_zone); - if (child_expression.is(nb::none())) { - continue; - } - if (expression.is(nb::none())) { - expression = std::move(child_expression); - } else { - expression = expression.attr("__and__")(child_expression); - } - } - return expression; -} - -} // namespace duckdb diff --git a/src/arrow/predicate_translation.cpp b/src/arrow/predicate_translation.cpp new file mode 100644 index 00000000..40f5c2b8 --- /dev/null +++ b/src/arrow/predicate_translation.cpp @@ -0,0 +1,609 @@ +#include "duckdb_python/arrow/predicate_translation.hpp" + +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/common/sql_identifier.hpp" +#include "duckdb/common/types/date.hpp" +#include "duckdb/common/types/time.hpp" +#include "duckdb/common/types/timestamp.hpp" +#include "duckdb/common/types/uuid.hpp" +#include "duckdb/function/scalar/comparison_functions.hpp" +#include "duckdb/function/scalar/struct_functions.hpp" +#include "duckdb/function/scalar/struct_utils.hpp" +#include "duckdb/function/table/arrow.hpp" +#include "duckdb/main/client_context.hpp" +#include "duckdb/main/connection.hpp" +#include "duckdb/optimizer/optimizer.hpp" +#include "duckdb/parser/parsed_data/create_table_info.hpp" +#include "duckdb/planner/expression/bound_between_expression.hpp" +#include "duckdb/planner/expression/bound_columnref_expression.hpp" +#include "duckdb/planner/expression/bound_comparison_expression.hpp" +#include "duckdb/planner/expression/bound_conjunction_expression.hpp" +#include "duckdb/planner/expression/bound_constant_expression.hpp" +#include "duckdb/planner/expression/bound_function_expression.hpp" +#include "duckdb/planner/expression/bound_operator_expression.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/planner/planner.hpp" +#include "duckdb_python/engine/connection.hpp" +#include "duckdb_python/pyconnection/pyconnection.hpp" +#include "duckdb_python/pyutil.hpp" + +#include +#include + +namespace duckdb { + +namespace { + +bool SameFieldNames(const ArrowSchema &field, const LogicalType &type) { + if (type.id() != LogicalTypeId::STRUCT) { + return true; + } + const ArrowSchema &node = field.dictionary ? *field.dictionary : field; + const auto child_count = StructType::GetChildCount(type); + if (!node.format || string(node.format) != "+s" || node.n_children < 0 || + NumericCast(node.n_children) != child_count) { + return false; + } + unordered_set seen; + for (idx_t i = 0; i < child_count; i++) { + const auto &child = *node.children[i]; + const string name = child.name ? child.name : ""; + if (name.empty() || !seen.insert(name).second || + name != StructType::GetChildName(type, i).GetIdentifierName() || + !SameFieldNames(child, StructType::GetChildType(type, i))) { + return false; + } + } + return true; +} + +vector UniqueNames(const vector &columns) { + unordered_map counts; + for (auto &column : columns) { + counts[column.name]++; + } + vector result; + for (auto &column : columns) { + result.push_back(counts[column.name] == 1); + } + return result; +} + +const char *ComparisonOperator(ExpressionType type) { + switch (type) { + case ExpressionType::COMPARE_EQUAL: + return "="; + case ExpressionType::COMPARE_NOTEQUAL: + return "!="; + case ExpressionType::COMPARE_LESSTHAN: + return "<"; + case ExpressionType::COMPARE_LESSTHANOREQUALTO: + return "<="; + case ExpressionType::COMPARE_GREATERTHAN: + return ">"; + case ExpressionType::COMPARE_GREATERTHANOREQUALTO: + return ">="; + default: + throw PredicateRefused(); + } +} + +nb::str PythonText(const string &text) { + return nb::str(text.c_str(), text.size()); +} + +bool PythonYear(int32_t year) { + return year >= 1 && year <= 9999; +} + +timestamp_t Microseconds(const Value &value) { + auto stored = value.GetValueUnsafe(); + if (!Value::IsFinite(stored)) { + throw PredicateRefused(); + } + switch (value.type().id()) { + case LogicalTypeId::TIMESTAMP_SEC: + return Timestamp::FromEpochSeconds(stored.value); + case LogicalTypeId::TIMESTAMP_MS: + return Timestamp::FromEpochMs(stored.value); + default: + return stored; + } +} + +using ScalarFunctionPointer = void (*)(DataChunk &, ExpressionState &, Vector &); + +//! The bind and execution callbacks of one built in struct field accessor overload. +struct StructAccessorCallbacks { + bind_scalar_function_t bind; + ScalarFunctionPointer execute; +}; + +ScalarFunctionPointer ExecutionPointer(const ScalarFunction &function) { + auto target = function.GetCallbacks().function.target(); + return target ? *target : nullptr; +} + +//! The field index of a struct field access, or false when the call is anything else. Another function may carry an +//! accessor's name, and a native one may even reuse its binder, so the call is recognized by both the bind and the +//! execution callback of the definition it was bound from, read from the engine's own accessor definitions. +bool StructFieldIndex(const BoundFunctionExpression &function, idx_t &index) { + static const vector built_in = [] { + vector accessors; + auto struct_extract = StructExtractFun::GetFunctions(); + for (idx_t i = 0; i < struct_extract.Size(); i++) { + auto &overload = *struct_extract.GetFunctionByOffset(i); + accessors.push_back({overload.GetBindCallback(), ExecutionPointer(overload)}); + } + auto struct_extract_at = StructExtractAtFun::GetFunction(); + accessors.push_back({struct_extract_at.GetBindCallback(), ExecutionPointer(struct_extract_at)}); + return accessors; + }(); + auto &name = function.Function().GetName(); + auto &definition = function.Function().GetDefinition(); + if ((name != StructExtractFun::Name && name != StructExtractAtFun::Name) || !definition || + function.GetChildren().size() != 2 || + function.GetChildren()[0]->GetReturnType().id() != LogicalTypeId::STRUCT) { + return false; + } + const auto bind = definition->GetBindCallback(); + const auto execute = ExecutionPointer(*definition); + auto matches = [&](const StructAccessorCallbacks &accessor) { + return accessor.bind && accessor.execute && accessor.bind == bind && accessor.execute == execute; + }; + if (std::none_of(built_in.begin(), built_in.end(), matches)) { + return false; + } + auto bind_data = dynamic_cast(function.BindInfo().get()); + if (!bind_data) { + return false; + } + index = bind_data->index; + return true; +} + +class PredicateWalker { +public: + explicit PredicateWalker(const std::function &resolve) + : resolve(resolve), nodes(nb::borrow(DuckDBPyConnection::ArrowSourcesNodes())) { + } + + nb::object Predicate(const Expression &expression) { + switch (expression.GetExpressionClass()) { + case ExpressionClass::BOUND_CONJUNCTION: + return Conjunction(expression.Cast()); + case ExpressionClass::BOUND_OPERATOR: + return Operator(expression.Cast()); + case ExpressionClass::BOUND_FUNCTION: + return Function(expression.Cast()); + default: + throw PredicateRefused(); + } + } + +private: + struct Column { + nb::object node; + LogicalType type; + }; + + nb::object Conjunction(const BoundConjunctionExpression &conjunction) { + const char *op; + switch (conjunction.GetExpressionType()) { + case ExpressionType::CONJUNCTION_AND: + op = "AND"; + break; + case ExpressionType::CONJUNCTION_OR: + op = "OR"; + break; + default: + throw PredicateRefused(); + } + auto &children = conjunction.GetChildren(); + if (children.empty()) { + throw PredicateRefused(); + } + nb::object folded = Predicate(*children[0]); + for (idx_t i = 1; i < children.size(); i++) { + folded = nodes.attr("Binary")(op, folded, Predicate(*children[i])); + } + return folded; + } + + nb::object Operator(const BoundOperatorExpression &operation) { + auto &children = operation.GetChildren(); + switch (operation.GetExpressionType()) { + case ExpressionType::OPERATOR_NOT: + if (children.size() != 1) { + throw PredicateRefused(); + } + return nodes.attr("Unary")("NOT", Predicate(*children[0])); + case ExpressionType::OPERATOR_IS_NULL: + case ExpressionType::OPERATOR_IS_NOT_NULL: { + if (children.size() != 1) { + throw PredicateRefused(); + } + const char *op = + operation.GetExpressionType() == ExpressionType::OPERATOR_IS_NULL ? "IS NULL" : "IS NOT NULL"; + return nodes.attr("Postfix")(op, ColumnOf(*children[0]).node); + } + case ExpressionType::COMPARE_IN: + case ExpressionType::COMPARE_NOT_IN: { + if (children.size() < 2) { + throw PredicateRefused(); + } + auto column = ColumnOf(*children[0]); + nb::list candidates; + for (idx_t i = 1; i < children.size(); i++) { + candidates.append(Constant(*children[i], column.type)); + } + auto membership = nodes.attr("In")(column.node, candidates); + if (operation.GetExpressionType() == ExpressionType::COMPARE_NOT_IN) { + return nodes.attr("Unary")("NOT", membership); + } + return membership; + } + default: + throw PredicateRefused(); + } + } + + nb::object Function(const BoundFunctionExpression &function) { + auto &children = function.GetChildren(); + if (BoundComparisonExpression::IsComparison(function.GetExpressionType()) && children.size() == 2) { + return Comparison(function.GetExpressionType(), *children[0], *children[1]); + } + if (function.GetExpressionType() == ExpressionType::COMPARE_BETWEEN && + function.Function().GetName() == BetweenFun::Name && BoundBetweenExpression::HasValidBindData(function) && + children.size() == 3) { + return Between(function); + } + throw PredicateRefused(); + } + + nb::object Comparison(ExpressionType type, const Expression &left, const Expression &right) { + if (left.GetExpressionClass() == ExpressionClass::BOUND_CONSTANT && + right.GetExpressionClass() != ExpressionClass::BOUND_CONSTANT) { + return Comparison(FlipComparisonExpression(type), right, left); + } + const char *op = ComparisonOperator(type); + auto column = ColumnOf(left); + return nodes.attr("Binary")(op, column.node, Constant(right, column.type)); + } + + nb::object Between(const BoundFunctionExpression &between) { + const char *lower_op = ComparisonOperator(BoundBetweenExpression::LowerComparisonType(between)); + const char *upper_op = ComparisonOperator(BoundBetweenExpression::UpperComparisonType(between)); + auto column = ColumnOf(BoundBetweenExpression::Input(between)); + auto lower = Constant(BoundBetweenExpression::LowerBound(between), column.type); + auto upper = Constant(BoundBetweenExpression::UpperBound(between), column.type); + nb::object binary = nodes.attr("Binary"); + return binary("AND", binary(lower_op, column.node, lower), binary(upper_op, column.node, upper)); + } + + Column ColumnOf(const Expression &expression) { + vector path; + const Expression *current = &expression; + while (current->GetExpressionClass() == ExpressionClass::BOUND_FUNCTION) { + auto &function = current->Cast(); + idx_t index; + if (!StructFieldIndex(function, index)) { + throw PredicateRefused(); + } + path.push_back(index); + current = function.GetChildren()[0].get(); + } + if (current->GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF) { + throw PredicateRefused(); + } + auto &reference = current->Cast(); + if (reference.Depth() != 0) { + throw PredicateRefused(); + } + auto column = resolve(reference.Binding().column_index.GetIndex()); + if (column.type != reference.GetReturnType()) { + throw PredicateRefused(); + } + nb::list parts; + parts.append(PythonText(column.name)); + auto type = column.type; + for (auto step = path.rbegin(); step != path.rend(); ++step) { + if (type.id() != LogicalTypeId::STRUCT || *step >= StructType::GetChildCount(type)) { + throw PredicateRefused(); + } + auto &name = StructType::GetChildName(type, *step).GetIdentifierName(); + if (name.empty()) { + throw PredicateRefused(); + } + parts.append(PythonText(name)); + type = StructType::GetChildType(type, *step); + } + if (type != expression.GetReturnType()) { + throw PredicateRefused(); + } + return Column {nodes.attr("Col")(nb::tuple(parts)), type}; + } + + nb::object Constant(const Expression &expression, const LogicalType &column_type) { + if (expression.GetExpressionClass() != ExpressionClass::BOUND_CONSTANT) { + throw PredicateRefused(); + } + auto &value = expression.Cast().GetValue(); + // An aliased or extension type, such as JSON, may compare differently from its storage type in the source. + if (value.IsNull() || value.type() != column_type || column_type.HasAlias() || column_type.HasExtensionInfo()) { + throw PredicateRefused(); + } + nb::object python; + try { + python = PythonValue(value); + } catch (nb::python_error &) { + throw PredicateRefused(); + } catch (Exception &) { + throw PredicateRefused(); + } + return nodes.attr("Lit")(python, PythonText(column_type.ToString())); + } + + //! Only types whose every value a Python value holds exactly, and that a source can compare as the engine does. + static nb::object PythonValue(const Value &value) { + switch (value.type().id()) { + case LogicalTypeId::BOOLEAN: + return nb::bool_(value.GetValue()); + case LogicalTypeId::TINYINT: + case LogicalTypeId::SMALLINT: + case LogicalTypeId::INTEGER: + case LogicalTypeId::BIGINT: + return nb::int_(value.GetValue()); + case LogicalTypeId::UTINYINT: + case LogicalTypeId::USMALLINT: + case LogicalTypeId::UINTEGER: + case LogicalTypeId::UBIGINT: + return nb::int_(value.GetValue()); + case LogicalTypeId::HUGEINT: + case LogicalTypeId::UHUGEINT: { + auto digits = value.ToString(); + auto number = PyLong_FromString(digits.c_str(), nullptr, 10); + if (!number) { + throw nb::python_error(); + } + return nb::steal(number); + } + case LogicalTypeId::FLOAT: + case LogicalTypeId::DOUBLE: { + const double number = + value.type().id() == LogicalTypeId::FLOAT ? value.GetValue() : value.GetValue(); + if (!std::isfinite(number)) { + throw PredicateRefused(); + } + return nb::float_(number); + } + case LogicalTypeId::DECIMAL: + return nb::module_::import_("decimal").attr("Decimal")(PythonText(value.ToString())); + case LogicalTypeId::VARCHAR: + return PythonText(StringValue::Get(value)); + case LogicalTypeId::BLOB: { + auto &blob = StringValue::Get(value); + return nb::bytes(blob.data(), blob.size()); + } + case LogicalTypeId::DATE: { + auto date = value.GetValueUnsafe(); + if (!Value::IsFinite(date)) { + throw PredicateRefused(); + } + int32_t year, month, day; + Date::Convert(date, year, month, day); + if (!PythonYear(year)) { + throw PredicateRefused(); + } + return nb::module_::import_("datetime").attr("date")(year, month, day); + } + case LogicalTypeId::TIME: { + int32_t hour, minute, second, micros; + Time::Convert(value.GetValueUnsafe(), hour, minute, second, micros); + if (hour > 23) { + throw PredicateRefused(); + } + return nb::module_::import_("datetime").attr("time")(hour, minute, second, micros); + } + case LogicalTypeId::TIMESTAMP: + case LogicalTypeId::TIMESTAMP_SEC: + case LogicalTypeId::TIMESTAMP_MS: + case LogicalTypeId::TIMESTAMP_TZ: { + auto timestamp = Microseconds(value); + date_t date; + dtime_t time; + Timestamp::Convert(timestamp, date, time); + int32_t year, month, day, hour, minute, second, micros; + Date::Convert(date, year, month, day); + if (!PythonYear(year)) { + throw PredicateRefused(); + } + if (value.type().id() == LogicalTypeId::TIMESTAMP_TZ) { + return nb::int_(timestamp.value); + } + Time::Convert(time, hour, minute, second, micros); + return nb::module_::import_("datetime").attr("datetime")(year, month, day, hour, minute, second, micros); + } + default: + throw PredicateRefused(); + } + } + + const std::function &resolve; + nb::module_ nodes; +}; + +struct PredicateProbe : public TableFunctionInfo { + PredicateProbe(const PredicateColumns &columns, nb::list &entries) : columns(columns), entries(entries) { + } + + const PredicateColumns &columns; + nb::list &entries; + unordered_set offered; +}; + +void CaptureOfferedPredicates(ClientContext &context, LogicalGet &get, FunctionData *, + vector> &filters) { + engine::CallbackScope callback(context); + auto &probe = get.function.function_info->Cast(); + auto &column_ids = get.GetColumnIds(); + auto resolve = [&](idx_t binding_index) { + if (binding_index >= column_ids.size()) { + throw PredicateRefused(); + } + return probe.columns.Resolve(column_ids[binding_index].GetPrimaryIndex()); + }; + for (auto &filter : filters) { + auto text = filter->ToString(); + // Later optimizer passes offer again the filters that stayed above the scan. + if (!probe.offered.insert(text).second) { + continue; + } + nb::object node = nb::none(); + try { + node = TranslatePredicate(*filter, resolve); + } catch (PredicateRefused &) { + } + probe.entries.append(nb::make_tuple(PythonText(text), node)); + } +} + +struct DropTable { + Connection &connection; + string table; + + ~DropTable() { + try { + connection.Query("DROP TABLE IF EXISTS temp.main." + table); + } catch (...) { + } + } +}; + +LogicalGet &OnlyGet(LogicalOperator &op) { + if (op.type == LogicalOperatorType::LOGICAL_GET) { + return op.Cast(); + } + if (op.children.size() != 1) { + throw InvalidInputException("The predicate must not change the shape of the plan"); + } + return OnlyGet(*op.children[0]); +} + +} // namespace + +PredicateColumns::PredicateColumns(const vector &names, const vector &types) { + D_ASSERT(names.size() == types.size()); + for (idx_t i = 0; i < names.size(); i++) { + columns.push_back(PredicateColumn {names[i], types[i]}); + } + nameable = UniqueNames(columns); +} + +PredicateColumns::PredicateColumns(const ArrowSchema &schema, const vector &types) { + if (schema.n_children < 0 || NumericCast(schema.n_children) != types.size()) { + throw InvalidInputException("The Arrow schema has %d columns where the scan declares %d", schema.n_children, + types.size()); + } + for (idx_t i = 0; i < types.size(); i++) { + const auto &child = *schema.children[i]; + columns.push_back(PredicateColumn {child.name ? child.name : "", types[i]}); + } + nameable = UniqueNames(columns); + for (idx_t i = 0; i < types.size(); i++) { + nameable[i] = nameable[i] && SameFieldNames(*schema.children[i], types[i]); + } +} + +PredicateColumn PredicateColumns::Resolve(idx_t declared_index) const { + if (declared_index >= columns.size() || !nameable[declared_index]) { + throw PredicateRefused(); + } + return columns[declared_index]; +} + +nb::object TranslatePredicate(const Expression &predicate, + const std::function &resolve) { + DUCKDB_PY_ASSERT_GIL_HELD(); + return PredicateWalker(resolve).Predicate(predicate); +} + +nb::list TranslatePredicateForTesting(DuckDBPyConnection &py_connection, const string &predicate, + const nb::object &columns, const nb::object &schema, const string &select) { + DuckDBPyConnection::ConnectionLockGuard lock(py_connection); + auto &connection = py_connection.con.GetConnection(); + auto &context = *connection.context; + vector names; + vector type_texts; + nb::object capsule; + const ArrowSchema *arrow_schema = nullptr; + { + engine::CallbackScope callback(context); + if (schema.is_none() == columns.is_none()) { + throw InvalidInputException("Give either the columns or an Arrow schema"); + } + if (!schema.is_none()) { + capsule = schema.attr("__arrow_c_schema__")(); + arrow_schema = static_cast(PyCapsule_GetPointer(capsule.ptr(), "arrow_schema")); + if (!arrow_schema) { + throw nb::python_error(); + } + } else { + for (auto item : nb::cast(columns)) { + names.push_back(nb::cast(item.first)); + type_texts.push_back(nb::cast(nb::str(item.second))); + } + } + } + auto uuid = UUID::ToString(UUID::GenerateRandomUUID()); + std::replace(uuid.begin(), uuid.end(), '-', '_'); + const string table = "__duckdb_predicate_probe_" + uuid; + DropTable drop {connection, SQLIdentifier::ToString(table)}; + auto statements = connection.ExtractStatements("SELECT " + select + " FROM temp.main." + + SQLIdentifier::ToString(table) + " WHERE " + predicate); + if (statements.size() != 1) { + throw InvalidInputException("The predicate must be one expression"); + } + nb::list entries; + auto probe_plan = [&]() { + vector types; + if (arrow_schema) { + ArrowTableSchema arrow_table; + ArrowTableFunction::PopulateArrowTableSchema(context, arrow_table, *arrow_schema); + names = arrow_table.GetNames(); + types = arrow_table.GetTypes(); + } else { + for (auto &type_text : type_texts) { + types.push_back(TransformStringToLogicalType(type_text, context)); + } + } + if (names.empty()) { + throw InvalidInputException("A predicate needs at least one column"); + } + auto info = make_uniq( + QualifiedName(Identifier::TempCatalog(), Identifier::DefaultSchema(), Identifier(table))); + info->temporary = true; + for (idx_t i = 0; i < names.size(); i++) { + info->columns.AddColumn(ColumnDefinition(Identifier(names[i]), types[i])); + } + Catalog::GetCatalog(context, Identifier::TempCatalog()).CreateTable(context, std::move(info)); + + Planner planner(context); + planner.CreatePlan(std::move(statements[0])); + if (!planner.plan) { + throw InvalidInputException("The predicate's parameter types could not be resolved"); + } + auto &get = OnlyGet(*planner.plan); + auto resolver = arrow_schema ? PredicateColumns(*arrow_schema, get.returned_types) + : PredicateColumns(names, get.returned_types); + get.function.filter_pushdown = false; + get.function.pushdown_complex_filter = CaptureOfferedPredicates; + get.function.function_info = make_shared_ptr(resolver, entries); + Optimizer optimizer(*planner.binder, context); + optimizer.Optimize(std::move(planner.plan)); + }; + context.RunFunctionInTransaction(probe_plan); + return entries; +} + +} // namespace duckdb diff --git a/src/arrow/pyarrow_filter_pushdown.cpp b/src/arrow/pyarrow_filter_pushdown.cpp deleted file mode 100644 index 3fc513b2..00000000 --- a/src/arrow/pyarrow_filter_pushdown.cpp +++ /dev/null @@ -1,336 +0,0 @@ -#include "duckdb_python/arrow/pyarrow_filter_pushdown.hpp" - -#include "duckdb_python/arrow/filter_pushdown_visitor.hpp" -#include "duckdb_python/pyconnection/pyconnection.hpp" -#include "duckdb_python/python_objects.hpp" -#include "duckdb_python/pyrelation.hpp" -#include "duckdb/function/table/arrow.hpp" -#include - -namespace duckdb { - -namespace { - -string ColumnTimeZone(const ArrowSchema *schema) { - const string format = schema && schema->format ? schema->format : ""; - const auto colon = format.find(':'); - if (format.rfind("ts", 0) != 0 || colon == string::npos || colon + 1 == format.size()) { - // A TIMESTAMP_TZ constant only meets a zoned timestamp column, so any other schema node is one the client - // cannot read. - throw NotImplementedException("Cannot read the time zone of column \"%s\" from the Arrow schema", - schema && schema->name ? schema->name : ""); - } - return format.substr(colon + 1); -} - -// Build a pyarrow.dataset scalar matching the given DuckDB Value and (optionally) ArrowType. -// The ArrowType is needed for decimal precision / blob-view disambiguation; the -// DuckDB Value alone is not sufficient. -nb::object MakePyArrowScalar(const Value &constant, const string &timezone_config, const ArrowType *arrow_type, - const ArrowSchema *arrow_schema) { - auto &import_cache = *DuckDBPyConnection::ImportCache(); - auto scalar = import_cache.pyarrow.scalar(); - nb::handle dataset_scalar = import_cache.pyarrow.dataset().attr("scalar"); - - switch (constant.type().id()) { - case LogicalTypeId::BOOLEAN: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::TINYINT: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::SMALLINT: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::INTEGER: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::BIGINT: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::DATE: { - nb::handle date_type = import_cache.pyarrow.date32(); - return dataset_scalar(scalar(constant.GetValue(), date_type())); - } - case LogicalTypeId::TIME: { - nb::handle date_type = import_cache.pyarrow.time64(); - return dataset_scalar(scalar(constant.GetValue(), date_type("us"))); - } - case LogicalTypeId::TIME_NS: { - // Polars TIME columns round-trip through arrow as time64("ns"). - // `Value::GetValue()` has a hand-rolled fast-path switch for TIME but not - // TIME_NS — it falls through to GetValueInternal, which then tries - // Cast::Operation for which no specialization exists, and - // throws "Unimplemented type for cast (INT64 -> INT64)". Use the type-strong - // GetValueUnsafe() which reads `value_.time_ns` from the union - // directly. dtime_ns_t.value holds nanoseconds (see arrow_conversion.cpp:432). - nb::handle date_type = import_cache.pyarrow.time64(); - return dataset_scalar(scalar(constant.GetValueUnsafe().value, date_type("ns"))); - } - case LogicalTypeId::TIMESTAMP: { - nb::handle date_type = import_cache.pyarrow.timestamp(); - return dataset_scalar(scalar(constant.GetValue(), date_type("us"))); - } - case LogicalTypeId::TIMESTAMP_MS: { - nb::handle date_type = import_cache.pyarrow.timestamp(); - return dataset_scalar(scalar(constant.GetValue(), date_type("ms"))); - } - case LogicalTypeId::TIMESTAMP_NS: { - nb::handle date_type = import_cache.pyarrow.timestamp(); - return dataset_scalar(scalar(constant.GetValue(), date_type("ns"))); - } - case LogicalTypeId::TIMESTAMP_SEC: { - nb::handle date_type = import_cache.pyarrow.timestamp(); - return dataset_scalar(scalar(constant.GetValue(), date_type("s"))); - } - case LogicalTypeId::TIMESTAMP_TZ: { - // The microsecond scalar keeps its UTC instant under the column's own zone, and pyarrow has no kernel for a - // comparison that differs in both unit and zone. - nb::handle date_type = import_cache.pyarrow.timestamp(); - auto zone = ColumnTimeZone(arrow_schema); - return dataset_scalar(scalar(constant.GetValue(), date_type("us", nb::arg("tz") = zone))); - } - case LogicalTypeId::TIMESTAMP_TZ_NS: { - nb::handle date_type = import_cache.pyarrow.timestamp(); - return dataset_scalar(scalar(constant.GetValue(), date_type("ns", nb::arg("tz") = timezone_config))); - } - case LogicalTypeId::UTINYINT: { - nb::handle integer_type = import_cache.pyarrow.uint8(); - return dataset_scalar(scalar(constant.GetValue(), integer_type())); - } - case LogicalTypeId::USMALLINT: { - nb::handle integer_type = import_cache.pyarrow.uint16(); - return dataset_scalar(scalar(constant.GetValue(), integer_type())); - } - case LogicalTypeId::UINTEGER: { - nb::handle integer_type = import_cache.pyarrow.uint32(); - return dataset_scalar(scalar(constant.GetValue(), integer_type())); - } - case LogicalTypeId::UBIGINT: { - nb::handle integer_type = import_cache.pyarrow.uint64(); - return dataset_scalar(scalar(constant.GetValue(), integer_type())); - } - case LogicalTypeId::FLOAT: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::DOUBLE: - return dataset_scalar(constant.GetValue()); - case LogicalTypeId::VARCHAR: - return dataset_scalar(constant.ToString()); - case LogicalTypeId::BLOB: { - if (arrow_type && arrow_type->GetTypeInfo().GetSizeType() == ArrowVariableSizeType::VIEW) { - nb::handle binary_view_type = import_cache.pyarrow.binary_view(); - { - auto blob = constant.GetValueUnsafe(); - return dataset_scalar(scalar(nb::bytes(blob.data(), blob.size()), binary_view_type())); - } - } - { - auto blob = constant.GetValueUnsafe(); - return dataset_scalar(nb::bytes(blob.data(), blob.size())); - } - } - case LogicalTypeId::DECIMAL: { - if (!arrow_type) { - throw NotImplementedException("Cannot push down DECIMAL filter without an arrow type"); - } - nb::handle decimal_type; - auto &decimal_info = arrow_type->GetTypeInfo(); - auto bit_width = decimal_info.GetBitWidth(); - switch (bit_width) { - case DecimalBitWidth::DECIMAL_32: - decimal_type = import_cache.pyarrow.decimal32(); - break; - case DecimalBitWidth::DECIMAL_64: - decimal_type = import_cache.pyarrow.decimal64(); - break; - case DecimalBitWidth::DECIMAL_128: - decimal_type = import_cache.pyarrow.decimal128(); - break; - default: - throw NotImplementedException("Unsupported precision for Arrow Decimal Type."); - } - - uint8_t width; - uint8_t scale; - constant.type().GetDecimalProperties(width, scale); - auto val = import_cache.decimal.Decimal()(constant.ToString()); - return dataset_scalar( - scalar(std::move(val), decimal_type(nb::arg("precision") = width, nb::arg("scale") = scale))); - } - default: - throw NotImplementedException("Unimplemented type \"%s\" for Arrow Filter Pushdown", - constant.type().ToString()); - } -} - -struct PyArrowBackend : public FilterBackend { - explicit PyArrowBackend(const ClientProperties &client_properties_p) : client_properties(client_properties_p) { - auto &import_cache = *DuckDBPyConnection::ImportCache(); - field_factory = import_cache.pyarrow.dataset().attr("field"); - dataset_scalar = import_cache.pyarrow.dataset().attr("scalar"); - } - - nb::object MakeColumnRef(const vector &path) override { - vector str_path; - std::transform(path.begin(), path.end(), std::back_inserter(str_path), - [](const Identifier &segment) { return segment.GetIdentifierName(); }); - return field_factory(nb::tuple(nb::cast(str_path))); - } - - nb::object MakeScalar(const Value &v, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, - const string &timezone_config) override { - return MakePyArrowScalar(v, timezone_config, arrow_type, arrow_schema); - } - - nb::object Compare(ExpressionType op, nb::object col, nb::object scalar) override { - switch (op) { - case ExpressionType::COMPARE_EQUAL: - return col.attr("__eq__")(scalar); - case ExpressionType::COMPARE_NOTEQUAL: - return col.attr("__ne__")(scalar); - case ExpressionType::COMPARE_LESSTHAN: - return col.attr("__lt__")(scalar); - case ExpressionType::COMPARE_GREATERTHAN: - return col.attr("__gt__")(scalar); - case ExpressionType::COMPARE_LESSTHANOREQUALTO: - return col.attr("__le__")(scalar); - case ExpressionType::COMPARE_GREATERTHANOREQUALTO: - return col.attr("__ge__")(scalar); - default: - throw NotImplementedException("Comparison Type %s can't be an Arrow Scan Pushdown Filter", - ExpressionTypeToString(op)); - } - } - - nb::object NaNCompare(ExpressionType op, nb::object col) override { - switch (op) { - case ExpressionType::COMPARE_EQUAL: - case ExpressionType::COMPARE_GREATERTHANOREQUALTO: - return IsNaN(std::move(col)); - case ExpressionType::COMPARE_LESSTHAN: - case ExpressionType::COMPARE_NOTEQUAL: - // Every value but NaN equals itself, and a NULL row stays NULL. - return col.attr("__eq__")(col); - case ExpressionType::COMPARE_GREATERTHAN: - // Nothing is greater than NaN. - return dataset_scalar(false); - case ExpressionType::COMPARE_LESSTHANOREQUALTO: - // NaN sorts above every other value, so only NULL rows fail. - return IsNotNull(std::move(col)); - default: - throw NotImplementedException("Unsupported comparison type (%s) for NaN values", - ExpressionTypeToString(op)); - } - } - - nb::object IsNaN(nb::object col) override { - // NaN is the only value unequal to itself, and pyarrow has no is_nan kernel for a dictionary column. - return col.attr("__ne__")(col); - } - - nb::object IsNull(nb::object col) override { - return col.attr("is_null")(); - } - - nb::object IsNotNull(nb::object col) override { - return col.attr("is_valid")(); - } - - nb::object IsIn(nb::object col, const vector &values, const LogicalType &col_logical_type, - const string &timezone_config) override { - (void)timezone_config; - const auto col_type_id = col_logical_type.id(); - const bool is_float = col_type_id == LogicalTypeId::FLOAT || col_type_id == LogicalTypeId::DOUBLE; - - // isin compares float bits (missing -0.0 and NaN payloads), and is_nan has no dictionary kernel. - if (is_float) { - bool has_zero = false; - nb::list non_zero_values; - - for (auto &val : values) { - if (val.IsNull()) { - non_zero_values.append(PythonObject::FromValue(val, val.type(), client_properties)); - continue; - } - - double dbl = col_type_id == LogicalTypeId::FLOAT ? static_cast(val.GetValue()) - : val.GetValue(); - - if (std::isnan(dbl)) { - throw NotImplementedException( - "Arrow filter pushdown: IN list with NaN over float/double is not pushed"); - } - - if (dbl == 0.0) { - has_zero = true; - } else { - non_zero_values.append(PythonObject::FromValue(val, val.type(), client_properties)); - } - } - - if (has_zero) { - auto zero_expr = Compare(ExpressionType::COMPARE_EQUAL, col, dataset_scalar(0.0)); - if (nb::len(non_zero_values) > 0) { - auto isin_expr = col.attr("isin")(non_zero_values); - return Or(zero_expr, isin_expr); - } - return zero_expr; - } - - if (nb::len(non_zero_values) > 0) { - return col.attr("isin")(non_zero_values); - } - - return dataset_scalar(false); - } - - nb::list py_values; - for (auto &val : values) { - py_values.append(PythonObject::FromValue(val, val.type(), client_properties)); - } - return col.attr("isin")(std::move(py_values)); - } - - nb::object And(nb::object a, nb::object b) override { - return a.attr("__and__")(b); - } - - nb::object Or(nb::object a, nb::object b) override { - return a.attr("__or__")(b); - } - -private: - const ClientProperties &client_properties; - nb::object field_factory; - nb::object dataset_scalar; -}; - -} // anonymous namespace - -nb::object PyArrowFilterPushdown::TransformFilter(TableFilterSet &filter_collection, - unordered_map &columns, - unordered_map filter_to_col, - const ClientProperties &config, const ArrowTableSchema &arrow_table, - const ArrowSchema &arrow_schema) { - PyArrowBackend backend(config); - nb::object expression = nb::none(); - for (auto &entry : filter_collection) { - auto column_idx = entry.GetIndex(); - auto &column_name = columns[column_idx]; - D_ASSERT(columns.find(column_idx) != columns.end()); - - vector column_path = {Identifier(column_name)}; - auto arrow_column_idx = filter_to_col.at(column_idx); - auto &arrow_type = arrow_table.GetColumns().at(arrow_column_idx); - D_ASSERT(arrow_column_idx < NumericCast(arrow_schema.n_children)); - auto column_schema = arrow_schema.children[arrow_column_idx]; - nb::object child_expression = duckdb::TransformFilter(entry.Filter(), std::move(column_path), backend, - arrow_type.get(), column_schema, config.time_zone); - if (child_expression.is(nb::none())) { - continue; - } - if (expression.is(nb::none())) { - expression = std::move(child_expression); - } else { - expression = expression.attr("__and__")(child_expression); - } - } - return expression; -} - -} // namespace duckdb diff --git a/src/arrow/python_arrow_scan.cpp b/src/arrow/python_arrow_scan.cpp new file mode 100644 index 00000000..242e22f5 --- /dev/null +++ b/src/arrow/python_arrow_scan.cpp @@ -0,0 +1,822 @@ +#include "duckdb_python/arrow/python_arrow_scan.hpp" + +#include "duckdb_python/arrow/arrow_importer.hpp" +#include "duckdb_python/arrow/predicate_translation.hpp" +#include "duckdb_python/arrow/single_use.hpp" + +#include "duckdb/catalog/catalog.hpp" +#include "duckdb/catalog/catalog_transaction.hpp" +#include "duckdb/common/arrow/arrow.hpp" +#include "duckdb/common/arrow/arrow_wrapper.hpp" +#include "duckdb/common/atomic.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/common/mutex.hpp" +#include "duckdb/common/optional_idx.hpp" +#include "duckdb/common/string_util.hpp" +#include "duckdb/function/table/arrow.hpp" +#include "duckdb/main/client_context.hpp" +#include "duckdb/main/external_dependencies.hpp" +#include "duckdb/parser/expression/function_expression.hpp" +#include "duckdb/parser/parsed_data/create_table_function_info.hpp" +#include "duckdb/planner/operator/logical_get.hpp" +#include "duckdb/storage/statistics/node_statistics.hpp" +#include "duckdb_python/engine/connection.hpp" +#include "duckdb_python/exceptions.hpp" +#include "duckdb_python/pyconnection/pyconnection.hpp" +#include "duckdb_python/python_dependency.hpp" + +#include +#include +#include +#include +#ifndef _WIN32 +#include +#endif + +namespace duckdb { + +namespace { + +constexpr const char *STREAM_CAPSULE = "arrow_array_stream"; +constexpr const char *SCHEMA_CAPSULE = "arrow_schema"; +constexpr const char *ARRAY_CAPSULE = "arrow_array"; +//! The engine clamps the cap to its scheduler's thread count, so a large one asks for every thread it has +constexpr idx_t STREAM_MAX_THREADS = 1024; + +string StreamError(ArrowArrayStream &stream, int code) { + const char *text = stream.get_last_error ? stream.get_last_error(&stream) : nullptr; + if (!text || !*text) { + return StringUtil::Format("no error message (error code %d)", code); + } + return text; +} + +//! Set by an atexit hook: from then on no scan calls into a source, since a thread that waits for the GIL while the +//! interpreter finalizes never gets it, and the thread finalizing may itself be waiting for that scan to end. A scan +//! an exit handler starts after the hook is still read, as the interpreter only finalizes once every handler ran. +atomic interpreter_exiting {false}; +//! Calls into a source that were let in and have not returned; the atexit hook waits for them with the GIL released. +atomic calls_in_flight {0}; +//! The calls among them made on this thread, the only ones still in flight in a child this thread forks +thread_local idx_t calls_on_this_thread = 0; + +class SourceCall { +public: + explicit SourceCall(bool started_during_exit = false) { + calls_in_flight++; + admitted = !interpreter_exiting.load() || started_during_exit; + if (!admitted) { + calls_in_flight -= 1; + } else { + calls_on_this_thread++; + } + } + ~SourceCall() { + if (admitted) { + calls_on_this_thread -= 1; + calls_in_flight -= 1; + } + } + SourceCall(const SourceCall &) = delete; + SourceCall &operator=(const SourceCall &) = delete; + + explicit operator bool() const { + return admitted; + } + +private: + bool admitted; +}; + +void OnInterpreterExit() { + interpreter_exiting = true; + nb::gil_scoped_release release; + // Not bounded: a call let in may need the GIL again before it returns, which it never gets once the interpreter + // finalizes, so a slow source delays the exit instead of hanging it. + while (calls_in_flight.load() > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } +} + +#ifndef _WIN32 +//! A forked child keeps the exit flag, as it continues from the parent's place among the exit handlers. +void AfterForkInChild() { + calls_in_flight = calls_on_this_thread; +} +#endif + +//! Whether this call comes from an exit handler running after the exit hook, which the hook lets through; call with +//! the GIL held +bool StartedDuringExit() { + return interpreter_exiting.load() && !nb::cast(nb::module_::import_("sys").attr("is_finalizing")()); +} + +//! Releases an Arrow struct taken from a source: under the GIL when the source asked for it, directly otherwise, and +//! not at all once the interpreter is exiting, since even pyarrow's own release can enter Python. +template +void ReleaseFromSource(T &value, bool under_gil, bool admitted_during_exit) noexcept { + if (!value.release) { + return; + } + if (SourceCall call {admitted_during_exit}) { + if (!under_gil) { + value.release(&value); + } else if (nb::detail::cleanup_guard guard {}) { + value.release(&value); + } + } + value.release = nullptr; +} + +template +struct ArrowOwned { + ArrowOwned() { + value.release = nullptr; + } + ArrowOwned(ArrowOwned &&other) noexcept + : value(other.value), release_under_gil(other.release_under_gil), + admitted_during_exit(other.admitted_during_exit) { + other.value.release = nullptr; + } + ArrowOwned &operator=(ArrowOwned &&other) noexcept { + if (this != &other) { + Release(); + value = other.value; + release_under_gil = other.release_under_gil; + admitted_during_exit = other.admitted_during_exit; + other.value.release = nullptr; + } + return *this; + } + ArrowOwned(const ArrowOwned &) = delete; + ArrowOwned &operator=(const ArrowOwned &) = delete; + ~ArrowOwned() { + Release(); + } + + void Release() noexcept { + ReleaseFromSource(value, release_under_gil, admitted_during_exit); + } + explicit operator bool() const { + return value.release != nullptr; + } + + T value {}; + bool release_under_gil = true; + bool admitted_during_exit = false; +}; + +struct RoutedRelease { + void (*release)(ArrowArray *); + void *private_data; + bool under_gil; + bool admitted_during_exit; +}; + +void ReleaseRouted(ArrowArray *array) { + auto routed = static_cast(array->private_data); + array->release = routed->release; + array->private_data = routed->private_data; + const bool under_gil = routed->under_gil; + const bool admitted_during_exit = routed->admitted_during_exit; + delete routed; + ReleaseFromSource(*array, under_gil, admitted_during_exit); +} + +//! The engine releases an adopted array wherever its last vector goes, on any thread and after the scan is gone, so +//! its release is routed through ReleaseFromSource here. +void RouteRelease(ArrowArray &array, bool under_gil, bool admitted_during_exit) { + array.private_data = new RoutedRelease {array.release, array.private_data, under_gil, admitted_during_exit}; + array.release = &ReleaseRouted; +} + +template +ArrowOwned TakeFromCapsule(nb::handle capsule, const char *kind, const string &subject) { + auto &held = InCapsule(capsule, kind, subject); + ArrowOwned taken; + taken.value = held; + held.release = nullptr; + return taken; +} + +bool IsBatch(const ArrowSchema &schema) { + return schema.format && std::strcmp(schema.format, "+s") == 0; +} + +string NameOf(const ArrowSchema &schema, idx_t index) { + auto child = schema.children[index]; + return child->name ? child->name : ""; +} + +//! A struct array with no validity buffer still declares one buffer slot, holding null. +const void *NO_BUFFERS[1] = {nullptr}; + +template +struct Wrapped { + T *child; + T **children; +}; + +template +void ReleaseWrapped(T *wrapper) { + auto owned = static_cast *>(wrapper->private_data); + if (owned->child->release) { + owned->child->release(owned->child); + } + delete owned->child; + delete[] owned->children; + delete owned; + wrapper->release = nullptr; +} + +void WrapAsBatch(ArrowSchema &schema) { + auto child = new ArrowSchema(schema); + auto children = new ArrowSchema *[1] {child}; + auto wrapped = new Wrapped {child, children}; + schema = ArrowSchema {}; + schema.format = "+s"; + schema.name = ""; + schema.n_children = 1; + schema.children = children; + schema.release = &ReleaseWrapped; + schema.private_data = wrapped; +} + +void WrapAsBatch(ArrowArray &array) { + auto child = new ArrowArray(array); + auto children = new ArrowArray *[1] {child}; + auto wrapped = new Wrapped {child, children}; + array = ArrowArray {}; + array.length = child->length; + array.null_count = 0; + array.offset = 0; + array.n_buffers = 1; + array.buffers = NO_BUFFERS; + array.n_children = 1; + array.children = children; + array.release = &ReleaseWrapped; + array.private_data = wrapped; +} + +//! A predicate the source promised to apply, whole, since one may read several columns, and its text for the plan +struct AcceptedPredicate { + shared_ptr node; + string text; +}; + +//! What one bind of the scan knows; shared with copies of the plan, so nothing of one execution lives here +struct PythonArrowScanBindData : public TableFunctionData { + explicit PythonArrowScanBindData(shared_ptr info_p) : info(std::move(info_p)) { + } + + unique_ptr Copy() const override { + return make_uniq(*this); + } + + shared_ptr info; + //! The Arrow names of the declared columns, empty for a nameless array + vector raw_names; + vector types; + ArrowTableSchema arrow_table; + PredicateColumns predicate_columns; + optional_idx rows; + //! Accumulated across the optimizer's calls + vector predicates; +}; + +struct Exported { + nb::object stream; + nb::object schema; + nb::object array; + bool projected = false; + + bool IsArray() const { + return array.is_valid(); + } +}; + +//! Call with the GIL held +Exported ExportStream(const PythonArrowSourceInfo &info, const vector *columns, + const vector &predicates) { + nb::object request = nb::none(); + if (columns) { + nb::list wanted; + for (auto column : *columns) { + wanted.append(nb::int_(static_cast(column))); + } + request = std::move(wanted); + } + nb::list promised; + for (auto &predicate : predicates) { + promised.append(predicate.node->obj); + } + auto answer = info.source.obj.attr("stream")(request, nb::tuple(promised)); + if (!PyTuple_Check(answer.ptr()) || PyTuple_Size(answer.ptr()) != 2 || + !PyBool_Check(PyTuple_GetItem(answer.ptr(), 1))) { + throw InvalidInputException("%s answered stream() with something other than a (data, projected) pair", + Subject(info.name, info.kind, "source")); + } + auto data = nb::borrow(PyTuple_GetItem(answer.ptr(), 0)); + Exported exported; + exported.projected = PyTuple_GetItem(answer.ptr(), 1) == Py_True; + if (PyTuple_Check(data.ptr()) && PyTuple_Size(data.ptr()) == 2) { + exported.schema = nb::borrow(PyTuple_GetItem(data.ptr(), 0)); + exported.array = nb::borrow(PyTuple_GetItem(data.ptr(), 1)); + } else { + exported.stream = std::move(data); + } + return exported; +} + +} // namespace + +struct CachedArrowSchema { + ArrowOwned schema; + bool wrapped = false; +}; + +namespace { + +//! Resolves a batch schema whose columns may carry a NULL name, which the Arrow C data interface allows and the engine +//! cannot read: such a column resolves as one named "", through copies that leave the source's structs untouched. +void ResolveBatchSchema(ClientContext &context, ArrowTableSchema &table, const ArrowSchema &batch) { + const auto count = NumericCast(batch.n_children); + vector columns(count); + vector pointers(count); + for (idx_t i = 0; i < count; i++) { + columns[i] = *batch.children[i]; + if (!columns[i].name) { + columns[i].name = ""; + } + pointers[i] = &columns[i]; + } + auto named = batch; + named.children = pointers.data(); + ArrowTableFunction::PopulateArrowTableSchema(context, table, named); +} + +//! The source's schema, read without consuming a stream: from a schema hook, asked once for the life of a source whose +//! schema is stable and at every bind otherwise, else from a stream that is only borrowed; call with the GIL held +shared_ptr SchemaOf(PythonArrowSourceInfo &info) { + if (info.schema) { + return info.schema; + } + auto &source = info.source.obj; + const auto subject = Subject(info.name, info.kind, "object"); + auto result = make_shared_ptr(); + bool from_hook = true; + auto underlying = nb::getattr(source, "obj", nb::none()); + if (PyUtil::HasAttribute(source, "__arrow_c_schema__")) { + result->schema = TakeFromCapsule(source.attr("__arrow_c_schema__")(), SCHEMA_CAPSULE, subject); + } else if (!underlying.is_none() && PyUtil::HasAttribute(underlying, "schema") && + PyUtil::HasAttribute(underlying.attr("schema"), "_export_to_c")) { + underlying.attr("schema").attr("_export_to_c")(reinterpret_cast(&result->schema.value)); + } else { + from_hook = false; + auto exported = ExportStream(info, nullptr, {}); + if (exported.IsArray()) { + // The array capsule is dropped unread, which releases it. + result->schema = TakeFromCapsule(exported.schema, SCHEMA_CAPSULE, subject); + } else { + // Only borrowed, so a stream the source keeps, such as a bare capsule, is still there for the scan. + auto &stream = InCapsule(exported.stream, STREAM_CAPSULE, subject); + auto code = stream.get_schema(&stream, &result->schema.value); + if (code != 0) { + result->schema.value.release = nullptr; + throw InvalidInputException("reading the schema of %s failed: %s", + Subject(info.name, info.kind, "stream"), StreamError(stream, code)); + } + } + } + if (!result->schema) { + throw InvalidInputException("reading the schema of %s failed: it gave a released schema", subject); + } + result->schema.admitted_during_exit = StartedDuringExit(); + result->wrapped = !IsBatch(result->schema.value); + if (result->wrapped) { + WrapAsBatch(result->schema.value); + } + if (from_hook && nb::getattr(source, "schema_is_stable", nb::handle(Py_False)).ptr() == Py_True) { + info.schema = result; + } + return result; +} + +unique_ptr PythonArrowScanBind(ClientContext &context, TableFunctionBindInput &input, + vector &return_types, vector &names) { + if (!input.ref.bind_info) { + throw BinderException("python_arrow_scan requires a Python Arrow source bind input"); + } + DynamicCastCheck(input.ref.bind_info.get()); + auto info = shared_ptr_cast(input.ref.bind_info); + auto result = make_uniq(info); + shared_ptr schema; + { + nb::gil_scoped_acquire gil; + engine::CallbackScope callback(context); + try { + schema = SchemaOf(*info); + } catch (nb::python_error &error) { + throw InvalidInputException("reading the schema of %s failed: %s", + Subject(info->name, info->kind, "object"), FormatUDFPythonError(error)); + } + try { + auto rows = info->source.obj.attr("rows")(); + if (PyLong_Check(rows.ptr()) && !PyBool_Check(rows.ptr())) { + auto count = PyLong_AsUnsignedLongLong(rows.ptr()); + if (PyErr_Occurred()) { + PyErr_Clear(); + } else { + result->rows = NumericCast(count); + } + } + } catch (nb::python_error &error) { + throw InvalidInputException("counting the rows of %s failed: %s", Subject(info->name, info->kind, "object"), + FormatUDFPythonError(error)); + } + } + auto &arrow_schema = schema->schema.value; + ResolveBatchSchema(context, result->arrow_table, arrow_schema); + result->types = result->arrow_table.GetTypes(); + if (result->types.empty()) { + throw InvalidInputException("Provided table/dataframe must have at least one column"); + } + auto &resolved_names = result->arrow_table.GetNames(); + for (idx_t i = 0; i < result->types.size(); i++) { + auto raw = NameOf(arrow_schema, i); + // A plain array has no column name of its own; the engine would call it v0. + names.emplace_back(schema->wrapped && raw.empty() ? string("value") : resolved_names[i]); + result->raw_names.push_back(std::move(raw)); + } + result->predicate_columns = PredicateColumns(arrow_schema, result->types); + return_types = result->types; + return std::move(result); +} + +unique_ptr PythonArrowScanCardinality(ClientContext &context, const FunctionData *bind_data_p) { + auto &bind_data = bind_data_p->Cast(); + if (!bind_data.rows.IsValid()) { + return make_uniq(); + } + // An estimate only: an object can grow in place after bind, and a maximum would let a reused plan drop a LIMIT. + return make_uniq(bind_data.rows.GetIndex()); +} + +//! Offers each predicate the walker can carry exactly to the source; one it accepts leaves the plan, which is why only +//! a True answer accepts. The optimizer may call this more than once, each time with the predicates still in the plan. +void PythonArrowScanPushdown(ClientContext &context, LogicalGet &get, FunctionData *bind_data_p, + vector> &filters) { + auto &bind_data = bind_data_p->Cast(); + if (filters.empty()) { + return; + } + auto &column_ids = get.GetColumnIds(); + nb::gil_scoped_acquire gil; + engine::CallbackScope callback(context); + try { + nb::object fragment; + for (idx_t i = 0; i < filters.size(); i++) { + auto resolve = [&](idx_t binding_index) { + if (binding_index >= column_ids.size()) { + throw PredicateRefused(); + } + auto &column = column_ids[binding_index]; + if (!column.HasPrimaryIndex() || column.IsVirtualColumn()) { + throw PredicateRefused(); + } + return bind_data.predicate_columns.Resolve(column.GetPrimaryIndex()); + }; + nb::object node; + try { + node = TranslatePredicate(*filters[i], resolve); + } catch (const PredicateRefused &) { + continue; + } + auto answer = bind_data.info->source.obj.attr("accepts")(node); + if (answer.ptr() != Py_True) { + continue; + } + if (!fragment.is_valid()) { + fragment = DuckDBPyConnection::ArrowSourcesNodes().attr("fragment"); + } + AcceptedPredicate accepted; + accepted.text = ToUTF8(fragment(node)); + accepted.node = make_shared_ptr(std::move(node)); + bind_data.predicates.push_back(std::move(accepted)); + filters.erase_at(i); + i -= 1; + } + } catch (nb::python_error &error) { + throw InvalidInputException("offering a filter to %s failed: %s", + Subject(bind_data.info->name, bind_data.info->kind, "source"), + FormatUDFPythonError(error)); + } +} + +InsertionOrderPreservingMap PythonArrowScanToString(TableFunctionToStringInput &input) { + InsertionOrderPreservingMap result; + result["Function"] = "PYTHON_ARROW_SCAN"; + if (!input.bind_data) { + return result; + } + auto &bind_data = input.bind_data->Cast(); + string filters; + for (auto &predicate : bind_data.predicates) { + if (!filters.empty()) { + filters += "\n"; + } + filters += predicate.text; + } + if (!filters.empty()) { + result["Filters"] = filters; + } + return result; +} + +struct PythonArrowScanGlobalState : public GlobalTableFunctionState { + idx_t MaxThreads() const override { + return max_threads; + } + + string name; + ArrowSourceKind kind = ArrowSourceKind::REGISTERED; + ArrowOwned stream; + //! The one batch, until a thread claims it + ArrowOwned single; + ArrowOwned schema; + shared_ptr arrow_table; + //! The stream column each output column reads + vector picks; + //! Whether the stream's arrays are plain values, each to be read as a batch of one column + bool wrap = false; + //! Whether the source asked for the GIL around pulling the next array and releasing one + bool python_callbacks = false; + //! Started by an exit handler that runs after the exit hook, while other threads can still take the GIL + bool started_during_exit = false; + idx_t max_threads = 1; + + //! Guards the pull, the claim of a batch index and the claim of the single batch + mutex pull_lock; + idx_t next_batch = 0; + bool exhausted = false; +}; + +//! Exports the stream, narrowed to the requested columns when the source can, and checks it against the declared +//! columns by width, name and type: the scan reads columns by position, so a mismatch is refused rather than read. +unique_ptr PythonArrowScanInitGlobal(ClientContext &context, TableFunctionInitInput &input) { + auto &bind_data = input.bind_data->Cast(); + auto &info = *bind_data.info; + const auto declared = bind_data.types.size(); + vector requested; + bool identity = input.column_ids.size() == declared; + for (idx_t i = 0; i < input.column_ids.size(); i++) { + auto column = input.column_ids[i]; + if (column >= declared) { + throw InternalException("python_arrow_scan was asked for column %d, which it does not declare", column); + } + bind_data.arrow_table.GetColumns().at(column)->ThrowIfInvalid(); + requested.push_back(column); + identity = identity && column == i; + } + auto result = make_uniq(); + result->name = info.name; + result->kind = info.kind; + const auto subject = Subject(info.name, info.kind, "object"); + const auto source_subject = Subject(info.name, info.kind, "source"); + + nb::gil_scoped_acquire gil; + engine::CallbackScope callback(context); + result->started_during_exit = StartedDuringExit(); + // Kept only once every check passed, so a scan refused before reading a row leaves a single use reader readable. + unique_ptr mark; + Exported exported; + try { + auto &source = info.source.obj; + auto underlying = nb::getattr(source, "obj", source); + // A producer that runs Python takes the GIL itself, as pyarrow's, ctypes' and Cython's callbacks do; holding it + // across a pull instead deadlocks a producer whose own threads need Python, such as a reader over a Python + // file. + result->python_callbacks = nb::getattr(source, "pull_under_gil", nb::handle(Py_False)).ptr() == Py_True; + mark = make_uniq(underlying, IsSingleUseArrowReader(underlying)); + exported = ExportStream(info, identity ? nullptr : &requested, bind_data.predicates); + } catch (nb::python_error &error) { + throw InvalidInputException("exporting a stream from %s failed: %s", subject, FormatUDFPythonError(error)); + } + if (exported.IsArray()) { + result->schema = TakeFromCapsule(exported.schema, SCHEMA_CAPSULE, subject); + auto single = TakeFromCapsule(exported.array, ARRAY_CAPSULE, subject); + RouteRelease(single.value, result->python_callbacks, result->started_during_exit); + single.release_under_gil = false; + single.admitted_during_exit = result->started_during_exit; + result->single = std::move(single); + } else { + // Only borrowed until every check passed, so a scan that fails here leaves a stream the source keeps in place. + auto &borrowed = InCapsule(exported.stream, STREAM_CAPSULE, subject); + auto code = borrowed.get_schema(&borrowed, &result->schema.value); + if (code != 0) { + result->schema.value.release = nullptr; + throw InvalidInputException("reading the schema of %s failed: %s", Subject(info.name, info.kind, "stream"), + StreamError(borrowed, code)); + } + } + result->schema.release_under_gil = result->python_callbacks; + result->schema.admitted_during_exit = result->started_during_exit; + auto &schema = result->schema.value; + result->wrap = !IsBatch(schema); + if (result->wrap) { + WrapAsBatch(schema); + if (result->single) { + WrapAsBatch(result->single.value); + } + } + const auto fields = NumericCast(schema.n_children); + if (exported.projected) { + if (fields != requested.size()) { + throw InvalidInputException("%s answered with %d columns where %d were requested", source_subject, fields, + requested.size()); + } + for (idx_t i = 0; i < fields; i++) { + result->picks.push_back(i); + } + } else { + if (fields != declared) { + throw InvalidInputException("%s answered with %d columns where it declared %d", source_subject, fields, + declared); + } + result->picks = requested; + } + result->arrow_table = make_shared_ptr(); + ResolveBatchSchema(context, *result->arrow_table, schema); + auto &types = result->arrow_table->GetTypes(); + for (idx_t i = 0; i < fields; i++) { + const auto expected = exported.projected ? requested[i] : i; + const auto name = NameOf(schema, i); + if (name != bind_data.raw_names[expected]) { + throw InvalidInputException("%s answered with column '%s' at position %d where '%s' was expected", + source_subject, name, i, bind_data.raw_names[expected]); + } + if (types[i] != bind_data.types[expected]) { + throw InvalidInputException("%s answered with column '%s' of type %s where %s was declared", source_subject, + name, types[i].ToString(), bind_data.types[expected].ToString()); + } + } + // A single batch is one claim, which only one thread can ever work on. + result->max_threads = exported.IsArray() ? 1 : STREAM_MAX_THREADS; + if (!exported.IsArray()) { + auto stream = TakeFromCapsule(exported.stream, STREAM_CAPSULE, subject); + stream.release_under_gil = result->python_callbacks; + stream.admitted_during_exit = result->started_during_exit; + result->stream = std::move(stream); + } + mark->Keep(); + return std::move(result); +} + +struct PythonArrowScanLocalState : public LocalTableFunctionState { + PythonArrowScanLocalState(ClientContext &context, shared_ptr schema, vector columns) + : importer(context, std::move(schema), std::move(columns)) { + } + + ArrowBatchImporter importer; + idx_t batch_index = 0; +}; + +unique_ptr PythonArrowScanInitLocal(ExecutionContext &context, TableFunctionInitInput &input, + GlobalTableFunctionState *global_state) { + auto &global = global_state->Cast(); + return make_uniq(context.client, global.arrow_table, global.picks); +} + +void PullNext(ClientContext &context, PythonArrowScanGlobalState &global, ArrowArray &out) { + auto &stream = global.stream.value; + // A pyarrow reader over a generator runs that Python code on this thread, so the scope is set either way. + engine::CallbackScope callback(context); + SourceCall call(global.started_during_exit); + if (!call) { + throw InvalidInputException("reading %s failed: the Python interpreter is shutting down", + Subject(global.name, global.kind, "stream")); + } + int code; + string error; + if (global.python_callbacks) { + nb::gil_scoped_acquire gil; + code = stream.get_next(&stream, &out); + if (code != 0) { + error = StreamError(stream, code); + } + } else { + code = stream.get_next(&stream, &out); + if (code != 0) { + error = StreamError(stream, code); + } + } + if (code != 0) { + out.release = nullptr; + throw InvalidInputException("reading %s failed: %s", Subject(global.name, global.kind, "stream"), error); + } + if (out.release) { + RouteRelease(out, global.python_callbacks, global.started_during_exit); + } +} + +void PythonArrowScanExec(ClientContext &context, TableFunctionInput &data, DataChunk &output) { + auto &global = data.global_state->Cast(); + auto &local = data.local_state->Cast(); + while (!local.importer.Next(output)) { + ArrowArrayWrapper pulled; + { + lock_guard guard(global.pull_lock); + if (global.exhausted) { + return; + } + if (!global.stream) { + pulled.arrow_array = global.single.value; + global.single.value.release = nullptr; + global.exhausted = true; + } else { + try { + PullNext(context, global, pulled.arrow_array); + } catch (...) { + // A stream that failed is not pulled again by another thread while the query is cancelled. + global.exhausted = true; + throw; + } + if (!pulled.arrow_array.release) { + global.exhausted = true; + } else if (global.wrap) { + WrapAsBatch(pulled.arrow_array); + } + } + if (pulled.arrow_array.release) { + local.batch_index = global.next_batch++; + } + } + if (!pulled.arrow_array.release) { + continue; + } + local.importer.Append(pulled.arrow_array); + } +} + +//! A thread claims increasing indexes from one counter, so its batch index never goes down, as the engine requires. +OperatorPartitionData PythonArrowScanPartitionData(ClientContext &context, TableFunctionGetPartitionInput &input) { + if (input.partition_info.RequiresPartitionColumns()) { + throw InternalException("python_arrow_scan: partition columns are not supported"); + } + return OperatorPartitionData(input.local_state->Cast().batch_index); +} + +} // namespace + +void WatchInterpreterExit() { +#ifndef _WIN32 + pthread_atfork(nullptr, nullptr, &AfterForkInChild); +#endif + nb::module_::import_("atexit").attr("register")(nb::cpp_function(&OnInterpreterExit)); +} + +void RegisterPythonArrowScan(DatabaseInstance &db) { + TableFunction function("python_arrow_scan", {}, PythonArrowScanExec, PythonArrowScanBind, PythonArrowScanInitGlobal, + PythonArrowScanInitLocal); + function.cardinality = PythonArrowScanCardinality; + function.pushdown_complex_filter = PythonArrowScanPushdown; + function.to_string = PythonArrowScanToString; + function.get_partition_data = PythonArrowScanPartitionData; + function.projection_pushdown = true; + function.filter_pushdown = false; + function.filter_prune = false; + + TableFunctionSet set(function.name); + set.AddFunction(std::move(function)); + CreateTableFunctionInfo info(std::move(set)); + info.on_conflict = OnCreateConflict::ALTER_ON_CONFLICT; + auto &catalog = Catalog::GetSystemCatalog(db); + auto transaction = CatalogTransaction::GetSystemTransaction(db); + catalog.CreateFunction(transaction, info); +} + +string Subject(const string &name, ArrowSourceKind kind, const char *noun) { + if (name.empty() || kind == ArrowSourceKind::NAMELESS) { + return "the Arrow source"; + } + if (kind == ArrowSourceKind::VARIABLE) { + if (std::strcmp(noun, "stream") == 0) { + return StringUtil::Format("the stream of the Python object '%s'", name); + } + return StringUtil::Format("the Python object '%s'", name); + } + return StringUtil::Format("the %s registered as '%s'", noun, name); +} + +unique_ptr CreatePythonArrowScanRef(nb::object source, const string &name, ArrowSourceKind kind) { + auto ref = make_uniq(); + vector> children; + ref->function = make_uniq("python_arrow_scan", std::move(children)); + auto dependency = make_shared_ptr(); + dependency->AddDependency("replacement_cache", PythonDependencyItem::Create(source)); + ref->external_dependency = std::move(dependency); + ref->bind_info = make_shared_ptr(std::move(source), name, kind); + return ref; +} + +} // namespace duckdb diff --git a/src/arrow/single_use.cpp b/src/arrow/single_use.cpp new file mode 100644 index 00000000..06ca7650 --- /dev/null +++ b/src/arrow/single_use.cpp @@ -0,0 +1,108 @@ +#include "duckdb_python/arrow/single_use.hpp" + +#include "duckdb/common/exception.hpp" +#include "duckdb/common/unordered_set.hpp" +#include "duckdb/common/vector.hpp" +#include "duckdb_python/exceptions.hpp" +#include "duckdb_python/pyconnection/pyconnection.hpp" + +namespace duckdb { + +static constexpr const char *CONSUMED_READER_ERROR = + "This RecordBatchReader has already been consumed and cannot be scanned again. A RecordBatchReader can be read " + "only once; to scan its data more than once, read it into a table first, for example with read_all()."; + +namespace { + +//! The walk only compares names, so a missing or unreadable one is an empty text rather than an error +string TypeText(nb::object value) { + if (!value.is_valid() || !nb::str::check_(value)) { + PyErr_Clear(); + return string(); + } + try { + return ToUTF8(value); + } catch (nb::python_error &) { + return string(); + } +} + +//! The fallback for a reader whose pyarrow left sys.modules: an extension type's module and name come from tp_name, +//! while a Python class can forge both through __module__ and __qualname__, so it is never the reader. +bool IsPyarrowReaderType(PyTypeObject *type) { + if (PyType_GetFlags(type) & Py_TPFLAGS_HEAPTYPE) { + return false; + } + auto name = TypeText(nb::steal(PyType_GetName(type))); + auto module = TypeText(nb::steal(PyObject_GetAttrString(reinterpret_cast(type), "__module__"))); + return name == "RecordBatchReader" && module == "pyarrow.lib"; +} + +//! Walks the bases the type was created with, not __mro__, which a metaclass can replace; each ancestor once, since +//! a diamond hierarchy repeats them +bool DerivesFromPyarrowReader(PyTypeObject *type) { + vector pending {type}; + unordered_set seen; + while (!pending.empty()) { + auto current = pending.back(); + pending.pop_back(); + if (!seen.insert(current).second) { + continue; + } + if (IsPyarrowReaderType(current)) { + return true; + } + auto bases = static_cast(PyType_GetSlot(current, Py_tp_bases)); + if (!bases || !PyTuple_Check(bases)) { + PyErr_Clear(); + continue; + } + auto count = PyTuple_Size(bases); + for (Py_ssize_t i = 0; i < count; i++) { + pending.push_back(reinterpret_cast(PyTuple_GetItem(bases, i))); + } + } + return false; +} + +} // namespace + +bool IsSingleUseArrowReader(nb::handle object) { + if (ModuleIsLoaded()) { + auto reader_type = DuckDBPyConnection::ImportCache()->pyarrow.RecordBatchReader(); + if (reader_type.ptr() && PyType_Check(reader_type.ptr())) { + return PyType_IsSubtype(Py_TYPE(object.ptr()), reinterpret_cast(reader_type.ptr())) != 0; + } + } + return DerivesFromPyarrowReader(Py_TYPE(object.ptr())); +} + +ConsumedReaderMark::ConsumedReaderMark(nb::handle reader, bool single_use) { + if (!single_use) { + return; + } + auto readers = DuckDBPyConnection::ConsumedArrowReaders(); + key = nb::steal(PyLong_FromVoidPtr(reader.ptr())); + if (!key) { + throw nb::python_error(); + } + auto forget = nb::module_::import_("functools").attr("partial")(readers.attr("pop"), key); + auto entry = nb::steal(PyWeakref_NewRef(reader.ptr(), forget.ptr())); + if (!entry) { + throw nb::python_error(); + } + // One setdefault tests and records at once; the entry leaves with the reader, before its id can be reused. + auto recorded = readers.attr("setdefault")(key, entry); + if (!recorded.is(entry)) { + throw InvalidInputException(CONSUMED_READER_ERROR); + } + record = std::move(readers); +} + +ConsumedReaderMark::~ConsumedReaderMark() { + if (record && !kept && PyDict_DelItem(record.ptr(), key.ptr()) < 0) { + PyErr_Clear(); + } +} + +} // namespace duckdb diff --git a/src/duckdb_python.cpp b/src/duckdb_python.cpp index 898b6b4c..be01c506 100644 --- a/src/duckdb_python.cpp +++ b/src/duckdb_python.cpp @@ -19,6 +19,8 @@ #include "duckdb_python/nb/conversions/python_csv_line_terminator_enum.hpp" #include "duckdb/common/enums/statement_type.hpp" #include "duckdb/common/adbc/adbc-init.hpp" +#include "duckdb_python/arrow/predicate_translation.hpp" +#include "duckdb_python/arrow/python_arrow_scan.hpp" #ifndef DUCKDB_PYTHON_LIB_NAME #define DUCKDB_PYTHON_LIB_NAME _duckdb @@ -1136,6 +1138,9 @@ NB_MODULE(DUCKDB_PYTHON_LIB_NAME, m) { // NOLINT "Tokenizes a SQL string, returning a list of (position, type) tuples that can be " "used for e.g., syntax highlighting", nb::arg("query")); + m.def("_set_arrow_sources", &DuckDBPyConnection::SetArrowSources, nb::arg("sources")); + m.def("_translate_predicate", &TranslatePredicateForTesting, nb::arg("connection"), nb::arg("predicate"), + nb::arg("columns").none() = nb::none(), nb::arg("schema").none() = nb::none(), nb::arg("select") = "*"); nb::enum_(m, "token_type") .value("identifier", PySQLTokenType::PY_SQL_TOKEN_IDENTIFIER) .value("numeric_const", PySQLTokenType::PY_SQL_TOKEN_NUMERIC_CONSTANT) @@ -1151,6 +1156,7 @@ NB_MODULE(DUCKDB_PYTHON_LIB_NAME, m) { // NOLINT static char clean_default_connection_sentinel; m.attr("_clean_default_connection") = nb::capsule(&clean_default_connection_sentinel, [](void *) noexcept { DuckDBPyConnection::Cleanup(); }); + WatchInterpreterExit(); } } // namespace duckdb diff --git a/src/include/duckdb_python/arrow/arrow_array_stream.hpp b/src/include/duckdb_python/arrow/arrow_array_stream.hpp deleted file mode 100644 index 87ff1f79..00000000 --- a/src/include/duckdb_python/arrow/arrow_array_stream.hpp +++ /dev/null @@ -1,114 +0,0 @@ -//===----------------------------------------------------------------------===// -// DuckDB -// -// duckdb_python/arrow/arrow_array_stream.hpp -// -// -//===----------------------------------------------------------------------===// - -#pragma once - -#include "duckdb/common/arrow/arrow_wrapper.hpp" -#include "duckdb/common/atomic.hpp" -#include "duckdb/common/constants.hpp" -#include "duckdb/function/table_function.hpp" -#include "duckdb/function/table/arrow.hpp" -#include "duckdb/main/client_config.hpp" -#include "duckdb/main/config.hpp" -#include "duckdb_python/nb/casters.hpp" -#include "duckdb_python/registered_py_object.hpp" - -#include "duckdb/common/string.hpp" -#include "duckdb/common/vector.hpp" -#include "duckdb/main/client_properties.hpp" - -namespace duckdb { - -namespace pyarrow { - -class RecordBatchReader : public nb::object { -public: - RecordBatchReader(const nb::object &o) : nb::object(o, nb::detail::borrow_t {}) { - } - using nb::object::object; - -public: - static bool check_(const nb::handle &object) { - return !nb::none().is(object); - } -}; - -class Table : public nb::object { -public: - Table(const nb::object &o) : nb::object(o, nb::detail::borrow_t {}) { - } - using nb::object::object; - -public: - static bool check_(const nb::handle &object) { - return !nb::none().is(object); - } -}; - -} // namespace pyarrow - -enum class PyArrowObjectType { - Invalid, - Table, - Scanner, - Dataset, - PyCapsule, - PyCapsuleInterface, - MessageReader, - PolarsLazyFrame -}; - -void TransformDuckToArrowChunk(nb::object pyarrow_schema, ArrowArray &data, nb::list &batches); - -PyArrowObjectType GetArrowType(const nb::handle &obj); - -class PythonTableArrowArrayStreamFactory : public ArrowScanFactory { -public: - //! Must be constructed while holding the GIL. - PythonTableArrowArrayStreamFactory(nb::object arrow_object_p, const ClientProperties &client_properties_p, - PyArrowObjectType arrow_type_p) - : arrow_object(std::move(arrow_object_p)), client_properties(client_properties_p), - cached_arrow_type(arrow_type_p), single_use(IsSingleUse(arrow_object.obj, arrow_type_p)) { - cached_schema.release = nullptr; - } - - ~PythonTableArrowArrayStreamFactory() override { - // The release callback of a schema taken from a Python producer may itself be Python. - if (cached_schema.release) { - if (nb::detail::cleanup_guard guard {}) { - cached_schema.release(&cached_schema); - } - } - } - - void GetSchema(ArrowSchema &schema) override; - unique_ptr ProduceStream(ArrowStreamParameters ¶meters) override; - - static void GetSchemaInternal(nb::handle arrow_object, ArrowSchema &schema); - - //! Arrow Object (i.e., Scanner, Record Batch Reader, Table, Dataset) - PyObjectHolder arrow_object; - - const ClientProperties client_properties; - const PyArrowObjectType cached_arrow_type; - //! A pyarrow RecordBatchReader, which hands out its batches once however many streams it exports - const bool single_use; - - //! Cached Arrow table from an unfiltered .collect().to_arrow() on a LazyFrame. - //! Avoids re-reading from source and re-converting on repeated scans without filters. - PyObjectHolder cached_arrow_table; - -private: - ArrowSchema cached_schema; - atomic schema_cached {false}; - - static nb::object ProduceScanner(nb::object &arrow_scanner, nb::handle &arrow_obj_handle, - ArrowStreamParameters ¶meters, const ClientProperties &client_properties); - static bool IsSingleUse(nb::handle arrow_object, PyArrowObjectType type); -}; -} // namespace duckdb diff --git a/src/include/duckdb_python/arrow/arrow_export_utils.hpp b/src/include/duckdb_python/arrow/arrow_export_utils.hpp index e5514e0a..26636e03 100644 --- a/src/include/duckdb_python/arrow/arrow_export_utils.hpp +++ b/src/include/duckdb_python/arrow/arrow_export_utils.hpp @@ -1,11 +1,37 @@ #pragma once +#include "duckdb/common/arrow/arrow.hpp" +#include "duckdb/main/client_properties.hpp" #include "duckdb_python/nb/casters.hpp" namespace duckdb { namespace pyarrow { +class RecordBatchReader : public nb::object { +public: + RecordBatchReader(const nb::object &o) : nb::object(o, nb::detail::borrow_t {}) { + } + using nb::object::object; + +public: + static bool check_(const nb::handle &object) { + return !nb::none().is(object); + } +}; + +class Table : public nb::object { +public: + Table(const nb::object &o) : nb::object(o, nb::detail::borrow_t {}) { + } + using nb::object::object; + +public: + static bool check_(const nb::handle &object) { + return !nb::none().is(object); + } +}; + nb::object ToPyArrowSchema(const ArrowSchema &schema); nb::object ToArrowTable(const vector &types, const vector &names, const nb::list &batches, @@ -15,4 +41,6 @@ nb::object ToArrowTable(const nb::list &batches, nb::object pyarrow_schema); } // namespace pyarrow +void TransformDuckToArrowChunk(nb::object pyarrow_schema, ArrowArray &data, nb::list &batches); + } // namespace duckdb diff --git a/src/include/duckdb_python/arrow/arrow_importer.hpp b/src/include/duckdb_python/arrow/arrow_importer.hpp new file mode 100644 index 00000000..cfc49290 --- /dev/null +++ b/src/include/duckdb_python/arrow/arrow_importer.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include "duckdb/common/arrow/arrow_wrapper.hpp" +#include "duckdb/common/types/data_chunk.hpp" +#include "duckdb/common/unordered_map.hpp" +#include "duckdb/function/table/arrow.hpp" +#include "duckdb/function/table/arrow/arrow_duck_schema.hpp" + +namespace duckdb { + +//! Converts Arrow record batches into DataChunks of at most STANDARD_VECTOR_SIZE rows against one resolved schema, one +//! batch at a time, referencing the batch's buffers where the engine's conversion allows. The batch's own validity +//! bitmap makes a whole row NULL. Not thread safe: each scanning thread owns one. +class ArrowBatchImporter { +public: + //! Output column i is read from batch column columns[i] + ArrowBatchImporter(ClientContext &context, shared_ptr schema, vector columns); + + //! Takes the batch over, clearing its release, after checking its shape against the schema; the previous batch + //! is dropped, drained or not + void Append(ArrowArray &batch); + //! Fills `output`, whose columns match the requested ones, with the next rows of the current batch; false once + //! the batch is drained + bool Next(DataChunk &output); + +private: + ArrowArrayScanState &StateOf(idx_t column); + + ClientContext &context; + const shared_ptr schema; + const arrow_column_map_t &arrow_columns; + const vector columns; + shared_ptr batch; + idx_t offset = 0; + idx_t length = 0; + //! Kept while one batch is drained, so its dictionaries are decoded once rather than once per chunk + unordered_map> states; +}; + +} // namespace duckdb diff --git a/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp b/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp deleted file mode 100644 index 0f4b205e..00000000 --- a/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp +++ /dev/null @@ -1,100 +0,0 @@ -//===----------------------------------------------------------------------===// -// DuckDB -// -// duckdb_python/arrow/filter_pushdown_visitor.hpp -// -// -//===----------------------------------------------------------------------===// - -#pragma once - -#include "duckdb/common/types/value.hpp" -#include "duckdb/function/table/arrow/arrow_duck_schema.hpp" -#include "duckdb/planner/expression.hpp" -#include "duckdb/planner/table_filter.hpp" -#include "duckdb_python/nb/casters.hpp" - -namespace duckdb { - -// A FilterBackend abstracts the Python side of an `ExpressionFilter` → -// expression translation. The shared walker in this file handles the -// structural recursion (CONJUNCTION_AND/OR, struct_extract column paths, the -// optional / selectivity-optional filter wrappers, and the internal runtime -// filter functions) and dispatches leaf operations to the backend. -// -// Two backends exist today: PyArrowBackend (emits pyarrow.dataset.Expression) -// and PolarsBackend (emits polars.Expr). Adding a new backend is purely a -// matter of implementing this interface; the walker itself is reused. -// -// Convention: a backend method that cannot push the given filter must throw -// `NotImplementedException`. The optional-filter wrapper swallows it (an optional -// filter is not required for correctness) and returns `nb::none()`, dropping the -// filter. Throwing keeps the "I can't push this" path uniform across backends, -// replacing the old polars walker's ad hoc `return nb::none()` style. -struct FilterBackend { - virtual ~FilterBackend() = default; - - // Build a column expression from an accumulated path. `path` always has - // at least one element (the top-level column). For nested struct - // references the path accumulates one entry per `struct_extract`. - virtual nb::object MakeColumnRef(const vector &path) = 0; - - // Convert a constant to a backend scalar, where the column's `arrow_type` (nullable) gives the decimal bit width - // and the blob view distinction, its `arrow_schema` node (nullable) the time zone, and `timezone_config` the - // session's. - virtual nb::object MakeScalar(const Value &v, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, - const string &timezone_config) = 0; - - // Apply a comparison operator. `op` is one of the COMPARE_* ExpressionTypes. - // `scalar` is what MakeScalar returned. NaN special cases go through - // NaNCompare instead. - virtual nb::object Compare(ExpressionType op, nb::object col, nb::object scalar) = 0; - - // NaN-specific comparison. DuckDB treats NaN as the greatest value, so - // each operator decomposes into is_nan / ~is_nan / is_valid / lit(true|false). - virtual nb::object NaNCompare(ExpressionType op, nb::object col) = 0; - - // Column-side NaN predicate, true on NaN rows. Used to re-include NaN rows for `>` / `>=` against a - // finite float constant, since DuckDB orders NaN as the greatest value (so `nan > finite` is TRUE) - // while IEEE comparisons make them FALSE. - virtual nb::object IsNaN(nb::object col) = 0; - - virtual nb::object IsNull(nb::object col) = 0; - virtual nb::object IsNotNull(nb::object col) = 0; - - // IN list. `col_logical_type` is the column's DuckDB logical type: polars - // uses it to construct a typed Series with matching precision/scale for - // decimal columns; pyarrow uses it to detect float columns and decline NaN - // lists, and to match zero constants with == 0.0. - virtual nb::object IsIn(nb::object col, const vector &values, const LogicalType &col_logical_type, - const string &timezone_config) = 0; - - virtual nb::object And(nb::object a, nb::object b) = 0; - virtual nb::object Or(nb::object a, nb::object b) = 0; -}; - -// Walk a TableFilter and emit a backend-specific expression. Since the -// table-filter -> expression-filter migration in core, the only runtime filter -// type is `EXPRESSION_FILTER`; this unwraps it and walks the expression tree. -// - `column_path` is the top-level column name; struct paths are accumulated -// inside the expression walk via struct_extract. -// - `arrow_type` is the column's root ArrowType (nullable for backends that don't track Arrow types); ResolveColumn -// finds the leaf. -// - `arrow_schema` is the column's root node in the ArrowSchema the scan exported (nullable like `arrow_type`). -// - Returns `nb::none()` if no part of the filter could be pushed. -nb::object TransformFilter(const TableFilter &filter, const vector &column_path, FilterBackend &backend, - const ArrowType *arrow_type, const ArrowSchema *arrow_schema, const string &timezone_config); - -// Walk a bound Expression tree (the contents of an `ExpressionFilter`) and emit -// a backend-specific expression. Handles BOUND_FUNCTION comparisons, -// BOUND_OPERATOR (IS_NULL / IS_NOT_NULL / COMPARE_IN), BOUND_CONJUNCTION -// (AND/OR), struct_extract column chains, the optional / selectivity-optional -// wrappers (unwrapped from `bind_info`; an untranslatable child is swallowed), -// and the internal runtime filter functions (dynamic / bloom / perfect-hash-join -// / prefix-range, which are skipped). Returns `nb::none()` for an optional or -// runtime filter that can't be pushed. -nb::object TransformExpression(const Expression &expression, const vector &column_path, - FilterBackend &backend, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, - const string &timezone_config); - -} // namespace duckdb diff --git a/src/include/duckdb_python/arrow/polars_filter_pushdown.hpp b/src/include/duckdb_python/arrow/polars_filter_pushdown.hpp deleted file mode 100644 index 5fb37ca7..00000000 --- a/src/include/duckdb_python/arrow/polars_filter_pushdown.hpp +++ /dev/null @@ -1,23 +0,0 @@ -//===----------------------------------------------------------------------===// -// DuckDB -// -// duckdb_python/arrow/polars_filter_pushdown.hpp -// -// -//===----------------------------------------------------------------------===// - -#pragma once - -#include "duckdb/planner/table_filter_set.hpp" -#include "duckdb/main/client_properties.hpp" -#include "duckdb_python/nb/casters.hpp" - -namespace duckdb { - -struct PolarsFilterPushdown { - static nb::object TransformFilter(const TableFilterSet &filter_collection, unordered_map &columns, - const unordered_map &filter_to_col, - const ClientProperties &client_properties); -}; - -} // namespace duckdb diff --git a/src/include/duckdb_python/arrow/predicate_translation.hpp b/src/include/duckdb_python/arrow/predicate_translation.hpp new file mode 100644 index 00000000..d4c6c606 --- /dev/null +++ b/src/include/duckdb_python/arrow/predicate_translation.hpp @@ -0,0 +1,63 @@ +#pragma once + +#include "duckdb/common/arrow/arrow.hpp" +#include "duckdb/common/types.hpp" +#include "duckdb/planner/expression.hpp" +#include "duckdb_python/nb/casters.hpp" + +#include +#include + +namespace duckdb { + +struct DuckDBPyConnection; + +//! A column a predicate refers to: the name the source itself uses for it, and the type the scan declared. +struct PredicateColumn { + string name; + LogicalType type; +}; + +//! A predicate, or a part of one, that cannot be handed to a source exactly; the engine keeps applying all of it. +struct PredicateRefused : std::exception {}; + +//! The declared columns of a scan as its source names them. A source can only apply a predicate to a column it can +//! name, so a column whose source name is repeated, or whose struct fields the importer renamed (an empty or +//! repeated field name inside a struct), resolves to a refusal. An empty source name is kept: the one column of a +//! nameless array is named by it. +class PredicateColumns { +public: + PredicateColumns() = default; + //! Columns whose source names the scan declared unchanged, struct fields included. + PredicateColumns(const vector &names, const vector &types); + //! Columns of an Arrow source: `schema` is the source's top level struct schema and `types` the DuckDB types the + //! scan declared for its children, in the same order. + PredicateColumns(const ArrowSchema &schema, const vector &types); + + //! @throws PredicateRefused For a column the source cannot name exactly. + PredicateColumn Resolve(idx_t declared_index) const; + +private: + vector columns; + vector nameable; +}; + +//! Turns a predicate the optimizer offers to a scan into duckdb._sources.nodes objects: comparisons of a column with +//! a constant (the column put on the left), BETWEEN as two comparisons under AND, IN and NOT IN over constants, +//! IS NULL, IS NOT NULL, AND, OR and NOT. A column is a column reference, or a chain of struct field accesses on one. +//! A constant is carried only when it has exactly the column's type and a Python value holds it without loss. +//! `resolve` maps a column reference's binding column index to the declared column and may throw PredicateRefused. +//! Call with the GIL held. +//! @throws PredicateRefused For any other shape, and for a constant without an exact Python value of the column's type. +nb::object TranslatePredicate(const Expression &predicate, + const std::function &resolve); + +//! Plans `SELECT