Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/arrow/arrow_array_stream.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
64 changes: 45 additions & 19 deletions src/arrow/filter_pushdown_visitor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<Identifier> path;
const ArrowType *leaf_type;
const ArrowSchema *leaf_schema;
};

ResolvedColumn ResolveColumn(const Expression &expr, const vector<Identifier> &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<Identifier> &path) {
if (type && type->RunEndEncoded()) {
vector<string> 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<Identifier> &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",
Expand All @@ -47,20 +66,24 @@ ResolvedColumn ResolveColumn(const Expression &expr, const vector<Identifier> &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<ArrowStructInfo>().GetChild(child_idx);
inner.leaf_type = ValueType(&inner.leaf_type->GetTypeInfo<ArrowStructInfo>().GetChild(child_idx), inner.path);
}
if (inner.leaf_schema) {
D_ASSERT(child_idx < NumericCast<idx_t>(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
Expand All @@ -85,7 +108,8 @@ nb::object EmitCompare(FilterBackend &backend, ExpressionType op, nb::object col
} // anonymous namespace

nb::object TransformExpression(const Expression &expression, const vector<Identifier> &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();

Expand All @@ -109,9 +133,9 @@ nb::object TransformExpression(const Expression &expression, const vector<Identi
throw NotImplementedException("Can only push down constant comparisons.");
}

auto resolved = ResolveColumn(*column_side, column_path, arrow_type);
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.leaf_type,
return EmitCompare(backend, expression_type, std::move(col), constant_side->GetValue(), resolved,
timezone_config);
}

Expand Down Expand Up @@ -143,7 +167,7 @@ nb::object TransformExpression(const Expression &expression, const vector<Identi
return nb::none();
}
try {
return TransformExpression(*child, column_path, backend, arrow_type, timezone_config);
return TransformExpression(*child, column_path, backend, arrow_type, arrow_schema, timezone_config);
} catch (const NotImplementedException &) {
return nb::none();
}
Expand All @@ -160,17 +184,17 @@ nb::object TransformExpression(const Expression &expression, const vector<Identi
if (expression_class == ExpressionClass::BOUND_OPERATOR) {
auto &op_expr = expression.Cast<BoundOperatorExpression>();
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<Value> values;
for (idx_t i = 1; i < op_expr.GetChildren().size(); i++) {
Expand All @@ -188,8 +212,8 @@ nb::object TransformExpression(const Expression &expression, const vector<Identi
auto &conj_expr = expression.Cast<BoundConjunctionExpression>();
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
Expand Down Expand Up @@ -219,11 +243,13 @@ nb::object TransformExpression(const Expression &expression, const vector<Identi
}

nb::object TransformFilter(const TableFilter &filter, const vector<Identifier> &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<ExpressionFilter>();
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",
Expand Down
8 changes: 5 additions & 3 deletions src/arrow/polars_filter_pushdown.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -137,9 +139,9 @@ nb::object PolarsFilterPushdown::TransformFilter(const TableFilterSet &filter_co

vector<Identifier> 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;
}
Expand Down
82 changes: 31 additions & 51 deletions src/arrow/pyarrow_filter_pushdown.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(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");
Expand Down Expand Up @@ -101,16 +79,11 @@ nb::object MakePyArrowScalar(const Value &constant, const string &timezone_confi
return dataset_scalar(scalar(constant.GetValue<int64_t>(), 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<ArrowDateTimeInfo>();
auto base_value = constant.GetValue<int64_t>();
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<int64_t>(), date_type("us", nb::arg("tz") = zone)));
}
case LogicalTypeId::TIMESTAMP_TZ_NS: {
nb::handle date_type = import_cache.pyarrow.timestamp();
Expand Down Expand Up @@ -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 {
Expand All @@ -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);
Expand All @@ -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 {
Expand Down Expand Up @@ -330,7 +306,8 @@ struct PyArrowBackend : public FilterBackend {
nb::object PyArrowFilterPushdown::TransformFilter(TableFilterSet &filter_collection,
unordered_map<idx_t, string> &columns,
unordered_map<idx_t, idx_t> 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) {
Expand All @@ -339,9 +316,12 @@ nb::object PyArrowFilterPushdown::TransformFilter(TableFilterSet &filter_collect
D_ASSERT(columns.find(column_idx) != columns.end());

vector<Identifier> 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<idx_t>(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;
}
Expand Down
Loading
Loading