cool9850311 commented on issue #27127:
URL: https://github.com/apache/datafusion/issues/27127#issuecomment-6657864852
I dug into the root cause and found the scope is wider than my original
report, plus one imprecision in it that I want to correct. Summary: the
proximate mechanism I described (a `RepartitionExec` blocking the projection)
is right, but the **root cause is one node type missing from the
leaf-expression extraction allowlist**, and it affects two more functions.
## Correction to my repro: plain `LIMIT n` masks the bug
My first test matrix put `LIMIT 3` on every query and almost everything
passed, which nearly led me to the wrong conclusion that only the `ORDER BY`
shape breaks. A bare `LIMIT n` is pushed into the scan, so the row estimate
drops to `n`, falls below `batch_size`, and no `RepartitionExec` is inserted:
| query | result |
|---|---|
| `SELECT file_row_index() + 1 FROM big LIMIT 3` | **ok** — limit pushed
into scan, bug masked |
| `SELECT file_row_index() + 1 FROM big` | **error** |
| `SELECT file_row_index() + 1, column1 FROM big ORDER BY column1 LIMIT 3` |
**error** |
| `SELECT file_row_index() + 1 FROM big OFFSET 5 LIMIT 3` | **error** |
So the **minimal repro is simpler than the one I filed** — no `ORDER BY` or
`LIMIT` needed:
```sql
SELECT file_row_index() + 1 FROM big; -- 20 000 rows -> error
```
This also means "tables larger than one batch" in my original title is a
necessary but not sufficient condition, so I have retitled the issue.
**This matters for whoever writes the regression test:** a fixture with
`LIMIT 3`, or with fewer than `batch_size` rows, will pass against broken code.
## Root cause: `Projection` is not in the extraction allowlist
DataFusion already has the general machinery for this —
`ExtractLeafExpressions` / `PushDownLeafProjections` in
`datafusion/optimizer/src/extract_leaf_expressions.rs` — which pulls
`ExpressionPlacement::MoveTowardsLeafNodes` sub-expressions out into a
`__datafusion_extracted_N` projection that the source can absorb. And
`file_row_index()` already declares that placement
(`functions/src/core/file_row_index.rs:85`), as does `input_file_name()`.
The same wrapped shape works in a `WHERE` clause, because `Filter` **is** on
the allowlist:
```sql
SELECT column1 FROM big WHERE file_row_index() + 1 = 3; -- ok
```
```
Filter: __datafusion_extracted_1 + Int64(1) = Int64(3)
Projection: file_row_index() AS __datafusion_extracted_1, big.column1
TableScan: big projection=[column1]
```
Note this plan contains a `RepartitionExec` too and still works, because by
then the marker is an ordinary column and the `all_columns` gate no longer
applies.
In a `SELECT` list there is no extraction at all, because of the allowlist
in `extract_from_plan` (`extract_leaf_expressions.rs:259`):
```rust
if !matches!(
&plan,
LogicalPlan::Aggregate(_)
| LogicalPlan::Filter(_)
| LogicalPlan::Sort(_)
| LogicalPlan::Limit(_)
| LogicalPlan::Join(_)
) {
return Ok(Transformed::no(plan));
}
```
Measured, same `marker + 1` shape in each node type:
| node type | on allowlist? | result |
|---|---|---|
| `Filter` (`WHERE`) | yes | ok |
| `Sort` (`ORDER BY`) | yes | ok |
| `Aggregate` (`sum`/`GROUP BY`/`HAVING`) | yes | ok |
| `Join` (`ON`) | **yes** | **error** — see below |
| `Window` (`OVER`) | no | **error** |
| **`Projection` (`SELECT` list)** | **no** | **error — this issue** |
The module docs give the rationale for excluding `Projection`:
> The split has no use there: the source absorbs the leaf expressions of the
original projection just as well.
That holds only while the projection actually reaches the source. Physical
planning then inserts the `RepartitionExec` per `n_rows > batch_size`,
`RepartitionExec::try_swapping_with_projection` refuses projections that are
not `all_columns`, and nothing re-establishes the assumption.
This also explains why the bare call works: a bare `file_row_index()`
projection is already in the shape the source can absorb.
## Scope is wider than the title
**`input_file_name()` is affected identically** (all without `LIMIT`):
| shape | result |
|---|---|
| `input_file_name()` | ok |
| `input_file_name() \|\| 'x'` | **error** |
| `upper(input_file_name())` | **error** |
| `input_file_name() \|\| CAST(file_row_index() AS VARCHAR)` | **error** |
**`get_field()` is the third member of the `MoveTowardsLeafNodes` family,
and it silently loses Parquet nested-column pruning** rather than erroring,
because it *is* evaluable. On a 20 000-row table with a struct column:
| query | scan's physical projection | reads |
|---|---|---|
| `SELECT s['a'] FROM t` | `projection=[get_field(s@0, a) as v]` | just the
`a` leaf |
| `SELECT s['a'] + 1 FROM t` | `projection=[s]` | **the whole struct** |
Same root cause, confirmed by removing the repartition — pruning comes back:
| setting | scan projection for `SELECT s['a'] + 1` |
|---|---|
| default (`RepartitionExec` present) | `projection=[s]` |
| `target_partitions = 1` | `projection=[get_field(s@0, a) + 1 as v]` |
| `batch_size = 30000` | `projection=[get_field(s@0, a) + 1 as v]` |
So the general statement is: **any `MoveTowardsLeafNodes` expression wrapped
in another expression loses leaf pushdown once a `RepartitionExec` is
inserted.** For the two marker UDFs that is an execution error; for `get_field`
it is a silent performance regression in what is arguably the extraction rule's
headline feature.
## Not the same gap as #23531
`datafusion.execution.parquet.pushdown_filters` makes no difference here —
the projection case fails with it both on and off — so this is a distinct gap
from #23531 / #23638.
## Suggested fix direction
Extend extraction to projections whose expressions *contain* a
`MoveTowardsLeafNodes` sub-expression, turning
```
Projection: file_row_index() + 1 AS v, column1
```
into
```
Projection: __datafusion_extracted_1 + 1 AS v, column1
Projection: file_row_index() AS __datafusion_extracted_1, column1
```
The inner pure extraction projection is absorbable by the source (exactly
the shape that works today), and the outer one references only columns. This
reuses the already-accepted design from the `Filter` path and needs no
loosening of the `all_columns` guard in `RepartitionExec`.
Two constraints for whoever implements it:
- The allowlist comment excludes node types whose output schema depends on
their expressions, and `Projection` is one of those. The split must keep the
outer projection's alias list and output field names byte-identical, changing
only the inner marker into a column reference.
- Detection should go through `ExpressionPlacement::MoveTowardsLeafNodes`
rather than another `try_downcast_func` chain. #23638 added a hardcoded chain
and its author noted: *"If we keep adding this sort of metadata functions, we
might want a better API to detect them"* — that API now exists and all three
functions already declare it, so keying off placement also covers future
metadata functions.
## Two separate gaps found along the way
These are distinct from this issue and I don't think they belong in the same
fix. Happy to file them separately if useful:
1. **`Join` is on the allowlist but still fails.** Extraction routes each
extracted sub-expression to an input by checking which input's schema holds all
of the expression's column references. A marker call has *no* column
references, so it cannot be routed. Pre-extracting it in a subquery works:
```sql
SELECT count(*) FROM (SELECT file_row_index() AS r, column1 FROM big) a
JOIN big b ON a.column1 = b.column1 AND a.r >= 0; -- ok
```
2. **`file_row_index()` does not work on CSV or JSON at all** — even the
bare form errors, independently of row count or `target_partitions` (verified
on a 3-row file and with `target_partitions = 1`). The only production caller
of the rewrite is the Parquet source. `input_file_name()` bare does work on
CSV. So `file_row_index()` is effectively Parquet-only, which the documentation
does not say.
All of the above verified on `c120619de` (`48.0.0-rc2-4901-gc120619de`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
