diff --git a/src/arrow/arrow_array_stream.cpp b/src/arrow/arrow_array_stream.cpp index eb2df04d..04299941 100644 --- a/src/arrow/arrow_array_stream.cpp +++ b/src/arrow/arrow_array_stream.cpp @@ -167,8 +167,9 @@ nb::object PythonTableArrowArrayStreamFactory::ProduceScanner(nb::object &arrow_ } if (has_filter) { - auto filter = PyArrowFilterPushdown::TransformFilter(*filters, parameters.projected_columns.projection_map, - filter_to_col, client_properties, arrow_table); + 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; } diff --git a/src/arrow/filter_pushdown_visitor.cpp b/src/arrow/filter_pushdown_visitor.cpp index 5311b685..385a8540 100644 --- a/src/arrow/filter_pushdown_visitor.cpp +++ b/src/arrow/filter_pushdown_visitor.cpp @@ -23,18 +23,37 @@ bool ValueIsNan(const Value &value) { return false; } -// ResolveColumn walks a column-side expression to extract the (full path, leaf -// ArrowType) pair. Accepts a bare BoundReferenceExpression and (nested) +// 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; }; -ResolvedColumn ResolveColumn(const Expression &expr, const vector &root_path, const ArrowType *root_type) { +// 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, root_type}; + 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", @@ -47,20 +66,24 @@ ResolvedColumn ResolveColumn(const Expression &expr, const vector &r ExpressionTypeToString(expr.GetExpressionType())); } // Recurse innermost-first so names accumulate root → leaf. - auto inner = ResolveColumn(*func.GetChildren()[0], root_path, root_type); + 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 = &inner.leaf_type->GetTypeInfo().GetChild(child_idx); + 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 ArrowType *arrow_type, const string &timezone_config) { + const ResolvedColumn &resolved, const string &timezone_config) { if (ValueIsNan(constant)) { return backend.NaNCompare(op, std::move(col)); } - auto scalar = backend.MakeScalar(constant, arrow_type, timezone_config); + 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 @@ -85,7 +108,8 @@ nb::object EmitCompare(FilterBackend &backend, ExpressionType op, nb::object col } // anonymous namespace nb::object TransformExpression(const Expression &expression, const vector &column_path, - FilterBackend &backend, const ArrowType *arrow_type, const string &timezone_config) { + FilterBackend &backend, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, + const string &timezone_config) { auto expression_class = expression.GetExpressionClass(); auto expression_type = expression.GetExpressionType(); @@ -109,9 +133,9 @@ nb::object TransformExpression(const Expression &expression, const vectorGetValue(), resolved.leaf_type, + return EmitCompare(backend, expression_type, std::move(col), constant_side->GetValue(), resolved, timezone_config); } @@ -143,7 +167,7 @@ nb::object TransformExpression(const Expression &expression, const vector(); if (expression_type == ExpressionType::OPERATOR_IS_NULL) { - auto resolved = ResolveColumn(*op_expr.GetChildren()[0], column_path, arrow_type); + 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); + 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); + 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++) { @@ -188,8 +212,8 @@ nb::object TransformExpression(const Expression &expression, const vector(); 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, timezone_config); + 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 @@ -219,11 +243,13 @@ nb::object TransformExpression(const Expression &expression, const vector &column_path, FilterBackend &backend, - const ArrowType *arrow_type, const string &timezone_config) { + 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, timezone_config); + 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", diff --git a/src/arrow/polars_filter_pushdown.cpp b/src/arrow/polars_filter_pushdown.cpp index acad083d..f23a5449 100644 --- a/src/arrow/polars_filter_pushdown.cpp +++ b/src/arrow/polars_filter_pushdown.cpp @@ -24,9 +24,11 @@ struct PolarsBackend : public FilterBackend { return col; } - nb::object MakeScalar(const Value &v, const ArrowType *arrow_type, const string &timezone_config) override { + 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); } @@ -137,9 +139,9 @@ nb::object PolarsFilterPushdown::TransformFilter(const TableFilterSet &filter_co vector column_path = {Identifier(column_name)}; // Polars does not need ArrowType information — `nullptr` here propagates through the - // shared walker; the PolarsBackend ignores the parameter in MakeScalar. + // 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, - client_properties.time_zone); + nullptr, client_properties.time_zone); if (child_expression.is(nb::none())) { continue; } diff --git a/src/arrow/pyarrow_filter_pushdown.cpp b/src/arrow/pyarrow_filter_pushdown.cpp index f4f6a2d5..e2a57d19 100644 --- a/src/arrow/pyarrow_filter_pushdown.cpp +++ b/src/arrow/pyarrow_filter_pushdown.cpp @@ -11,45 +11,23 @@ namespace duckdb { namespace { -string ConvertTimestampUnit(ArrowDateTimeType unit) { - switch (unit) { - case ArrowDateTimeType::MICROSECONDS: - return "us"; - case ArrowDateTimeType::MILLISECONDS: - return "ms"; - case ArrowDateTimeType::NANOSECONDS: - return "ns"; - case ArrowDateTimeType::SECONDS: - return "s"; - default: - throw NotImplementedException("DatetimeType not recognized in ConvertTimestampUnit: %d", - static_cast(unit)); - } -} - -int64_t ConvertTimestampTZValue(int64_t base_value, ArrowDateTimeType datetime_type) { - auto input = timestamp_t(base_value); - if (!Value::IsFinite(input)) { - return base_value; - } - switch (datetime_type) { - case ArrowDateTimeType::MICROSECONDS: - return Timestamp::GetEpochMicroSeconds(input); - case ArrowDateTimeType::MILLISECONDS: - return Timestamp::GetEpochMs(input); - case ArrowDateTimeType::NANOSECONDS: - return Timestamp::GetEpochNanoSeconds(input); - case ArrowDateTimeType::SECONDS: - return Timestamp::GetEpochSeconds(input); - default: - throw NotImplementedException("DatetimeType not recognized in ConvertTimestampTZValue"); +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 timestamp unit / decimal precision / blob-view disambiguation; the +// 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) { +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"); @@ -101,16 +79,11 @@ nb::object MakePyArrowScalar(const Value &constant, const string &timezone_confi return dataset_scalar(scalar(constant.GetValue(), date_type("s"))); } case LogicalTypeId::TIMESTAMP_TZ: { - if (!arrow_type) { - throw NotImplementedException("Cannot push down TIMESTAMP_TZ filter without an arrow type"); - } - auto &datetime_info = arrow_type->GetTypeInfo(); - auto base_value = constant.GetValue(); - auto arrow_datetime_type = datetime_info.GetDateTimeType(); - auto time_unit_string = ConvertTimestampUnit(arrow_datetime_type); - auto converted_value = ConvertTimestampTZValue(base_value, arrow_datetime_type); + // 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(); - return dataset_scalar(scalar(converted_value, date_type(time_unit_string, nb::arg("tz") = timezone_config))); + 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(); @@ -200,8 +173,9 @@ struct PyArrowBackend : public FilterBackend { return field_factory(nb::tuple(nb::cast(str_path))); } - nb::object MakeScalar(const Value &v, const ArrowType *arrow_type, const string &timezone_config) override { - return MakePyArrowScalar(v, timezone_config, arrow_type); + 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 { @@ -228,10 +202,11 @@ struct PyArrowBackend : public FilterBackend { switch (op) { case ExpressionType::COMPARE_EQUAL: case ExpressionType::COMPARE_GREATERTHANOREQUALTO: - return col.attr("is_nan")(); + return IsNaN(std::move(col)); case ExpressionType::COMPARE_LESSTHAN: case ExpressionType::COMPARE_NOTEQUAL: - return col.attr("is_nan")().attr("__invert__")(); + // 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); @@ -245,7 +220,8 @@ struct PyArrowBackend : public FilterBackend { } nb::object IsNaN(nb::object col) override { - return col.attr("is_nan")(); + // 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 { @@ -330,7 +306,8 @@ struct PyArrowBackend : public FilterBackend { nb::object PyArrowFilterPushdown::TransformFilter(TableFilterSet &filter_collection, unordered_map &columns, unordered_map filter_to_col, - const ClientProperties &config, const ArrowTableSchema &arrow_table) { + const ClientProperties &config, const ArrowTableSchema &arrow_table, + const ArrowSchema &arrow_schema) { PyArrowBackend backend(config); nb::object expression = nb::none(); for (auto &entry : filter_collection) { @@ -339,9 +316,12 @@ nb::object PyArrowFilterPushdown::TransformFilter(TableFilterSet &filter_collect D_ASSERT(columns.find(column_idx) != columns.end()); vector column_path = {Identifier(column_name)}; - auto &arrow_type = arrow_table.GetColumns().at(filter_to_col.at(column_idx)); + 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(), config.time_zone); + arrow_type.get(), column_schema, config.time_zone); if (child_expression.is(nb::none())) { continue; } diff --git a/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp b/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp index 65758d6f..0f4b205e 100644 --- a/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp +++ b/src/include/duckdb_python/arrow/filter_pushdown_visitor.hpp @@ -39,11 +39,11 @@ struct FilterBackend { // references the path accumulates one entry per `struct_extract`. virtual nb::object MakeColumnRef(const vector &path) = 0; - // Convert a DuckDB Value to a backend-native Python scalar. `arrow_type` - // may be nullptr for backends that don't need Arrow type information - // (polars relies on DuckDB LogicalType only). `timezone_config` is the - // active session's time zone for `TIMESTAMP_TZ` handling. - virtual nb::object MakeScalar(const Value &v, const ArrowType *arrow_type, const string &timezone_config) = 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 @@ -54,7 +54,7 @@ struct FilterBackend { // 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: `col.is_nan()`. Used to re-include NaN rows for `>` / `>=` against a + // 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; @@ -78,11 +78,12 @@ struct FilterBackend { // 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 ArrowType for the current path leaf (nullable for -// backends that don't track Arrow types). +// - `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 string &timezone_config); + 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, @@ -93,6 +94,7 @@ nb::object TransformFilter(const TableFilter &filter, const vector & // / 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 string &timezone_config); + FilterBackend &backend, const ArrowType *arrow_type, const ArrowSchema *arrow_schema, + const string &timezone_config); } // namespace duckdb diff --git a/src/include/duckdb_python/arrow/pyarrow_filter_pushdown.hpp b/src/include/duckdb_python/arrow/pyarrow_filter_pushdown.hpp index d5d34fe5..0c388faf 100644 --- a/src/include/duckdb_python/arrow/pyarrow_filter_pushdown.hpp +++ b/src/include/duckdb_python/arrow/pyarrow_filter_pushdown.hpp @@ -18,7 +18,7 @@ namespace duckdb { struct PyArrowFilterPushdown { static nb::object TransformFilter(TableFilterSet &filter_collection, unordered_map &columns, unordered_map filter_to_col, const ClientProperties &config, - const ArrowTableSchema &arrow_table); + const ArrowTableSchema &arrow_table, const ArrowSchema &arrow_schema); }; } // namespace duckdb diff --git a/tests/fast/arrow/_pushdown_helpers.py b/tests/fast/arrow/_pushdown_helpers.py index cdc68b1b..037e5fe6 100644 --- a/tests/fast/arrow/_pushdown_helpers.py +++ b/tests/fast/arrow/_pushdown_helpers.py @@ -247,3 +247,19 @@ def was_pushed(con, query: str) -> bool: elif isinstance(filters, list) and len(filters) > 0: return True return False + + +def assert_fully_pushed(con, query: str, expected_count: int) -> None: + """Assert that EXPLAIN ANALYZE of `query` shows its filter inside the only Arrow scan, with nothing above it. + + The scan must emit exactly `expected_count` rows, which shows the pushed filter ran inside the scan. + """ + import json + + plan = json.loads(con.execute(f"EXPLAIN (ANALYZE, FORMAT JSON) {query}").fetchone()[1]) + scans = _find_arrow_scan_nodes(plan["operator"]) + assert len(scans) == 1 + scan, has_filter_above = scans[0] + assert scan["extra_info"].get("Filters"), "filter was not pushed into the Arrow scan" + assert not has_filter_above, "a FILTER operator stayed above the Arrow scan" + assert scan["intermediate_rows"] == expected_count, "the Arrow scan emitted rows the pushed filter rejects" diff --git a/tests/fast/arrow/test_dictionary_filter_pushdown.py b/tests/fast/arrow/test_dictionary_filter_pushdown.py new file mode 100644 index 00000000..dcaba3b2 --- /dev/null +++ b/tests/fast/arrow/test_dictionary_filter_pushdown.py @@ -0,0 +1,784 @@ +"""Filter pushdown on dictionary encoded Arrow columns. + +The Arrow scan never applies a pushed filter again, so a filter pushed on a dictionary encoded column must +compare the dictionary's values, never its integer indices. Every case compares the rows against a native +DuckDB table holding the same data. Where pyarrow cannot evaluate a pushed filter, the query must fail with +an error that leaves the connection usable. +""" + +from __future__ import annotations + +import datetime as dt +from decimal import Decimal +from typing import NamedTuple + +import pytest +from _pushdown_helpers import assert_fully_pushed + +import duckdb + +pa = pytest.importorskip("pyarrow") +pc = pytest.importorskip("pyarrow.compute") +pytest.importorskip("pyarrow.dataset") + +# Dictionary entry 3 is NULL and row 5 has a NULL index, so both kinds of NULL meet every predicate. +INDICES = [0, 1, 2, 3, None, 1, 0] +ROWS = list(range(1, len(INDICES) + 1)) + +UTC = dt.timezone.utc +NEW_YORK = "America/New_York" + + +class DictCase(NamedTuple): + id: str + value_type: object + values: tuple # low, mid, high + literals: tuple # SQL literals for low, mid, high + session_zone: str | None = None + value_field: str | None = None + # "pushed": the engine pushes the comparison into the scan before and after it judges a dictionary by its + # value type. "either": that change moves the comparison above the scan, so the plan is not asserted. + plan: str = "pushed" + + +def _tstz(year, month, day): + return dt.datetime(year, month, day, tzinfo=UTC) + + +def _tstz_literal(value): + return f"TIMESTAMPTZ '{value:%Y-%m-%d %H:%M:%S}+00'" + + +TSTZ_VALUES = (_tstz(2020, 1, 1), _tstz(2020, 6, 1), _tstz(2021, 1, 1)) +TSTZ_LITERALS = tuple(_tstz_literal(v) for v in TSTZ_VALUES) +BLOB_LITERALS = ("'aa'::BLOB", "'bb'::BLOB", "'cc'::BLOB") +DECIMALS = (Decimal("1.25"), Decimal("2.50"), Decimal("3.75")) +STRUCT_TYPE = pa.struct([("i", pa.int64()), ("s", pa.string())]) +STRUCT_VALUES = ({"i": 1, "s": "a"}, {"i": 2, "s": "b"}, {"i": 3, "s": "c"}) + + +def _decimal_literals(width, scale): + return tuple(f"{v}::DECIMAL({width},{scale})" for v in DECIMALS) + + +UNTYPED_SCALARS = [ + DictCase("utf8", pa.utf8(), ("a", "b", "c"), ("'a'", "'b'", "'c'")), + DictCase("large_utf8", pa.large_utf8(), ("a", "b", "c"), ("'a'", "'b'", "'c'")), + DictCase("int64", pa.int64(), (5, 7, 9), ("5", "7", "9")), + DictCase( + "date32", + pa.date32(), + (dt.date(2020, 1, 1), dt.date(2020, 6, 1), dt.date(2021, 1, 1)), + ("DATE '2020-01-01'", "DATE '2020-06-01'", "DATE '2021-01-01'"), + ), + DictCase("bool", pa.bool_(), (False, True, True), ("false", "true", "true")), + DictCase( + "timestamp_us", + pa.timestamp("us"), + (dt.datetime(2020, 1, 1), dt.datetime(2020, 6, 1), dt.datetime(2021, 1, 1)), + ("TIMESTAMP '2020-01-01'", "TIMESTAMP '2020-06-01'", "TIMESTAMP '2021-01-01'"), + ), +] + +FLOATS = [ + DictCase("float32", pa.float32(), (1.5, 2.5, float("nan")), ("1.5::FLOAT", "2.5::FLOAT", "'nan'::FLOAT")), + DictCase("float64", pa.float64(), (1.5, 2.5, float("nan")), ("1.5::DOUBLE", "2.5::DOUBLE", "'nan'::DOUBLE")), +] + +TYPED_SCALARS = [ + DictCase("binary", pa.binary(), (b"aa", b"bb", b"cc"), BLOB_LITERALS), + DictCase("large_binary", pa.large_binary(), (b"aa", b"bb", b"cc"), BLOB_LITERALS), + DictCase( + "fixed_size_binary", + pa.binary(4), + (b"aaaa", b"bbbb", b"cccc"), + ("'aaaa'::BLOB", "'bbbb'::BLOB", "'cccc'::BLOB"), + ), + DictCase("decimal128_30_2", pa.decimal128(30, 2), DECIMALS, _decimal_literals(30, 2)), + DictCase("decimal128_10_2", pa.decimal128(10, 2), DECIMALS, _decimal_literals(10, 2), plan="either"), + DictCase("timestamptz_utc", pa.timestamp("us", tz="UTC"), TSTZ_VALUES, TSTZ_LITERALS, session_zone="UTC"), + DictCase( + "timestamptz_new_york", + pa.timestamp("us", tz=NEW_YORK), + TSTZ_VALUES, + TSTZ_LITERALS, + session_zone=NEW_YORK, + ), +] + +# pyarrow cannot evaluate a comparison on these dictionaries. +NO_PYARROW_KERNEL = [ + DictCase("decimal32", pa.decimal32(5, 2), DECIMALS, _decimal_literals(5, 2)), + DictCase("decimal64", pa.decimal64(15, 2), DECIMALS, _decimal_literals(15, 2)), + DictCase("binary_view", pa.binary_view(), (b"aa", b"bb", b"cc"), BLOB_LITERALS), + DictCase("string_view", pa.string_view(), ("a", "b", "c"), ("'a'", "'b'", "'c'")), +] + +# The session zone differs from the column's zone. +ZONE_MISMATCH = [ + DictCase( + "timestamptz_utc_in_new_york", + pa.timestamp("us", tz="UTC"), + TSTZ_VALUES, + TSTZ_LITERALS, + session_zone=NEW_YORK, + ), + DictCase( + "timestamptz_new_york_in_utc", + pa.timestamp("us", tz=NEW_YORK), + TSTZ_VALUES, + TSTZ_LITERALS, + session_zone="UTC", + ), +] + +# pyarrow cannot reference a field inside a dictionary of structs. +STRUCT_DICTIONARIES = [ + DictCase("struct_of_int", STRUCT_TYPE, STRUCT_VALUES, ("1", "2", "3"), value_field="i"), + DictCase("struct_of_string", STRUCT_TYPE, STRUCT_VALUES, ("'a'", "'b'", "'c'"), value_field="s"), +] + +LIST_DICTIONARIES = [ + DictCase( + "list_of_int", + pa.list_(pa.int64()), + ([1], [1, 2], [3]), + ("[1]::BIGINT[]", "[1, 2]::BIGINT[]", "[3]::BIGINT[]"), + ), +] + +PREDICATES = { + "eq": "{col} = {mid}", + "ne": "{col} <> {mid}", + "lt": "{col} < {mid}", + "le": "{col} <= {mid}", + "gt": "{col} > {mid}", + "ge": "{col} >= {mid}", + "between": "{col} BETWEEN {low} AND {mid}", + "in": "{col} IN ({low}, {high})", + "is_null": "{col} IS NULL", + "eq_nan": "{col} = {high}", + "ne_nan": "{col} <> {high}", + "lt_nan": "{col} < {high}", + "ge_nan": "{col} >= {high}", +} +COMPARISONS = ["eq", "ne", "lt", "le", "gt", "ge", "between"] +# The float cases hold NaN as their high value, so these compare with a NaN constant. +NAN_COMPARISONS = ["eq_nan", "ne_nan", "lt_nan", "ge_nan"] +ALL_OPS = [*COMPARISONS, "in", "is_null"] +SHAPES = ["bare", "in_struct"] +NESTED_SHAPES = ["in_null_struct", "two_structs_deep"] +# The struct is NULL on rows whose dictionary value matches, so a filter ignoring the parent's validity returns them. +NULL_STRUCT_ROWS = [False, True, False, False, False, False, True] + +INVALIDATING = (duckdb.InternalException, duckdb.FatalException) + + +def _dictionary(case): + dictionary = pa.array([*case.values, None], case.value_type) + return pa.DictionaryArray.from_arrays(pa.array(INDICES, pa.int8()), dictionary) + + +def _register(con, case, column, column_name): + if case.session_zone: + con.execute(f"SET TimeZone = '{case.session_zone}'") + con.register("arrow_t", pa.table({"x": ROWS, column_name: column})) + con.execute("CREATE TABLE native_t AS SELECT * FROM arrow_t") + + +def _register_shape(con, case, shape): + dictionary = _dictionary(case) + if shape == "bare": + _register(con, case, dictionary, "d") + col = "d" + elif shape == "in_struct": + _register(con, case, pa.StructArray.from_arrays([dictionary], ["d"]), "s") + col = "s.d" + elif shape == "in_null_struct": + mask = pa.array(NULL_STRUCT_ROWS) + _register(con, case, pa.StructArray.from_arrays([dictionary], ["d"], mask=mask), "s") + col = "s.d" + else: + inner = pa.StructArray.from_arrays([dictionary], ["d"]) + _register(con, case, pa.StructArray.from_arrays([inner], ["t"]), "s") + col = "s.t.d" + return f"{col}.{case.value_field}" if case.value_field else col + + +def _predicate(case, op, col): + low, mid, high = case.literals + return PREDICATES[op].format(col=col, low=low, mid=mid, high=high) + + +def _rows(con, table, predicate): + return con.execute(f"SELECT x FROM {table} WHERE {predicate} ORDER BY x").fetchall() + + +def _arrow_query(predicate): + return f"SELECT x FROM arrow_t WHERE {predicate}" + + +def _arrow_rows(con, predicate): + try: + return _rows(con, "arrow_t", predicate) + except INVALIDATING as error: + # InternalException and FatalException derive from duckdb.Error, which some strict xfail markers expect. + pytest.fail(f"the Arrow query invalidated the database: {error}") + + +def assert_expected_outcome(con, predicate, outcome, error_text): + """Assert one of three outcomes for the Arrow query. + + "rows": the native rows. "error": an error that leaves the connection usable and whose message contains + error_text. "either": one of the two, for cases the planned engine pushdown gate change moves from error to rows. + """ + expected = _rows(con, "native_t", predicate) + failure = None + try: + got = _rows(con, "arrow_t", predicate) + except duckdb.Error as error: + if isinstance(error, INVALIDATING): + raise + failure = error + if failure is None: + assert outcome != "error", f"expected an error containing {error_text!r}, got rows {got}" + assert got == expected + return + assert outcome != "rows", f"expected the native rows, got {failure}" + assert error_text in str(failure) + assert con.execute("SELECT 42").fetchall() == [(42,)] + + +def _params(cases, ops, marks_for, shapes=SHAPES): + params = [] + for case in cases: + for shape in shapes: + for op in ops: + marks = marks_for(case, op) + params.append(pytest.param(case, shape, op, marks=marks, id=f"{case.id}-{shape}-{op}")) + return params + + +def _outcome_params(cases, ops, outcome_for, marks=()): + params = [] + for case in cases: + for shape in SHAPES: + for op in ops: + outcome, error_text = outcome_for(case, shape, op) + params.append( + pytest.param(case, shape, op, outcome, error_text, marks=list(marks), id=f"{case.id}-{shape}-{op}") + ) + return params + + +def _xfail(reason, raises): + return [pytest.mark.xfail(strict=True, raises=raises, reason=reason)] + + +def _no_marks(case, op): + return [] + + +def _case(cases, case_id): + return next(case for case in cases if case.id == case_id) + + +NESTING_CASES = [ + _case(TYPED_SCALARS, "binary"), + _case(UNTYPED_SCALARS, "utf8"), + _case(TYPED_SCALARS, "decimal128_30_2"), +] + +MATCHES_NATIVE = [ + *_params(UNTYPED_SCALARS, ["eq", "between", "in", "is_null"], _no_marks), + *_params(FLOATS, [*ALL_OPS, *NAN_COMPARISONS], _no_marks), + *_params(TYPED_SCALARS, ALL_OPS, _no_marks), + *_params(ZONE_MISMATCH, COMPARISONS, _no_marks), + *_params(NESTING_CASES, ["eq", "is_null"], _no_marks, shapes=NESTED_SHAPES), +] + +# pyarrow's message for a comparison and for an IN list on each dictionary. +NO_KERNEL_TEXT = { + "decimal32": ( + "has no kernel matching input types (decimal32(5, 2), decimal32(5, 2))", + "'is_in' has no kernel matching input types (decimal32(5, 2))", + ), + "decimal64": ( + "has no kernel matching input types (decimal64(15, 2), decimal64(15, 2))", + "'is_in' has no kernel matching input types (decimal64(15, 2))", + ), + "binary_view": ( + "Unsupported cast from dictionary 2.00::DECIMAL(30,2)") + assert duckdb_cursor.execute("SELECT 42").fetchall() == [(42,)] + + +FLOAT_CHUNK_PREDICATES = { + "eq_nan": "= 'nan'::{type}", + "ne_nan": "<> 'nan'::{type}", + "lt_nan": "< 'nan'::{type}", + "ge_nan": ">= 'nan'::{type}", + "gt": "> 1.0::{type}", +} + + +@pytest.mark.parametrize("op", list(FLOAT_CHUNK_PREDICATES)) +@pytest.mark.parametrize(("value_type", "sql_type"), [(pa.float32(), "FLOAT"), (pa.float64(), "DOUBLE")]) +def test_float_dictionary_chunks_with_different_dictionaries_match_native(duckdb_cursor, value_type, sql_type, op): + nan = float("nan") + first = pa.DictionaryArray.from_arrays( + pa.array([0, 1, 2, None, 0], pa.int8()), pa.array([1.5, nan, None], value_type) + ) + second = pa.DictionaryArray.from_arrays( + pa.array([3, 2, 1, 0, None], pa.int8()), pa.array([0.5, None, nan, 2.5], value_type) + ) + column = pa.chunked_array([first, second]) + duckdb_cursor.register("arrow_t", pa.table({"x": list(range(1, 11)), "d": column})) + duckdb_cursor.execute("CREATE TABLE native_t AS SELECT * FROM arrow_t") + predicate = f"d {FLOAT_CHUNK_PREDICATES[op].format(type=sql_type)}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +UNIT_MICROS = {"s": 1_000_000, "ms": 1_000, "us": 1} +EPOCH = dt.datetime(1970, 1, 1, tzinfo=UTC) +# Constants on the column's grid, and half a unit above and below it (the latter before 1970). +GRID_OFFSETS = {"on_grid": 0, "half_above": 1, "half_below": -1} +ZONED_COMPARISONS = {"eq": "=", "ne": "<>", "lt": "<", "le": "<=", "gt": ">", "ge": ">="} +EPOCH_LITERAL = "TIMESTAMPTZ '1970-01-01 00:00:00+00'" + + +def _micros_literal(micros): + return f"TIMESTAMPTZ '{EPOCH + dt.timedelta(microseconds=micros):%Y-%m-%d %H:%M:%S.%f}+00'" + + +def _register_zoned(con, unit, column_zone, session_zone, shape): + step = UNIT_MICROS[unit] + values = tuple(EPOCH + dt.timedelta(microseconds=k * step) for k in (-1, 0, 1)) + case = DictCase(f"{unit}_{column_zone}", pa.timestamp(unit, tz=column_zone), values, (), session_zone) + if shape == "plain": + _register(con, case, _dictionary(case).cast(case.value_type), "p") + return "p" + return _register_shape(con, case, shape) + + +@pytest.mark.parametrize("op", [*ZONED_COMPARISONS, "between"]) +@pytest.mark.parametrize("shape", [*SHAPES, "plain"]) +@pytest.mark.parametrize("offset", list(GRID_OFFSETS)) +@pytest.mark.parametrize("session_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("unit", ["s", "ms"]) +def test_zoned_timestamp_comparison_matches_native(duckdb_cursor, unit, column_zone, session_zone, offset, shape, op): + step = UNIT_MICROS[unit] + col = _register_zoned(duckdb_cursor, unit, column_zone, session_zone, shape) + constant = GRID_OFFSETS[offset] * step // 2 + if op == "between": + predicate = f"{col} BETWEEN {_micros_literal(constant)} AND {_micros_literal(constant + step)}" + else: + predicate = f"{col} {ZONED_COMPARISONS[op]} {_micros_literal(constant)}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +INFINITE_BOUNDS = { + "minus_infinity_to_infinity": ("'-infinity'::TIMESTAMPTZ", "'infinity'::TIMESTAMPTZ"), + "minus_infinity_to_epoch": ("'-infinity'::TIMESTAMPTZ", EPOCH_LITERAL), + "epoch_to_infinity": (EPOCH_LITERAL, "'infinity'::TIMESTAMPTZ"), +} + + +@pytest.mark.parametrize("bounds", list(INFINITE_BOUNDS)) +@pytest.mark.parametrize("shape", [*SHAPES, "plain"]) +@pytest.mark.parametrize("session_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("unit", list(UNIT_MICROS)) +def test_zoned_timestamp_between_infinities_matches_native( + duckdb_cursor, unit, column_zone, session_zone, shape, bounds +): + col = _register_zoned(duckdb_cursor, unit, column_zone, session_zone, shape) + low, high = INFINITE_BOUNDS[bounds] + predicate = f"{col} BETWEEN {low} AND {high}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +INFINITE_IN_LISTS = { + "infinities": "'-infinity'::TIMESTAMPTZ, 'infinity'::TIMESTAMPTZ", + "infinity_and_epoch": f"'infinity'::TIMESTAMPTZ, {EPOCH_LITERAL}", +} +INFINITE_IN_REASON = ( + "client IsIn: PythonObject::FromValue turns an infinite TIMESTAMPTZ into a zoneless datetime, " + "which pyarrow refuses for a zoned column" +) +INFINITE_IN_CASES = [ + pytest.param( + unit, + column_zone, + session_zone, + shape, + values, + # The engine keeps an IN list on a struct field above the scan, so only bare and plain columns reach pyarrow. + marks=[] if shape == "in_struct" else _xfail(INFINITE_IN_REASON, duckdb.Error), + id=f"{unit}-{column_zone}-{session_zone}-{shape}-{values}", + ) + for unit in UNIT_MICROS + for column_zone in ["UTC", NEW_YORK] + for session_zone in ["UTC", NEW_YORK] + for shape in [*SHAPES, "plain"] + for values in INFINITE_IN_LISTS +] + + +@pytest.mark.parametrize(("unit", "column_zone", "session_zone", "shape", "values"), INFINITE_IN_CASES) +def test_zoned_timestamp_in_with_infinity_matches_native(duckdb_cursor, unit, column_zone, session_zone, shape, values): + col = _register_zoned(duckdb_cursor, unit, column_zone, session_zone, shape) + predicate = f"{col} IN ({INFINITE_IN_LISTS[values]})" + assert _arrow_rows(duckdb_cursor, predicate) == _rows(duckdb_cursor, "native_t", predicate) + + +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("session_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("unit", ["s", "ms"]) +def test_zoned_timestamp_in_with_off_grid_constants_matches_native( + duckdb_cursor, unit, session_zone, column_zone, shape +): + col = _register_zoned(duckdb_cursor, unit, column_zone, session_zone, "bare" if shape == "dictionary" else "plain") + step = UNIT_MICROS[unit] + predicate = f"{col} IN ({_micros_literal(step // 2)}, {_micros_literal(step)}, {_micros_literal(-step // 2)})" + assert _arrow_rows(duckdb_cursor, predicate) == _rows(duckdb_cursor, "native_t", predicate) + + +@pytest.mark.parametrize("op", ["eq", "gt"]) +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +def test_zoned_timestamp_two_structs_deep_matches_native(duckdb_cursor, shape, op): + step = UNIT_MICROS["s"] + values = tuple(EPOCH + dt.timedelta(microseconds=k * step) for k in (-1, 0, 1)) + case = DictCase("s_utc", pa.timestamp("s", tz="UTC"), values, (), NEW_YORK) + leaf = _dictionary(case) if shape == "dictionary" else _dictionary(case).cast(case.value_type) + inner = pa.StructArray.from_arrays([leaf], ["d"]) + _register(duckdb_cursor, case, pa.StructArray.from_arrays([inner], ["t"]), "s") + predicate = f"s.t.d {ZONED_COMPARISONS[op]} {_micros_literal(step // 2)}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +@pytest.mark.parametrize( + ("op", "constant", "expected"), + [ + pytest.param("gt", "'1970-01-01 00:00:00.0005+00'", [(3,)], id="gt"), + pytest.param("eq", "'1970-01-01 00:00:00+00'", [(2,)], id="eq"), + ], +) +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +def test_zoned_timestamp_beside_decoy_siblings_matches_native(duckdb_cursor, shape, op, constant, expected): + # A decoy before the real child at each struct level makes a wrong sibling index read a UTC microsecond column. + leaf = pa.array([-1, 0, 1], pa.int64()).cast(pa.timestamp("ms", tz="+05:45")) + if shape == "dictionary": + leaf = pa.DictionaryArray.from_arrays(pa.array([0, 1, 2], pa.int8()), leaf) + decoy = pa.array([10**6, 2 * 10**6, 3 * 10**6], pa.int64()).cast(pa.timestamp("us", tz="UTC")) + inner = pa.StructArray.from_arrays([decoy, leaf], ["decoy", "d"]) + duckdb_cursor.execute(f"SET TimeZone = '{NEW_YORK}'") + duckdb_cursor.register( + "arrow_t", pa.table({"x": [1, 2, 3], "s": pa.StructArray.from_arrays([decoy, inner], ["decoy", "t"])}) + ) + duckdb_cursor.execute("CREATE TABLE native_t AS SELECT * FROM arrow_t") + predicate = f"s.t.d {ZONED_COMPARISONS[op]} TIMESTAMPTZ {constant}" + assert _rows(duckdb_cursor, "native_t", predicate) == expected + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +@pytest.mark.parametrize("op", ["eq", "lt", "between"]) +@pytest.mark.parametrize("offset", ["on_grid", "half_above"]) +@pytest.mark.parametrize("shape", ["bare", "plain"]) +@pytest.mark.parametrize("column_zone", ["+05:00", "-03:30"]) +@pytest.mark.parametrize("unit", ["s", "ms"]) +def test_fixed_offset_zone_comparison_matches_native(duckdb_cursor, unit, column_zone, shape, offset, op): + step = UNIT_MICROS[unit] + col = _register_zoned(duckdb_cursor, unit, column_zone, NEW_YORK, shape) + constant = GRID_OFFSETS[offset] * step // 2 + if op == "between": + predicate = f"{col} BETWEEN {_micros_literal(constant)} AND {_micros_literal(constant + step)}" + else: + predicate = f"{col} {ZONED_COMPARISONS[op]} {_micros_literal(constant)}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +def _zoned_source(unit, column_zone, shape, source): + step = UNIT_MICROS[unit] + values = tuple(EPOCH + dt.timedelta(microseconds=k * step) for k in (-1, 0, 1)) + case = DictCase(f"{unit}_{column_zone}", pa.timestamp(unit, tz=column_zone), values, ()) + column = _dictionary(case) if shape == "dictionary" else _dictionary(case).cast(case.value_type) + table = pa.table({"x": ROWS, "p": column}) + if source == "scanner": + return table, pytest.importorskip("pyarrow.dataset").dataset(table).scanner() + return table, table.to_reader() + + +@pytest.mark.parametrize("op", ["eq", "gt"]) +@pytest.mark.parametrize("offset", ["on_grid", "half_above"]) +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +@pytest.mark.parametrize("session_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("unit", ["s", "ms"]) +@pytest.mark.parametrize("source", ["scanner", "record_batch_reader"]) +def test_zoned_timestamp_scanner_and_reader_sources_match_native( + duckdb_cursor, source, unit, column_zone, session_zone, shape, offset, op +): + step = UNIT_MICROS[unit] + predicate = f"p {ZONED_COMPARISONS[op]} {_micros_literal(GRID_OFFSETS[offset] * step // 2)}" + duckdb_cursor.execute(f"SET TimeZone = '{session_zone}'") + table, arrow_source = _zoned_source(unit, column_zone, shape, source) + duckdb_cursor.register("native_source", table) + duckdb_cursor.execute("CREATE TABLE native_t AS SELECT * FROM native_source") + duckdb_cursor.register("arrow_t", arrow_source) + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + # A record batch reader can be scanned once, so the plan check gets a fresh source. + duckdb_cursor.register("arrow_t", _zoned_source(unit, column_zone, shape, source)[1]) + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +class _SchemaExportOnly: + """A schema object that supports only the Arrow C export, as some dataset wrappers return.""" + + def __init__(self, schema) -> None: + self._schema = schema + + def _export_to_c(self, *args): + return self._schema._export_to_c(*args) + + +@pytest.mark.parametrize("op", ["eq", "gt"]) +@pytest.mark.parametrize("offset", ["on_grid", "half_above"]) +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +def test_dataset_with_an_export_only_schema_matches_native(duckdb_cursor, shape, offset, op): + pa_ds = pytest.importorskip("pyarrow.dataset") + + class ExportOnlySchemaDataset(pa_ds.InMemoryDataset): + @property + def schema(self): + return _SchemaExportOnly(super().schema) + + column = pa.array([-1, 0, 1, None], pa.timestamp("ms", tz="UTC")) + if shape == "dictionary": + column = pa.DictionaryArray.from_arrays(pa.array([0, 1, 2, None], pa.int8()), column[:3]) + table = pa.table({"x": [1, 2, 3, 4], "p": column}) + duckdb_cursor.execute(f"SET TimeZone = '{NEW_YORK}'") + duckdb_cursor.register("native_source", table) + duckdb_cursor.execute("CREATE TABLE native_t AS SELECT * FROM native_source") + duckdb_cursor.register("arrow_t", ExportOnlySchemaDataset(table)) + predicate = f"p {ZONED_COMPARISONS[op]} {_micros_literal(GRID_OFFSETS[offset] * UNIT_MICROS['ms'] // 2)}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +# The first stored values whose microsecond form overflows int64; one unit less still fits. +FIRST_OVERFLOWING = {"s": 9_223_372_036_855, "ms": 9_223_372_036_854_776} + + +def _register_raw_zoned(con, unit, raw_values, shape): + con.execute("SET TimeZone = 'UTC'") + column = pa.array([*raw_values, None], pa.int64()).cast(pa.timestamp(unit, tz="UTC")) + if shape == "dictionary": + column = pa.DictionaryArray.from_arrays(pa.array(range(len(column)), pa.int8()), column) + con.register("arrow_t", pa.table({"x": list(range(1, len(column) + 1)), "p": column})) + + +@pytest.mark.parametrize("op", list(ZONED_COMPARISONS)) +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +@pytest.mark.parametrize("unit", list(FIRST_OVERFLOWING)) +def test_zoned_timestamp_at_the_microsecond_limit_matches_hand_rows(duckdb_cursor, unit, shape, op): + _register_raw_zoned(duckdb_cursor, unit, [FIRST_OVERFLOWING[unit] - 1], shape) + expected = [(1,)] if op in ("ne", "gt", "ge") else [] + assert _arrow_rows(duckdb_cursor, f"p {ZONED_COMPARISONS[op]} {EPOCH_LITERAL}") == expected + + +@pytest.mark.parametrize("op", list(ZONED_COMPARISONS)) +@pytest.mark.parametrize("shape", ["plain", "dictionary"]) +@pytest.mark.parametrize("unit", list(FIRST_OVERFLOWING)) +def test_zoned_timestamp_beyond_the_microsecond_limit(duckdb_cursor, unit, shape, op): + first = FIRST_OVERFLOWING[unit] + _register_raw_zoned(duckdb_cursor, unit, [first - 1, first], shape) + predicate = f"p {ZONED_COMPARISONS[op]} {EPOCH_LITERAL}" + if op in ("ne", "gt", "ge"): + # pyarrow keeps the overflowing row, and the scan cannot convert it to a DuckDB TIMESTAMPTZ. + with pytest.raises(duckdb.ConversionException, match="Could not convert TimestampTZ to Microsecond"): + _rows(duckdb_cursor, "arrow_t", predicate) + assert duckdb_cursor.execute("SELECT 42").fetchall() == [(42,)] + else: + assert _arrow_rows(duckdb_cursor, predicate) == [] + + +@pytest.mark.parametrize("op", ["gt", "eq", "between"]) +@pytest.mark.parametrize("offset", ["on_grid", "half_above"]) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +@pytest.mark.parametrize("unit", ["s", "ms"]) +def test_zoned_timestamp_parquet_dataset_matches_native(duckdb_cursor, tmp_path, unit, column_zone, offset, op): + pq = pytest.importorskip("pyarrow.parquet") + pa_ds = pytest.importorskip("pyarrow.dataset") + step = UNIT_MICROS[unit] + column = pa.array([*range(20), None], pa.int64()).cast(pa.timestamp(unit, tz=column_zone)) + path = tmp_path / "zoned.parquet" + pq.write_table(pa.table({"x": list(range(1, 22)), "p": column}), path, row_group_size=4) + duckdb_cursor.execute(f"SET TimeZone = '{NEW_YORK}'") + duckdb_cursor.register("arrow_t", pa_ds.dataset(path)) + duckdb_cursor.execute("CREATE TABLE native_t AS SELECT * FROM arrow_t") + constant = 10 * step + GRID_OFFSETS[offset] * step // 2 + if op == "between": + predicate = f"p BETWEEN {_micros_literal(constant)} AND {_micros_literal(constant + 3 * step)}" + else: + predicate = f"p {ZONED_COMPARISONS[op]} {_micros_literal(constant)}" + expected = _rows(duckdb_cursor, "native_t", predicate) + assert _arrow_rows(duckdb_cursor, predicate) == expected + assert_fully_pushed(duckdb_cursor, _arrow_query(predicate), len(expected)) + + +OFF_GRID_PRUNING = "per unit rounding: an off grid constant prunes only once it is rounded into the column's unit" +PRUNING_CONSTANTS_MS = {"on_grid": 799, "off_grid": 799.5} + + +def _register_capturing_parquet(con, tmp_path, column_zone): + pq = pytest.importorskip("pyarrow.parquet") + pa_ds = pytest.importorskip("pyarrow.dataset") + captured = [] + + class FilterCapturingDataset(pa_ds.FileSystemDataset): + def scanner(self, **kwargs): + captured.append(kwargs.get("filter")) + return super().scanner(**kwargs) + + column = pa.array(range(1000), pa.int64()).cast(pa.timestamp("ms", tz=column_zone)) + path = tmp_path / "pruning.parquet" + pq.write_table(pa.table({"x": list(range(1000)), "p": column}), path, row_group_size=100) + source = pa_ds.dataset(path) + dataset = FilterCapturingDataset(list(source.get_fragments()), source.schema, source.format, source.filesystem) + con.execute(f"SET TimeZone = '{NEW_YORK}'") + con.register("arrow_t", dataset) + return next(iter(source.get_fragments())), captured + + +def _pruning_count(con, constant_ms): + predicate = f"p > {_micros_literal(int(constant_ms * 1000))}" + return con.execute(f"SELECT count(*) FROM arrow_t WHERE {predicate}").fetchone() + + +@pytest.mark.parametrize("constant", list(PRUNING_CONSTANTS_MS)) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +def test_zoned_timestamp_parquet_pruning_filter_matches_hand_rows(duckdb_cursor, tmp_path, column_zone, constant): + _register_capturing_parquet(duckdb_cursor, tmp_path, column_zone) + assert _pruning_count(duckdb_cursor, PRUNING_CONSTANTS_MS[constant]) == (200,) + + +@pytest.mark.parametrize( + "constant", + [ + pytest.param("on_grid"), + pytest.param("off_grid", marks=_xfail(OFF_GRID_PRUNING, AssertionError)), + ], +) +@pytest.mark.parametrize("column_zone", ["UTC", NEW_YORK]) +def test_zoned_timestamp_parquet_row_group_pruning(duckdb_cursor, tmp_path, column_zone, constant): + fragment, captured = _register_capturing_parquet(duckdb_cursor, tmp_path, column_zone) + _pruning_count(duckdb_cursor, PRUNING_CONSTANTS_MS[constant]) + assert fragment.subset(captured[-1]).num_row_groups == 2 diff --git a/tests/fast/arrow/test_nan_null_filter_pushdown.py b/tests/fast/arrow/test_nan_null_filter_pushdown.py index 914dfe76..a1040c7f 100644 --- a/tests/fast/arrow/test_nan_null_filter_pushdown.py +++ b/tests/fast/arrow/test_nan_null_filter_pushdown.py @@ -345,15 +345,15 @@ def test_le_nan_float_col_double_const(self, factory): ("coltype", "op", "factory"), _make_combined_params_coltype_op( [ - ("DOUBLE", "<", True), - ("DOUBLE", "=", True), - ("DOUBLE", "<>", True), - ("DOUBLE", ">=", True), + ("DOUBLE", "<", False), + ("DOUBLE", "=", False), + ("DOUBLE", "<>", False), + ("DOUBLE", ">=", False), ("DOUBLE", ">", False), - ("FLOAT", "<", True), - ("FLOAT", "=", True), - ("FLOAT", "<>", True), - ("FLOAT", ">=", True), + ("FLOAT", "<", False), + ("FLOAT", "=", False), + ("FLOAT", "<>", False), + ("FLOAT", ">=", False), ("FLOAT", ">", False), ] ), @@ -408,13 +408,13 @@ def test_le_nan_or_eq(self, coltype, factory): self._make_tables(coltype, VALUES_WITH_NULL, factory) assert_matches_native(self.con, f"b <= 'nan'::{coltype} OR b = 1.0") - @pytest.mark.parametrize("factory", _make_factory_params_with_xfail(xfail_on_dict=True)) + @pytest.mark.parametrize("factory", _make_factory_params_with_xfail(xfail_on_dict=False)) @pytest.mark.parametrize("coltype", ["DOUBLE", "FLOAT"]) def test_le_nan_and_gt(self, coltype, factory): self._make_tables(coltype, VALUES_WITH_NULL, factory) assert_matches_native(self.con, f"b <= 'nan'::{coltype} AND b > 0") - @pytest.mark.parametrize("factory", _make_factory_params_with_xfail(xfail_on_dict=True)) + @pytest.mark.parametrize("factory", _make_factory_params_with_xfail(xfail_on_dict=False)) @pytest.mark.parametrize("coltype", ["DOUBLE", "FLOAT"]) def test_between_zero_and_nan(self, coltype, factory): self._make_tables(coltype, VALUES_WITH_NULL, factory) @@ -485,38 +485,38 @@ def test_in_list_only_nan_and_zero(self, coltype, factory): ("DOUBLE", "<=", "0::DOUBLE", False), ("DOUBLE", "=", "0::DOUBLE", False), ("DOUBLE", "<>", "0::DOUBLE", False), - ("DOUBLE", ">=", "0::DOUBLE", True), - ("DOUBLE", ">", "0::DOUBLE", True), + ("DOUBLE", ">=", "0::DOUBLE", False), + ("DOUBLE", ">", "0::DOUBLE", False), ("DOUBLE", "<", "'inf'::DOUBLE", False), ("DOUBLE", "<=", "'inf'::DOUBLE", False), ("DOUBLE", "=", "'inf'::DOUBLE", False), ("DOUBLE", "<>", "'inf'::DOUBLE", False), - ("DOUBLE", ">=", "'inf'::DOUBLE", True), - ("DOUBLE", ">", "'inf'::DOUBLE", True), + ("DOUBLE", ">=", "'inf'::DOUBLE", False), + ("DOUBLE", ">", "'inf'::DOUBLE", False), ("DOUBLE", "<", "'-inf'::DOUBLE", False), ("DOUBLE", "<=", "'-inf'::DOUBLE", False), ("DOUBLE", "=", "'-inf'::DOUBLE", False), ("DOUBLE", "<>", "'-inf'::DOUBLE", False), - ("DOUBLE", ">=", "'-inf'::DOUBLE", True), - ("DOUBLE", ">", "'-inf'::DOUBLE", True), + ("DOUBLE", ">=", "'-inf'::DOUBLE", False), + ("DOUBLE", ">", "'-inf'::DOUBLE", False), ("FLOAT", "<", "0::FLOAT", False), ("FLOAT", "<=", "0::FLOAT", False), ("FLOAT", "=", "0::FLOAT", False), ("FLOAT", "<>", "0::FLOAT", False), - ("FLOAT", ">=", "0::FLOAT", True), - ("FLOAT", ">", "0::FLOAT", True), + ("FLOAT", ">=", "0::FLOAT", False), + ("FLOAT", ">", "0::FLOAT", False), ("FLOAT", "<", "'inf'::FLOAT", False), ("FLOAT", "<=", "'inf'::FLOAT", False), ("FLOAT", "=", "'inf'::FLOAT", False), ("FLOAT", "<>", "'inf'::FLOAT", False), - ("FLOAT", ">=", "'inf'::FLOAT", True), - ("FLOAT", ">", "'inf'::FLOAT", True), + ("FLOAT", ">=", "'inf'::FLOAT", False), + ("FLOAT", ">", "'inf'::FLOAT", False), ("FLOAT", "<", "'-inf'::FLOAT", False), ("FLOAT", "<=", "'-inf'::FLOAT", False), ("FLOAT", "=", "'-inf'::FLOAT", False), ("FLOAT", "<>", "'-inf'::FLOAT", False), - ("FLOAT", ">=", "'-inf'::FLOAT", True), - ("FLOAT", ">", "'-inf'::FLOAT", True), + ("FLOAT", ">=", "'-inf'::FLOAT", False), + ("FLOAT", ">", "'-inf'::FLOAT", False), ] ), )