cool9850311 opened a new issue, #27131:
URL: https://github.com/apache/datafusion/issues/27131
### Describe the bug
An expression with `ExpressionPlacement::MoveTowardsLeafNodes` placement is
supposed to be pushed down to the file scan, where the source can either supply
it as a column or prune to it. That pushdown is lost whenever the expression is
**wrapped in another expression** and a `RepartitionExec` has been inserted
between the `ProjectionExec` and the `DataSourceExec`.
All three functions that declare this placement are affected, in two symptom
classes:
| function | evaluable directly? | symptom when wrapped |
|---|---|---|
| `file_row_index()` | no | **execution error** |
| `input_file_name()` | no | **execution error** |
| `get_field()` (`s['a']`) | yes | **silently loses Parquet nested-column
pruning** |
```
Execution error: file_row_index() is source dependent and cannot be
evaluated directly
```
The bare call always works; wrapping it in *anything* breaks it:
```sql
SELECT file_row_index() FROM big; -- ok
SELECT file_row_index() + 1 FROM big; -- Execution error
```
#### Root cause
The general machinery for this already exists: `ExtractLeafExpressions` /
`PushDownLeafProjections` in
`datafusion/optimizer/src/extract_leaf_expressions.rs` pulls
`MoveTowardsLeafNodes` sub-expressions out into a `__datafusion_extracted_N`
projection that the source can absorb. All three functions already declare the
placement (e.g. `functions/src/core/file_row_index.rs:85`).
But `extract_from_plan` (`extract_leaf_expressions.rs:259`) only extracts
from an allowlist of node types, and **`Projection` is not on it**:
```rust
// Only extract from plan types whose output schema is predictable after
// expression rewriting. ...
if !matches!(
&plan,
LogicalPlan::Aggregate(_)
| LogicalPlan::Filter(_)
| LogicalPlan::Sort(_)
| LogicalPlan::Limit(_)
| LogicalPlan::Join(_)
) {
return Ok(Transformed::no(plan));
}
```
The same wrapped shape therefore works in a `WHERE` clause, where `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]
```
That plan contains a `RepartitionExec` too and still works, because by then
the marker is an ordinary column.
Measured, the 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** — separate cause, see Additional context |
| `Window` (`OVER`) | no | **error** |
| **`Projection` (`SELECT` list)** | **no** | **error — this issue** |
The module docs state 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.
`EnsureRequirements` then inserts a `RepartitionExec` when the estimated row
count exceeds `batch_size`
(`ensure_requirements/enforce_distribution.rs:1167`),
`RepartitionExec::try_swapping_with_projection` refuses any projection that is
not `all_columns` (`physical-plan/src/repartition/mod.rs:1949`), and nothing
re-establishes the assumption. The bare call works because a bare marker
projection is already in the shape the source absorbs.
### To Reproduce
In `datafusion-cli`, stock configuration:
```sql
COPY (SELECT value AS column1 FROM range(0, 20000))
TO '/tmp/df_leaf_repro/big.parquet' STORED AS PARQUET;
CREATE EXTERNAL TABLE big
STORED AS PARQUET LOCATION '/tmp/df_leaf_repro/big.parquet';
SELECT file_row_index() FROM big; -- ok
SELECT file_row_index() + 1 FROM big; -- Execution error
```
> [!IMPORTANT]
> **A bare `LIMIT n` masks this bug.** It is pushed into the scan, so the
row estimate drops to `n`, falls below `batch_size`, and no `RepartitionExec`
is inserted. A regression test written with `LIMIT 3`, or against a table
smaller than `batch_size`, **passes against broken code**.
>
> | query | result |
> |---|---|
> | `SELECT file_row_index() + 1 FROM big LIMIT 3` | **ok** — 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 |
Expression shapes, all without `LIMIT`:
| shape | result |
|---|---|
| `file_row_index()` | ok |
| `file_row_index() AS r` | ok |
| `file_row_index() + 1` | **error** |
| `arrow_typeof(file_row_index())` | **error** |
| `CAST(file_row_index() AS INT)` | **error** |
| `SELECT r + 1 FROM (SELECT file_row_index() AS r FROM big)` | **error** —
subquery is flattened |
`input_file_name()` behaves identically: bare is ok, while
`input_file_name() || 'x'`, `upper(input_file_name())` and `input_file_name()
|| CAST(file_row_index() AS VARCHAR)` all error.
#### `get_field()`: same cause, silent performance regression
```sql
COPY (SELECT named_struct('a', value, 'b', value * 2, 'c', CAST(value AS
VARCHAR)) AS s
FROM range(0, 20000)) TO '/tmp/df_leaf_repro/st.parquet' STORED AS
PARQUET;
CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION
'/tmp/df_leaf_repro/st.parquet';
EXPLAIN SELECT s['a'] AS v FROM t;
EXPLAIN SELECT s['a'] + 1 AS v FROM t;
```
| query | `DataSourceExec` projection | reads |
|---|---|---|
| `SELECT s['a'] FROM t` | `projection=[get_field(s@0, a) as v]` | only the
`a` leaf |
| `SELECT s['a'] + 1 FROM t` | `projection=[s]` | **the entire struct** |
No error, just substantially more I/O than necessary — and nested-column
pruning is the headline feature of the extraction rule.
#### Controls confirming the mechanism
Each of these removes the `RepartitionExec` and restores correct behaviour,
on the same table and data:
| setting | `file_row_index() + 1` | `s['a'] + 1` scan projection |
|---|---|---|
| default | error | `projection=[s]` |
| `SET datafusion.execution.batch_size = 30000` | ok |
`projection=[get_field(s@0, a) + 1 as v]` |
| `SET datafusion.execution.target_partitions = 1` | ok |
`projection=[get_field(s@0, a) + 1 as v]` |
Raising `batch_size` above the row count changes only the threshold at
`enforce_distribution.rs:1167`, which is what promotes this from correlation to
demonstrated cause.
`datafusion.execution.parquet.pushdown_filters` makes no difference — the
projection case fails with it both on and off.
`EXPLAIN` on the failing query **succeeds**; the failure appears only once
execution starts.
### Expected behavior
A wrapped `MoveTowardsLeafNodes` expression should reach the scan just as
the bare form does, so that:
- `file_row_index() + 1` and `input_file_name() || 'x'` return results
instead of erroring, and
- `s['a'] + 1` prunes to the `a` leaf instead of reading the whole struct.
Failing that, the two marker functions should at least be rejected at
**plan** time with an actionable message rather than accepted, passed by
`EXPLAIN`, and failed mid-execution.
### Additional context
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, so no
loosening of the `all_columns` guard in `RepartitionExec` is needed. Two
constraints:
- The allowlist comment excludes node types whose output schema depends on
their expressions, and `Projection` is one of them. The split must keep the
outer projection's alias list and output field names identical, changing only
the inner sub-expression into a column reference.
- Detection should key off `ExpressionPlacement::MoveTowardsLeafNodes`
rather than per-function downcasts. #23638 added a hardcoded
`try_downcast_func` chain for `FileRowIndexFunc`/`InputFileNameFunc`, and noted
at the time: *"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
declare it, so keying off placement also covers metadata functions added later.
Related: #23531 / #23638 are the same `"is source dependent and cannot be
evaluated directly"` error from a *different* pushdown gap (`input_file_name()`
inside a Parquet `RowFilter`), fixed by excluding those UDFs from filter
pushdown. `file_row_index()` was added in #22604, `input_file_name()` in #22978.
Two further gaps found while investigating, which I believe are **out of
scope here** and need their own issues:
1. **`Join` is on the allowlist but still fails.** Extraction routes each
extracted sub-expression to an input by checking which input's schema contains
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 and `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 state.
Verified on `c120619de` (`48.0.0-rc2-4901-gc120619de`), `datafusion-cli`
reporting `57.0.0-dev`, macOS arm64, default `target_partitions` (= core count).
### Willingness to contribute
I plan to contribute a fix for this bug independently.
--
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]
