Disproportionate memory use for `DISTINCT ON` query

[connec]

Describe the bug

We have a parquet file (built from public data) with 106 columns and 1M rows which is 131.14 MiB in size (compressed, 913.89 MiB uncompressed).

When running a DISTINCT ON query using the unbounded memory pool, memory use climbs to over 160 GiB for this query:

SELECT DISTINCT
  ON ("ADDRESS1", "ADDRESS2", "ADDRESS3", "POSTCODE") *
FROM
  table
ORDER BY
  "ADDRESS1",
  "ADDRESS2",
  "ADDRESS3",
  "POSTCODE",
  "INSPECTION_DATE" DESC

When using a fair spill pool with 10 GiB, memory usage reaches "only" 30 GiB.

These results were observed on my local machine (MacBook Pro). On a production machine with the same 10 GiB limit we have seen a graceful allocation failure:

Resources exhausted: Failed to allocate additional 55.0 MB for GroupedHashAggregateStream[3]

This makes me think it could be the same underlying issue as #13831, exacerbated by the many columns.

To Reproduce

See the parquet file and SQL query in the description above.

Expected behavior

In an ideal world, the memory usage for this query would respect the memory pool limit (or only use "small" allocations as described in the docs).

Additional context

I'm happy to help diagnose this further (and potentially fix) with some advice on how to profile the memory use or narrow down the cause. For now I just wanted to capture the issue to see if it's known as I imagine it won't be an easy fix 😄

I know there are a few issues related to memory management floating around atm but none that I could see directly mentioned DISTINCT ON, so apologies if this is a duplicate.

[adriangb]

I recommend using https://github.com/polarsignals/rust-jemalloc-pprof

[alamb]

It may also be related to the fact that the group by hash code can allocate large single swaths of memory

@Rachelint was working on improvements in this area

• Intermediate result blocked approach to aggregation memory management #15591

[alamb]

BTW I can easily reproduce the reported symptoms using datafusion-cli locally:

> SELECT DISTINCT
  ON ("ADDRESS1", "ADDRESS2", "ADDRESS3", "POSTCODE") *
FROM
  'https://digitalsociety-public.fsn1.your-objectstorage.com/f94d0c87-8798-4bf6-9c98-8d89971e2539.parquet'
ORDER BY
  "ADDRESS1",
  "ADDRESS2",
  "ADDRESS3",
  "POSTCODE",
  "INSPECTION_DATE" DESC;

[alamb]

I believe @Omega359 has also spoken about this usecase being a challenge (basically deduplicating large datasets).

While DataFusion's grouping operator is already pretty optimized, it clearly could be improved for this particular case.

For now I just wanted to capture the issue to see if it's known as I imagine it won't be an easy fix 😄

I agree this is not likely to be an easy fix

But frankly just having a good reproducer (as you have provided in this ticket) is a huge step forward

[alamb]

I think there are two issues we should explore:

1. Keeping within the memory budget (basically make the accounting better when possible)
2. Improving DataFusion's ability to deduplicate data with less memory

Maybe we could come up with a better operator that specializes with deduplication (basically for queries with only GROUP BY and no aggregations) 🤔

maybe @2010YOUY01 or @rluvaton have some thoughts on the latter

[rluvaton]

I'm surprised we got the error and we did not spill, also is the memory that you observed came from the machine or datafusion (so I can better know if this is over accounting issue or actual memory usage)

[Omega359]

fwiw I use duckdb for dedup and while it's noticeably better than DF memory wise it is still a concern for it as well.

[rluvaton]

Looking at the code I realized that for performance optimization we can do is to calculate the ordering once when we have order by and accumulators like first can use it

[Dandandan]

• Another thing is that datafusion will store the data for each unique combination, rather than only once for each unique value. For columns with larger data and muliple columns in group by it could save memory to dictionary encode the columns in the group by (I think there is an issue for it somewhere to explore the idea) and reference them using index
• For queries like distinct it could be better to remove duplicates through sorting rather than to use hash aggregation. In a query like the example, it would make even more sense as the result needs to be sorted anyway

[connec]

DISTINCT ON (...) requires the query to be sorted on those fields. I guess that should allow the deduplication to happen in a streaming fashion on the sorted results?

I.e. (very imprecise) calculate a key for each row as its visited. Remember just the latest key. Skip a row if the key matches the latest or update the latest key if not (plus merge deduplication if this is happening in multiple partitions).

[Dandandan]

I.e. (very imprecise) calculate a key for each row as its visited. Remember just the latest key. Skip a row if the key matches the latest or update the latest key if not (plus merge deduplication if this is happening in multiple partitions).

Yeah I think that could be how an efficient deduplicate operation could work.

[alamb]

Indeed -- if you already know the data is sorted on the DISTINCT ON keys, you can do deduplication with a single pass through the data with minimal memory requirements, following the suggestion from @connec

This is the exact strategy we use in InfluxDB IOx to deduplicate (which is how we implement insert order resolution)

The code for it is here

• https://github.com/influxdata/influxdb3_core/blob/main/iox_query/src/provider/deduplicate.rs

The biggest challenge is having the data sorted by the DISTINCT ON keys as (re)sorting large datasets is also resource intensive (memory, CPU and potentially IO)

[alamb]

I tried to get DataFusion to use the sort order and avoid using much memory but was not successful

I can resort the file quite efficiently with minimal RAM no problem:

$ datafusion-cli -m 4GB

> COPY (
  SELECT *
    FROM
      'https://digitalsociety-public.fsn1.your-objectstorage.com/f94d0c87-8798-4bf6-9c98-8d89971e2539.parquet'
    ORDER BY
      "ADDRESS1",
      "ADDRESS2",
      "ADDRESS3",
      "POSTCODE",
      "INSPECTION_DATE" DESC
  )
to 'f94d0c87-8798-4bf6-9c98-8d89971e2539.reordered.parquet';
+---------+
| count   |
+---------+
| 1523389 |
+---------+
1 row(s) fetched.
Elapsed 41.416 seconds.

Then I can tell DataFusion it is sorted:

> CREATE EXTERNAL TABLE t
STORED AS PARQUET
LOCATION 'f94d0c87-8798-4bf6-9c98-8d89971e2539.reordered.parquet'
WITH ORDER ( "ADDRESS1", "ADDRESS2", "ADDRESS3", "POSTCODE", "INSPECTION_DATE" DESC);
0 row(s) fetched.
Elapsed 0.004 seconds.

But when I try to deduplicate this I still run out of memory:

 SELECT DISTINCT
  ON ("ADDRESS1", "ADDRESS2", "ADDRESS3", "POSTCODE") *
FROM
  t
ORDER BY
  "ADDRESS1",
  "ADDRESS2",
  "ADDRESS3",
  "POSTCODE",
  "INSPECTION_DATE" DESC;

Resources exhausted: Additional allocation failed with top memory consumers (across reservations) as:
  ExternalSorter[7]#306(can spill: true) consumed 227.7 MB,
  ExternalSorter[13]#324(can spill: true) consumed 227.5 MB,
  ExternalSorter[14]#327(can spill: true) consumed 227.2 MB.
Error: Failed to allocate additional 679.8 MB for GroupedHashAggregateStream[11] (, , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , ) with 6.8 MB already allocated for this reservation - 448.6 MB remain available for the total pool

I would have expected that this could use GroupOrderingFull and avoid large amounts of memory usage, but that does not appear to be the case. Maybe it is a bug we should look into 🤔

datafusion/datafusion/physical-plan/src/aggregates/order/mod.rs

Lines 38 to 39 in 1daa5ed

	/// Groups are entirely contiguous,
	Full(GroupOrderingFull),