akupchinskiy opened a new issue, #27083:
URL: https://github.com/apache/datafusion/issues/27083
## Describe the bug
A recursive CTE that joins on string keys read from Parquet fails with
`Failed to allocate
additional ... for RecursiveQuery` even when the pool is several times
larger than what the process
actually uses. In the example below the query fails with a 4 GB pool and
completes with 16 GB, while
peaking at 1.2 GB of RSS; the same query over `Utf8` columns runs in a 1 GB
pool.
The cause is in `RecursiveQueryStream::push_batch`
(`datafusion/physical-plan/src/recursive_query.rs`).
It reserves `batch.get_array_memory_size()` for every batch it keeps for the
next iteration:
```rust
if let Err(e) = self.reservation.try_grow(batch.get_array_memory_size()) {
return Poll::Ready(Some(Err(e)));
}
```
The recursive term here is a hash join, and its output is built with `take`.
For a
`StringViewArray`, `take` copies the views but keeps references to all of
the build side's data
buffers. So every output batch reports the build side's string data as its
own, and the reservation
counts the same buffers again for each buffered batch.
I ran into this with a benchmark of mine: a span-tree walk over traces fails
at SF 10 with a 64 GB
pool, while DuckDB runs it in 1 GB.
## To Reproduce
`datafusion-cli` from `main` (f058e2d), default settings. An org chart of 2M
employees with
string ids:
```sql
COPY (
SELECT 'emp-' || md5(CAST(value AS VARCHAR)) AS employee_id,
CASE WHEN value = 0 THEN NULL ELSE 'emp-' || md5(CAST((value - 1) /
10 AS VARCHAR)) END AS manager_id
FROM range(2000000)
) TO 'employees.parquet' STORED AS PARQUET;
CREATE EXTERNAL TABLE employees STORED AS PARQUET LOCATION
'employees.parquet';
WITH RECURSIVE chain AS (
SELECT employee_id, 0 AS level FROM employees WHERE manager_id IS NULL
UNION ALL
SELECT e.employee_id, c.level + 1
FROM employees e JOIN chain c ON e.manager_id = c.employee_id
)
SELECT level, count(*) AS headcount FROM chain GROUP BY level ORDER BY level;
-- the same ids as Utf8; run the same query over employees_utf8
CREATE VIEW employees_utf8 AS
SELECT arrow_cast(employee_id, 'Utf8') AS employee_id,
arrow_cast(manager_id, 'Utf8') AS manager_id
FROM employees;
```
Run with `datafusion-cli -m <pool> -f repro.sql` (Apple M2 Pro, 10
partitions):
| columns | pool | result | peak RSS |
|---|---|---|---|
| Utf8View (default) | 4 GB | `Failed to allocate additional 24.1 MB for
RecursiveQuery with 579.0 MB already allocated` | 1.1 GB |
| Utf8View | 16 GB | OK | 1.2 GB |
| Utf8 (`employees_utf8`) | 4 GB | OK | 0.8 GB |
| Utf8 | 1 GB | OK | 0.8 GB |
It gets much worse when the planner picks a `CollectLeft` join, since then
every buffered batch
references the whole build side. On 4M spans in chains of 8, each buffered
batch was charged about
555 MB, and the query failed with every pool I tried up to 16 GB, at under
800 MB of RSS.
Setting `datafusion.execution.parquet.schema_force_view_types = false` works
around it.
## Expected behavior
`RecursiveQueryExec` reserves each buffer shared by its buffered batches
once per iteration. The
hash join, asof join and symmetric hash join already do this for their
buffered batches with `RecordBatchMemoryCounter`
(`datafusion/common/src/utils/memory.rs`).
## Additional context
Related issues, none of which covers this:
- #25712: `CollectLeft` joins multiply buffer references within one array.
Deduplicating within an
array doesn't stop the same allocation from being counted again in the
next batch.
- #13377: deduplicates buffers within one batch, not across batches.
- #16841: the general discussion of accounting for shared buffers.
- #25637: `RecursiveQueryExec` can't spill. That's a separate problem.
I have a fix ready (PR to follow): keep a `RecordBatchMemoryCounter` in
`RecursiveQueryStream`,
reserve `count_batch_with_array_overhead(&batch)` in `push_batch`, and start
a new counter for each
iteration. With it the `CollectLeft` example passes with a 1 GB pool (peak
RSS 1.16 GB) and the
org chart with 2 GB. The org chart still fails at 1 GB; the error reports
409 MB across
RecursiveQuery's reservations and 59.1 MB each for two `HashJoinInput`
partitions (only the top
three consumers are shown).
--
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]
