u70b3 opened a new pull request, #23684:
URL: https://github.com/apache/datafusion/pull/23684
## Which issue does this PR close?
- Closes #23440
- Closes #22190
Note: #22190 is currently assigned to @xiedeyantu (idle since 2026-05-17, no
PR). #23440 is the same root cause with a smaller reproducer; this PR fixes
both. Happy to coordinate if the assignee has work in progress.
## Rationale for this change
When `TopKAggregation` pushes a `LIMIT` into a MIN/MAX aggregate (`SELECT s,
MAX(y) AS max_y FROM t GROUP BY s ORDER BY max_y DESC NULLS FIRST LIMIT k`), a
group whose aggregate inputs are all NULL disappears from the result entirely,
instead of appearing with a NULL aggregate value.
Minimal reproducer (returns 0 rows on main; must return 1 row `NULL`):
```sql
CREATE TABLE t0 AS SELECT * FROM (VALUES ('gamma', CAST(NULL AS DOUBLE)))
v(s, y);
SELECT max_y FROM (
SELECT s, MAX(y) AS max_y FROM t0 GROUP BY s
) ORDER BY max_y DESC NULLS FIRST LIMIT 3;
```
DuckDB, PostgreSQL, Polars and SQLite all preserve the grouped row; only
DataFusion dropped it (confirmed by two independent reporters in #23440 /
#22190).
Root cause: `GroupedTopKAggregateStream::intern()` skips rows with NULL
aggregate inputs (correct: MIN/MAX ignore NULL inputs), but a group whose
inputs are ALL NULL then never enters the `PriorityMap` and is lost at emit
time. The existing `null_group_seen` recovery only covers the GROUP BY-only
(DISTINCT) path.
## What changes are included in this PR?
- `datafusion/physical-plan/src/aggregates/topk/hash_table.rs`
- Register all-NULL groups in the hash table with a sentinel `heap_idx`
(`NULL_HEAP_IDX`); at most `limit` such groups are tracked (they tie on the
sort key, so any `limit` of them is a valid top-k superset).
- `find_or_insert` converts a NULL-registered group into a valued group
when its first non-NULL value arrives, and reports the conversion via the new
`InsertKind` so callers can keep their accounting in sync without an extra
lookup.
- `remove_if_full` now counts only valued groups against the limit,
preserving the 1:1 heap-slot ↔ valued-group invariant.
- `datafusion/physical-plan/src/aggregates/topk/priority_map.rs`
- `PriorityMap::insert_null` for the stream.
- On the `is_worse` bail path, a NULL-registered group whose value loses
to the current top-k is unregistered (it can no longer reach the top-k by
MIN/MAX monotonicity, and must not be emitted as NULL either).
- The all-NULL group count is kept as a plain field (see Performance
below).
- `emit()` appends NULL groups after the heap contents (output ≤ 2×limit
superset; the parent `SortExec(TopK)` applies the final ordering/truncation
with full `nulls_first` semantics, so no null-ordering knowledge is needed
here).
- `datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs`
- `intern()` routes NULL rows of the MIN/MAX path to `insert_null` (7
lines). The GROUP BY-only path is unchanged.
- `datafusion/sqllogictest/test_files/aggregates_topk.slt`
- Regression coverage for #23440 / #22190, superset truncation (5 NULL
groups + LIMIT 2), NULLS FIRST/NULLS LAST variants, cross-batch NULL/value
mixing via `batch_size = 1`.
- **One existing expectation corrected** (`max(trace_id)` at ~L269): the
old expected output `c c, b b` encoded this very bug (the all-NULL group's
`max` is NULL and ranks first under `DESC NULLS FIRST`).
## Are these changes tested?
Yes:
- New sqllogictest cases in `aggregates_topk.slt` (including exact
reproducers from both issues).
- New unit tests in `priority_map.rs` / `hash_table.rs` (null emit, mixed
emit, cap, NULL→valued conversion, loser unregistration).
- Verified `./dev/rust_lint.sh`, the extended workspace suite (ci profile
with
`avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`),
and `cargo test -p datafusion-cli` — all green.
- Also cross-checked `target_partitions = 4` against
`enable_topk_aggregation = false` — results identical.
## Performance
`topk_aggregate` benchmarks were run against `main` (interleaved A/B, 3
rounds each). An initial version of this PR read the all-NULL group count
through a `dyn ArrowHashTable` call on the per-row bail path, which regressed
the worst-case benchmark (+13–38%). The count is now kept as a plain field in
`PriorityMap` and synchronized via `InsertKind`, so the hot path performs the
same `dyn` calls as `main` plus one predictable branch. Residual A/B deltas
after the fix are within the noise floor of the shared dev machine used for
testing (the untouched non-TopK control benchmark swings ±10% between identical
rounds); happy to re-run on dedicated hardware or incorporate reviewer
suggestions if any regression shows up there.
## Are there any user-facing changes?
Only the bug fix itself: queries that previously (incorrectly) dropped
all-NULL groups under `ORDER BY <min/max> ... LIMIT` now return those groups
with NULL aggregate values, matching PostgreSQL/DuckDB. No API or configuration
changes.
## Known limitation (intended as follow-up issue)
If a group enters the top-k with a non-NULL value, is later evicted (its map
slot is reused), and only afterwards produces a NULL row, it can be
re-registered as an all-NULL group. This only affects *which* tied NULL groups
appear (not the row count), only under `NULLS FIRST`, and requires the eviction
to happen before the NULL row arrives. We verified the common interleavings
(including `batch_size = 1` stress and partitioned execution) produce correct
results; closing this residual race requires the heap to carry NULL values
natively, which is a larger change.
--
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]
