|
1 | 1 | from __future__ import annotations # noqa: D100 |
2 | 2 |
|
3 | | -import contextlib |
4 | 3 | import datetime |
| 4 | +import io |
5 | 5 | import json |
6 | 6 | import typing |
7 | 7 | from decimal import Decimal |
@@ -35,16 +35,95 @@ def _predicate_to_expression(predicate: pl.Expr) -> duckdb.Expression | None: |
35 | 35 | """ |
36 | 36 | # Serialize the Polars expression tree to JSON |
37 | 37 | tree = json.loads(predicate.meta.serialize(format="json")) |
| 38 | + return _tree_to_sql_expression(tree) |
38 | 39 |
|
| 40 | + |
| 41 | +def _tree_to_sql_expression(tree: _ExpressionTree) -> duckdb.Expression | None: |
| 42 | + """Convert an already-parsed Polars expression tree to a DuckDB expression. |
| 43 | +
|
| 44 | + Returns None if the tree contains a node we cannot translate to SQL. |
| 45 | + """ |
39 | 46 | try: |
40 | | - # Convert the tree to SQL |
41 | | - sql_filter = _pl_tree_to_sql(tree) |
42 | | - return duckdb.SQLExpression(sql_filter) |
| 47 | + return duckdb.SQLExpression(_pl_tree_to_sql(tree)) |
43 | 48 | except Exception: |
44 | 49 | # If the conversion fails, we return None |
45 | 50 | return None |
46 | 51 |
|
47 | 52 |
|
| 53 | +# Polars "dynamic predicates" |
| 54 | +# --------------------------- |
| 55 | +# When a slice / TOP-N sits above a scan, polars' optimizer pushes a *dynamic |
| 56 | +# predicate* into the scan, AND-ed onto any real filter. It is an internal |
| 57 | +# optimizer node, not a materializable expression: it serializes as a `Display` |
| 58 | +# node with `fmt_str` "dynamic_pred: <uuid>", and feeding it back into |
| 59 | +# `DataFrame.filter()` panics polars (`unreachable!` in `expr_to_ir`). It is only |
| 60 | +# an early-pruning hint -- the limit above the scan still runs -- so we drop it |
| 61 | +# and keep the real predicate, which we MUST still apply (polars trusts the |
| 62 | +# source and does not re-filter above it). Reported as duckdb-python#460; for |
| 63 | +# polars-side context see: |
| 64 | +# - https://github.com/pola-rs/polars/issues/21665 (why real + dynamic arrive AND-ed) |
| 65 | +# - https://github.com/pola-rs/polars/issues/22252 (filter() panicking on un-lowerable nodes) |
| 66 | +def _is_dynamic_predicate_node(node: typing.Any) -> bool: # noqa: ANN401 |
| 67 | + """Return True if a serialized node is a dynamic predicate (see the note above). |
| 68 | +
|
| 69 | + Detected by shape: a ``Display`` node whose ``fmt_str`` starts with ``dynamic_pred``. |
| 70 | + """ |
| 71 | + if not isinstance(node, dict): |
| 72 | + return False |
| 73 | + display = node.get("Display") |
| 74 | + return ( |
| 75 | + isinstance(display, dict) |
| 76 | + and isinstance(display.get("fmt_str"), str) |
| 77 | + and display["fmt_str"].startswith("dynamic_pred") |
| 78 | + ) |
| 79 | + |
| 80 | + |
| 81 | +def _tree_contains_dynamic_predicate(node: typing.Any) -> bool: # noqa: ANN401 |
| 82 | + """Return True if the serialized expression tree contains a dynamic predicate anywhere.""" |
| 83 | + if _is_dynamic_predicate_node(node): |
| 84 | + return True |
| 85 | + if isinstance(node, dict): |
| 86 | + return any(_tree_contains_dynamic_predicate(child) for child in node.values()) |
| 87 | + if isinstance(node, list): |
| 88 | + return any(_tree_contains_dynamic_predicate(child) for child in node) |
| 89 | + return False |
| 90 | + |
| 91 | + |
| 92 | +def _strip_dynamic_predicates(tree: typing.Any) -> tuple[_ExpressionTree | None, bool]: # noqa: ANN401 |
| 93 | + """Remove dynamic-predicate conjuncts from a serialized predicate tree. |
| 94 | +
|
| 95 | + See the note above for what a dynamic predicate is and why we drop it. |
| 96 | +
|
| 97 | + Returns ``(stripped_tree, removed)``. ``stripped_tree`` is ``None`` when the |
| 98 | + predicate was purely dynamic. Raises ``NotImplementedError`` if a dynamic |
| 99 | + predicate appears anywhere other than a top-level ``And`` conjunct — a shape |
| 100 | + polars does not produce today, where the hint can neither be safely dropped |
| 101 | + nor applied. |
| 102 | + """ |
| 103 | + if _is_dynamic_predicate_node(tree): |
| 104 | + return None, True |
| 105 | + if isinstance(tree, dict) and "BinaryExpr" in tree: |
| 106 | + bin_expr = tree["BinaryExpr"] |
| 107 | + if isinstance(bin_expr, dict) and bin_expr.get("op") == "And": |
| 108 | + left, left_removed = _strip_dynamic_predicates(bin_expr["left"]) |
| 109 | + right, right_removed = _strip_dynamic_predicates(bin_expr["right"]) |
| 110 | + removed = left_removed or right_removed |
| 111 | + if left is None: |
| 112 | + return right, removed |
| 113 | + if right is None: |
| 114 | + return left, removed |
| 115 | + return {"BinaryExpr": {**bin_expr, "left": left, "right": right}}, removed |
| 116 | + if _tree_contains_dynamic_predicate(tree): |
| 117 | + msg = "Cannot handle a polars dynamic predicate outside a top-level AND conjunct" |
| 118 | + raise NotImplementedError(msg) |
| 119 | + return tree, False |
| 120 | + |
| 121 | + |
| 122 | +def _expression_from_tree(tree: _ExpressionTree) -> pl.Expr: |
| 123 | + """Rebuild a polars expression from a serialized tree (inverse of meta.serialize).""" |
| 124 | + return pl.Expr.deserialize(io.BytesIO(json.dumps(tree).encode()), format="json") |
| 125 | + |
| 126 | + |
48 | 127 | def _pl_operation_to_sql(op: str) -> str: |
49 | 128 | """Map Polars binary operation strings to SQL equivalents. |
50 | 129 |
|
@@ -286,26 +365,36 @@ def source_generator( |
286 | 365 | batch_size: int | None, |
287 | 366 | ) -> Iterator[pl.DataFrame]: |
288 | 367 | duck_predicate = None |
| 368 | + fallback_predicate = None |
289 | 369 | relation_final = relation |
290 | 370 | if with_columns is not None: |
291 | 371 | cols = ",".join(map(_escape_sql_identifier, with_columns)) |
292 | 372 | relation_final = relation_final.project(cols) |
293 | 373 | if n_rows is not None: |
294 | 374 | relation_final = relation_final.limit(n_rows) |
295 | 375 | if predicate is not None: |
296 | | - # We have a predicate, if possible, we push it down to DuckDB |
297 | | - with contextlib.suppress(AssertionError, KeyError): |
298 | | - duck_predicate = _predicate_to_expression(predicate) |
| 376 | + # Strip any dynamic-predicate hint (see the dynamic-predicate note |
| 377 | + # above); the real predicate must still be applied. |
| 378 | + tree = json.loads(predicate.meta.serialize(format="json")) |
| 379 | + real_tree, had_dynamic = _strip_dynamic_predicates(tree) |
| 380 | + if real_tree is not None: |
| 381 | + # We have a real predicate; if possible, push it down to DuckDB. |
| 382 | + duck_predicate = _tree_to_sql_expression(real_tree) |
| 383 | + if duck_predicate is None: |
| 384 | + # Could not push it down: re-apply it polars-side. Rebuild the |
| 385 | + # expression from the stripped tree so we never hand polars the |
| 386 | + # dynamic node it cannot lower. |
| 387 | + fallback_predicate = _expression_from_tree(real_tree) if had_dynamic else predicate |
299 | 388 | # Try to pushdown filter, if one exists |
300 | 389 | if duck_predicate is not None: |
301 | 390 | relation_final = relation_final.filter(duck_predicate) |
302 | 391 | results = relation_final.to_arrow_reader() if batch_size is None else relation_final.to_arrow_reader(batch_size) |
303 | 392 |
|
304 | 393 | for record_batch in iter(results.read_next_batch, None): |
305 | | - if predicate is not None and duck_predicate is None: |
| 394 | + if fallback_predicate is not None: |
306 | 395 | # We have a predicate, but did not manage to push it down, we fallback here |
307 | | - yield pl.from_arrow(record_batch).filter(predicate) # type: ignore[arg-type,misc,unused-ignore] |
| 396 | + yield pl.from_arrow(record_batch).filter(fallback_predicate) # type: ignore[arg-type,misc,unused-ignore] |
308 | 397 | else: |
309 | 398 | yield pl.from_arrow(record_batch) # type: ignore[misc,unused-ignore] |
310 | 399 |
|
311 | | - return register_io_source(source_generator, schema=schema) |
| 400 | + return register_io_source(source_generator, schema=schema, is_pure=True) |
0 commit comments