You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "ydgandhi (via GitHub)" <gi...@apache.org> on 2026/03/21 07:29:28 UTC
[PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
ydgandhi opened a new pull request, #21088:
URL: https://github.com/apache/datafusion/pull/21088
Add MultiDistinctCountRewrite in datafusion-optimizer and register it in Optimizer::new() after SingleDistinctToGroupBy. Rewrites 2+ simple COUNT(DISTINCT) on different args into a join of two-phase aggregates; filter distinct_arg IS NOT NULL on each branch for correct NULL semantics.
✅ Unit tests in datafusion-optimizer; ✅ SQL integration test (NULLs) in core_integration.
## Which issue does this PR close?
- Closes #21087.
## Rationale for this change
Queries with multiple COUNT(DISTINCT col_i) in the same GROUP BY can force independent distinct state per aggregate (e.g. separate hash sets), which scales poorly in memory when several high-cardinality distinct columns appear together.
DataFusion already optimizes the single shared distinct field case via SingleDistinctToGroupBy. This PR adds a conservative logical rewrite for multiple distinct COUNT(DISTINCT …) arguments by splitting work into per-distinct branches joined on the group keys, which reduces peak memory for eligible plans.
COUNT(DISTINCT x) must ignore NULL x; the rewrite applies x IS NOT NULL on each distinct branch before inner grouping so semantics stay aligned with count_distinct behavior.
## What changes are included in this PR?
New module:
- datafusion/optimizer/src/multi_distinct_count_rewrite.rs — MultiDistinctCountRewrite (OptimizerRule, bottom-up).
- Register rule in datafusion/optimizer/src/optimizer.rs immediately after SingleDistinctToGroupBy.
- Export module from datafusion/optimizer/src/lib.rs.
Tests:
- datafusion-optimizer: rewrite vs no-op cases (single distinct, mixed aggs, two distincts).
- datafusion core_integration: datafusion/core/tests/sql/aggregates/multi_distinct_count_rewrite.rs — SQL over MemTable with NULLs.
## Are these changes tested?
Yes.
- cargo test -p datafusion-optimizer (includes new unit tests).
- cargo test -p datafusion --test core_integration multi_count_distinct (SQL + NULLs).
## Are there any user-facing changes?
- Behavior / plans: For queries that match the rule, logical plans (and thus EXPLAIN output) can differ: eligible multi–COUNT(DISTINCT) aggregates may appear as joins of sub-aggregates instead of a single Aggregate with multiple distinct counts.
- Results: Intended to be semantics-preserving for supported patterns (including NULL handling via filters).
- Public API: No intentional breaking changes to public Rust APIs; this is an internal optimizer rule enabled by default.
- Docs: No user guide update required unless maintainers want an “optimizer behavior” note; can add if requested.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4108104692
Thanks @Dandandan @xiedeyantu for the discussion.
**Join cost / default behavior:** Agreed that joining on grouping keys isn’t free. We’ve added **`datafusion.optimizer.enable_multi_distinct_count_rewrite`** with default **`false`**, so the rewrite is **opt-in** until we have benchmarks that justify turning it on by default. Sessions can enable it explicitly when they want to try the plan shape.
**Benchmarks:** We’re aligned that we should measure **latency and memory** vs baseline across scenarios (e.g. multiple `COUNT(DISTINCT …)` with `GROUP BY`, varying group-key cardinality vs distinct cardinality). We’ll follow up with numbers on the PR.
**GroupsAccumulator / execution-layer improvements:** That work is complementary: better distinct accumulators help **how** each aggregate runs; this rule changes **logical plan shape** when several large distincts share one aggregate. Both can coexist on the roadmap.
**Tests:** We added optimizer + SQL integration coverage (including cases with `lower`/`CAST` inside `COUNT(DISTINCT …)` and a test that the rule is a no-op when the config is off). Happy to iterate on naming and placement with maintainers.
> > I think a rewrite like this might be useful, but I think it can also hurt performance because of the join on grouping keys. So I think it needs to have a config value (off by default) or when enabled some benchmarks showing that it is better in large majority of the cases.
> > I am also wondering if mostly for memory usage a `GroupsAccumulator` for distinct count / sum might give similar/more improvements.
>
> @Dandandan Thank you for the explanation. It’s true that this would add a hash join, but if aggregation can be performed in parallel, there might be advantages in scenarios with two or more COUNT(DISTINCT) operations. I agree to run performance tests across multiple scenarios to evaluate the actual results.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4115356321
Hi! @xiedeyantu, great feedback. Let me explain my thought process for some of the design choices regarding the tests. See inline
> Thank you for the detailed test! But I have a few questions.
>
> 1. Why add an extra
> "SUM" outside?
The outer `SELECT SUM(...) FROM ( ... GROUP BY ... )` is a **harness choice**:
- The work under measurement is the **inner** grouped aggregate (multiple `COUNT(DISTINCT …)`).
- The outer `SUM` collapses per-group results to **one row**, which avoids dumping thousands of rows when comparing wall time / peak RSS.
- To benchmark **only** the inner shape, run the subquery alone (or use `SELECT COUNT(*) FROM (inner) t` if you want a scalar sink without changing the inner aggregates).
> 2. I think it’s better for
> "CAST" not to involve expression calculations.
Agreed for a clean test: both distincts should use `CAST` on a column only. Our Scenario 2 used `CAST(l_extendedprice AS INT)` (column-only) and `CAST(l_discount * 100 AS INT)` (expression inside the cast). We'll switch the second to something like `CAST(l_discount AS INT)` (or another column-only cast) so we're not mixing in arithmetic inside `CAST`, and so it lines up with the rule's conservative “safe” distinct arguments.
> 3. It seems the third case is a low-cardinality test — could you show a high-cardinality scenario?
The `lower()` case used `GROUP BY l_returnflag`, which is only a few groups. For a high-cardinality variant we'll use the same select list but a denser key, e.g. `GROUP BY l_suppkey` (or `l_orderkey`), so the join after the rewrite has many more keys — still with `COUNT(DISTINCT lower(l_shipmode))` and `COUNT(DISTINCT l_shipinstruct)`.
We'll follow up with updated query text and numbers when we rerun.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4110284239
Following up with additional real-data benchmark results for `MultiDistinctCountRewrite` (config-gated by `datafusion.optimizer.enable_multi_distinct_count_rewrite`, default `false`).
## Dataset generation / setup
I used local TPCH SF1 parquet generated with:
```bash
cargo install tpchgen-cli
tpchgen-cli --scale-factor 1 --format parquet --parquet-compression 'ZSTD(1)' --parts 1 --output-dir benchmarks/data/tpch_sf1
```
Data path used in all queries:
`benchmarks/data/tpch_sf1/lineitem/*.parquet`
For timing/memory, each run used:
```bash
/usr/bin/time -l datafusion-cli -q --format csv -c "SET datafusion.optimizer.enable_multi_distinct_count_rewrite = <true|false>;" -c "<QUERY>"
```
(Reporting wall time + max RSS from `time -l`.)
## Queries tested
### Scenario 1: 3 distincts (grouped)
```sql
SELECT SUM(d1) AS s1, SUM(d2) AS s2, SUM(d3) AS s3
FROM (
SELECT l_suppkey,
COUNT(DISTINCT l_orderkey) AS d1,
COUNT(DISTINCT l_partkey) AS d2,
COUNT(DISTINCT l_linenumber) AS d3
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_suppkey
);
```
### Scenario 2: CAST-based distincts
```sql
SELECT SUM(cx) AS sx, SUM(cy) AS sy
FROM (
SELECT l_suppkey,
COUNT(DISTINCT CAST(l_extendedprice AS INT)) AS cx,
COUNT(DISTINCT CAST(l_discount * 100 AS INT)) AS cy
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_suppkey
);
```
### Scenario 3: lower() + string distinct
```sql
SELECT SUM(c1) AS s1, SUM(c2) AS s2
FROM (
SELECT l_returnflag,
COUNT(DISTINCT lower(l_shipmode)) AS c1,
COUNT(DISTINCT l_shipinstruct) AS c2
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_returnflag
);
```
## Results (latency + max RSS)
| Scenario | Rewrite OFF | Rewrite ON | Outcome |
|---|---:|---:|---|
| 1) three distincts | 0.28s, ~570 MB | 0.12s, ~730 MB | Faster, higher memory |
| 2) cast-based distincts | 0.17s, ~349 MB | 0.16s, ~350 MB | Near tie |
| 3) lower()+string distinct | 0.04s, ~99 MB | 0.06s, ~94 MB | Slightly slower, slightly lower memory |
All paired runs returned identical query results.
## Takeaway
The rewrite can improve latency in some multi-distinct grouped cases, but can also increase peak memory and/or be neutral/slightly slower depending on cardinality and expression shape. Keeping this rule opt-in (default off) remains appropriate while we expand benchmark coverage.
> Yes, I can see it would improve a lot of queries
>
> > Thanks @Dandandan @xiedeyantu for the discussion.
> > **Join cost / default behavior:** Agreed that joining on grouping keys isn’t free. We’ve added **`datafusion.optimizer.enable_multi_distinct_count_rewrite`** with default **`false`**, so the rewrite is **opt-in** until we have benchmarks that justify turning it on by default. Sessions can enable it explicitly when they want to try the plan shape.
> > **Benchmarks:** We’re aligned that we should measure **latency and memory** vs baseline across scenarios (e.g. multiple `COUNT(DISTINCT …)` with `GROUP BY`, varying group-key cardinality vs distinct cardinality). We’ll follow up with numbers on the PR.
> > **GroupsAccumulator / execution-layer improvements:** That work is complementary: better distinct accumulators help **how** each aggregate runs; this rule changes **logical plan shape** when several large distincts share one aggregate. Both can coexist on the roadmap.
> > **Tests:** We added optimizer + SQL integration coverage (including cases with `lower`/`CAST` inside `COUNT(DISTINCT …)` and a test that the rule is a no-op when the config is off). Happy to iterate on naming and placement with maintainers.
> > > > I think a rewrite like this might be useful, but I think it can also hurt performance because of the join on grouping keys. So I think it needs to have a config value (off by default) or when enabled some benchmarks showing that it is better in large majority of the cases.
> > > > I am also wondering if mostly for memory usage a `GroupsAccumulator` for distinct count / sum might give similar/more improvements.
> > >
> > >
> > > @Dandandan Thank you for the explanation. It’s true that this would add a hash join, but if aggregation can be performed in parallel, there might be advantages in scenarios with two or more COUNT(DISTINCT) operations. I agree to run performance tests across multiple scenarios to evaluate the actual results.
>
> Sounds good - agreed.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4133837810
> Hi @xiedeyantu update on the benchmarking. The new benchmarks use the **grouped aggregate only** (no outer `SUM` or other wrapper) so timings reflect the multi–`COUNT(DISTINCT …)` work directly. Note: Figures below are **warm-cache** repeats on one machine (relative A/B, not a CI guarantee). Also, these figures are from running each of the query 10 times and flipping the rewrite setting on and off using the grouped queries above (no outer `SUM`).
>
> ## Queries tested - original runs
> ### Scenario 1: three distincts (grouped)
> ```sql
> SELECT l_suppkey,
> COUNT(DISTINCT l_orderkey) AS d1,
> COUNT(DISTINCT l_partkey) AS d2,
> COUNT(DISTINCT l_linenumber) AS d3
> FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
> GROUP BY l_suppkey;
> ```
>
> ## Rerun: reviewer-aligned changes
> **CAST (Scenario 2):** both distincts use `CAST` on a **column** only (no `*` expression inside `CAST`):
>
> ```sql
> SELECT l_suppkey,
> COUNT(DISTINCT CAST(l_extendedprice AS INT)) AS cx,
> COUNT(DISTINCT CAST(l_discount AS INT)) AS cy
> FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
> GROUP BY l_suppkey;
> ```
>
> **High cardinality (Scenario 3):** same aggregates as before, but **`GROUP BY l_suppkey`** instead of `l_returnflag`:
>
> ```sql
> SELECT l_suppkey,
> COUNT(DISTINCT lower(l_shipmode)) AS c1,
> COUNT(DISTINCT l_shipinstruct) AS c2
> FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
> GROUP BY l_suppkey;
> ```
>
> ### Results (mean ± sample stddev, **n=10** per cell)
> Measured with `/usr/bin/time -l` and `target/release/datafusion-cli` from the repo root. RSS is **maximum resident set size** from `time -l` (MB here are decimal megabytes: bytes / 10⁶). Queries are the **reviewer-aligned grouped aggregates above** (no outer `SUM`).
>
> **wall (s) mean ± stddev**
>
> Scenario Rewrite OFF Rewrite ON Outcome
> 1) three distincts 0.2490 ± 0.0099 0.1000 ± 0.0000 Lower mean wall time with ON
> 2) cast-based (column-only `CAST`) 0.1350 ± 0.0053 0.0700 ± 0.0000 Lower mean wall time with ON
> 3) lower()+string (`GROUP BY l_suppkey`) 3.3920 ± 0.6496 0.0800 ± 0.0000 Much lower mean wall time with ON
> **max RSS (MB) mean ± stddev**
>
> Scenario Rewrite OFF Rewrite ON Outcome
> 1) three distincts 590.03 ± 10.12 722.62 ± 12.70 Higher mean peak RSS with ON
> 2) cast-based (column-only `CAST`) 331.83 ± 4.56 638.49 ± 8.29 Higher mean peak RSS with ON
> 3) lower()+string (`GROUP BY l_suppkey`) 9755.13 ± 114.94 277.66 ± 1.71 Much lower mean peak RSS with ON
> Row outputs matched OFF vs ON for every scenario (same grouped aggregate result). Wall time from `time` is low-precision on macOS (e.g. all runs may round to the same hundredth of a second, so stddev can be 0).
Thank you for the detailed testing! The test results look excellent to me. I have no further suggestions at this time.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4539163793
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4103382064
I personally think this PR is a more generalized rewrite approach. #20940 might cover a limited range of scenarios, but it is more rigorous. Could the current PR add some test cases?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4106184220
> I think a rewrite like this might be useful, but I think it can also hurt performance because of the join on grouping keys. So I think it needs to have a config value (off by default) or when enabled some benchmarks showing that it is better in large majority of the cases.
>
> I am also wondering if mostly for memory usage a `GroupsAccumulator` for distinct count / sum might give similar/more improvements.
@Dandandan Thank you for the explanation. It’s true that this would add a hash join, but if aggregation can be performed in parallel, there might be advantages in scenarios with two or more COUNT(DISTINCT) operations. I agree to run performance tests across multiple scenarios to evaluate the actual results.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4134601985
Perhaps the next PR could consider a simplification approach for injective functions, which might yield better performance gains. If we keep the original design, I personally don’t see the need to impose restrictions. We cannot guarantee that a rule is beneficial in all cases—cost-based selection should be integrated. Of course, this is just my personal suggestion.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4115456381
Hi @xiedeyantu update on the benchmarking. The new benchmarks use the **grouped aggregate only** (no outer `SUM` or other wrapper) so timings reflect the multi–`COUNT(DISTINCT …)` work directly. Note: Figures below are **warm-cache** repeats on one machine (relative A/B, not a CI guarantee). Also, these figures are from running each of the query 10 times and flipping the rewrite setting on and off using the grouped queries above (no outer `SUM`).
---
## Queries tested — original runs
### Scenario 1: three distincts (grouped)
```sql
SELECT l_suppkey,
COUNT(DISTINCT l_orderkey) AS d1,
COUNT(DISTINCT l_partkey) AS d2,
COUNT(DISTINCT l_linenumber) AS d3
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_suppkey;
```
### Scenario 2: CAST-based distincts (original)
```sql
SELECT l_suppkey,
COUNT(DISTINCT CAST(l_extendedprice AS INT)) AS cx,
COUNT(DISTINCT CAST(l_discount * 100 AS INT)) AS cy
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_suppkey;
```
### Scenario 3: `lower()` + string distinct (original, low-cardinality group key)
```sql
SELECT l_returnflag,
COUNT(DISTINCT lower(l_shipmode)) AS c1,
COUNT(DISTINCT l_shipinstruct) AS c2
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_returnflag;
```
---
## Rerun — reviewer-aligned changes
**CAST (Scenario 2):** both distincts use `CAST` on a **column** only (no `*` expression inside `CAST`):
```sql
SELECT l_suppkey,
COUNT(DISTINCT CAST(l_extendedprice AS INT)) AS cx,
COUNT(DISTINCT CAST(l_discount AS INT)) AS cy
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_suppkey;
```
**High cardinality (Scenario 3):** same aggregates as before, but **`GROUP BY l_suppkey`** instead of `l_returnflag`:
```sql
SELECT l_suppkey,
COUNT(DISTINCT lower(l_shipmode)) AS c1,
COUNT(DISTINCT l_shipinstruct) AS c2
FROM 'benchmarks/data/tpch_sf1/lineitem/*.parquet'
GROUP BY l_suppkey;
```
### Results (mean ± sample stddev, **n=10** per cell)
Measured with `/usr/bin/time -l` and `target/release/datafusion-cli` from the repo root. RSS is **maximum resident set size** from `time -l` (MB here are decimal megabytes: bytes / 10⁶). Queries are the **reviewer-aligned grouped aggregates above** (no outer `SUM`).
**wall (s) mean ± stddev**
| Scenario | Rewrite OFF | Rewrite ON | Outcome |
|----------|------------:|-----------:|---------|
| 1) three distincts | 0.2490 ± 0.0099 | 0.1000 ± 0.0000 | Lower mean wall time with ON |
| 2) cast-based (column-only `CAST`) | 0.1350 ± 0.0053 | 0.0700 ± 0.0000 | Lower mean wall time with ON |
| 3) lower()+string (`GROUP BY l_suppkey`) | 3.3920 ± 0.6496 | 0.0800 ± 0.0000 | Much lower mean wall time with ON |
**max RSS (MB) mean ± stddev**
| Scenario | Rewrite OFF | Rewrite ON | Outcome |
|----------|------------:|-----------:|---------|
| 1) three distincts | 590.03 ± 10.12 | 722.62 ± 12.70 | Higher mean peak RSS with ON |
| 2) cast-based (column-only `CAST`) | 331.83 ± 4.56 | 638.49 ± 8.29 | Higher mean peak RSS with ON |
| 3) lower()+string (`GROUP BY l_suppkey`) | 9755.13 ± 114.94 | 277.66 ± 1.71 | Much lower mean peak RSS with ON |
Row outputs matched OFF vs ON for every scenario (same grouped aggregate result). Wall time from `time` is low-precision on macOS (e.g. all runs may round to the same hundredth of a second, so stddev can be 0).
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on code in PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#discussion_r2969672012
##########
datafusion/optimizer/src/multi_distinct_count_rewrite.rs:
##########
@@ -0,0 +1,436 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Rewrites a single [`Aggregate`] that contains multiple `COUNT(DISTINCT ...)`
+//! into a join of smaller aggregates so each distinct is computed with one
+//! accumulator set, reducing peak memory for high-cardinality distincts.
+
+use std::sync::Arc;
+
+use crate::optimizer::ApplyOrder;
+use crate::{OptimizerConfig, OptimizerRule};
+
+use datafusion_common::{
+ Column, JoinConstraint, NullEquality, Result, tree_node::Transformed,
+};
+use datafusion_expr::builder::project;
+use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams, ScalarFunction};
+use datafusion_expr::logical_plan::{
+ Aggregate, Join, JoinType, LogicalPlan, SubqueryAlias,
+};
+use datafusion_expr::{Expr, col, lit, logical_plan::LogicalPlanBuilder};
+
+const MAX_DISTINCT_REWRITE_BRANCHES: usize = 8;
+
+/// Optimizer rule: multiple `COUNT(DISTINCT ...)` → join of per-distinct sub-aggregates.
+#[derive(Default, Debug)]
+pub struct MultiDistinctCountRewrite {}
+
+impl MultiDistinctCountRewrite {
+ /// Create a new rule instance.
+ pub fn new() -> Self {
+ Self {}
+ }
+
+ fn is_simple_count_distinct(
+ e: &Expr,
+ ) -> Option<(Arc<datafusion_expr::AggregateUDF>, Expr)> {
+ if let Expr::AggregateFunction(AggregateFunction { func, params }) = e {
+ let AggregateFunctionParams {
+ distinct,
+ args,
+ filter,
+ order_by,
+ ..
+ } = ¶ms;
+ if func.name().eq_ignore_ascii_case("count")
+ && *distinct
+ && args.len() == 1
+ && filter.is_none()
+ && order_by.is_empty()
+ {
+ let arg = args.first().cloned()?;
+ if Self::is_safe_distinct_arg(&arg) {
+ return Some((Arc::clone(func), arg));
+ }
+ }
+ }
+ None
+ }
+
+ fn is_safe_distinct_arg(e: &Expr) -> bool {
+ if e.is_volatile() {
+ return false;
+ }
+ match e {
+ Expr::Column(_) => true,
+ Expr::ScalarFunction(ScalarFunction { func, args }) => {
+ matches!(func.name().to_ascii_lowercase().as_str(), "lower" | "upper")
+ && args.len() == 1
+ && matches!(args.first(), Some(Expr::Column(_)))
+ }
+ Expr::Cast(cast) => matches!(cast.expr.as_ref(), Expr::Column(_)),
+ _ => false,
+ }
+ }
+
+ fn is_simple_group_expr(e: &Expr) -> bool {
+ matches!(e, Expr::Column(_))
+ }
+
+ fn contains_grouping_set(group_expr: &[Expr]) -> bool {
+ group_expr
+ .first()
+ .is_some_and(|e| matches!(e, Expr::GroupingSet(_)))
+ }
+}
+
+impl OptimizerRule for MultiDistinctCountRewrite {
+ fn name(&self) -> &str {
+ "multi_distinct_count_rewrite"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::BottomUp)
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
+ let LogicalPlan::Aggregate(Aggregate {
+ input,
+ aggr_expr,
+ schema,
+ group_expr,
+ ..
+ }) = plan
+ else {
+ return Ok(Transformed::no(plan));
+ };
+
+ if Self::contains_grouping_set(&group_expr) {
+ return Ok(Transformed::no(LogicalPlan::Aggregate(Aggregate::try_new(
+ input, group_expr, aggr_expr,
+ )?)));
+ }
+
+ if !group_expr.iter().all(Self::is_simple_group_expr) {
+ return Ok(Transformed::no(LogicalPlan::Aggregate(Aggregate::try_new(
+ input, group_expr, aggr_expr,
+ )?)));
+ }
+
+ let group_size = group_expr.len();
+ let mut distinct_list: Vec<(Expr, usize, Arc<datafusion_expr::AggregateUDF>)> =
+ vec![];
+ let mut other_list: Vec<(Expr, usize)> = vec![];
+
+ for (i, e) in aggr_expr.iter().enumerate() {
+ if let Some((func, arg)) = Self::is_simple_count_distinct(e) {
Review Comment:
Semantics are unchanged, distinct is on the inner expression; we added SQL tests for lower case-folding and CAST collision with two COUNT(DISTINCT …) so the rewrite matches vanilla aggregation.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on code in PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#discussion_r2994387765
##########
docs/source/user-guide/configs.md:
##########
@@ -144,6 +144,7 @@ The following configuration settings are available:
| datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown | true | When set to true, the optimizer will attempt to push down Aggregate dynamic filters into the file scan phase.
|
| datafusion.optimizer.enable_dynamic_filter_pushdown | true | When set to true attempts to push down dynamic filters generated by operators (TopK, Join & Aggregate) into the file scan phase. For example, for a query such as `SELECT * FROM t ORDER BY timestamp DESC LIMIT 10`, the optimizer will attempt to push down the current top 10 timestamps that the TopK operator references into the file scans. This means that if we already have 10 timestamps in the year 2025 any files that only have timestamps in the year 2024 can be skipped / pruned at various stages in the scan. The config will suppress `enable_join_dynamic_filter_pushdown`, `enable_topk_dynamic_filter_pushdown` & `enable_aggregate_dynamic_filter_pushdown` So if you disable `enable_topk_dynamic_filter_pushdown`, then enable `enable_dynamic_filter_pushdown`, the `enable_topk_dynamic_filter_pushdown` will be overridden.
|
| datafusion.optimizer.filter_null_join_keys | false | When set to true, the optimizer will insert filters before a join between a nullable and non-nullable column to filter out nulls on the nullable side. This filter can add additional overhead when the file format does not fully support predicate push down.
|
+| datafusion.optimizer.enable_multi_distinct_count_rewrite | false | When set to true, the optimizer may rewrite a single aggregate with multiple `COUNT(DISTINCT …)` (with `GROUP BY`) into joins of per-distinct sub-aggregates. This can reduce peak memory but adds join work; default off until benchmarks support enabling broadly.
|
Review Comment:
I have made small changes to this in my latest commit, Hope this helps.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4103001083
Hi @Dandandan — thanks for this work; the cross-join split for **multiple distinct aggregates with no `GROUP BY`** is a strong fit for workloads like ClickBench extended.
I’ve been working on a related but **different** pattern: **`GROUP BY` + several `COUNT(DISTINCT …)`** in the same aggregate (typical BI). In that situation, your rule is not quite applicable, because `MultiDistinctToCrossJoin` needs an **empty** `GROUP BY` and **all** aggregates to be distinct on different columns.
Here is a concrete example from our internal benchmark suite on an `ecommerce_orders` table:
```sql
SELECT
seller_name,
COUNT(*) as total_orders,
COUNT(DISTINCT delivery_city) as cities_served,
COUNT(DISTINCT state) as states_served,
SUM(CASE WHEN order_status = 'Completed' THEN 1 ELSE 0 END) as completed_orders,
SUM(CASE WHEN order_status = 'Cancelled' THEN 1 ELSE 0 END) as cancelled_orders,
ROUND(100.0 * SUM(CASE WHEN order_status = 'Completed' THEN 1 ELSE 0 END) / COUNT(*), 2) as success_rate
FROM orders_data
GROUP BY seller_name
HAVING COUNT(*) > 100
ORDER BY total_orders DESC
LIMIT 100;
```
This is **not** “global” multi-distinct: it’s **per `seller_name`**, with **multiple `COUNT(DISTINCT …)`** plus other aggregates. That’s the class my optimizer rule (`MultiDistinctCountRewrite`) targets — rewriting the **`COUNT(DISTINCT …)`** pieces into **joinable sub-aggregates aligned on the same `GROUP BY` keys**, with correct `NULL` handling where needed. For this table with 20m rows, on my M4 machine the times are 0.42s vs ~17s for the implementation in #20940
In my opinion, they’re **complementary**: different predicates, different plans, and they can **coexist** in the optimizer pipeline (we’d want to sanity-check rule order so we don’t double-rewrite the same node).
Happy to align naming, tests, and placement with you and the maintainers.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4105715551
I think a rewrite like this might be useful, but I think it can also hurt performance because of the join on grouping keys.
So I think it needs to have a config value (off by default) or when enabled some benchmarks showing that it is better in large majority of the cases.
I am also wondering if mostly for memory usage a `GroupsAccumulator` for distinct count / sum might give similar/more improvements.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4110389221
Thank you for the detailed test! But I have a few questions.
1. Why add an extra
"SUM" outside?
2. I think it’s better for
"CAST" not to involve expression calculations.
3. It seems the third case is a low-cardinality test — could you show a high-cardinality scenario?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4102785883
Hi @ydgandhi !
This does seem the same to me as https://github.com/apache/datafusion/pull/20940 ?
Can you maybe help review that one?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4133967253
You’re right that injectivity isn’t the right lens here. This rewrite doesn’t assume (f) is injective. We intentionally rewrite to GROUP BY g, f(a) (not GROUP BY g, a) so cases like lower('Abc') and lower('aBC') correctly collapse to one distinct bucket.
The current lower/upper/cast allowlist is a conservative scope guard, not a correctness requirement: it keeps the first version (default off) limited to a small set of well-understood, deterministic expression shapes we have tests/benchmarks for, while we gather more performance + edge-case coverage (e.g. other scalar functions, float/NaN semantics, dictionary types, etc.). We already block volatile expressions, but “any non-volatile expr” still opens a large surface area.
> I have one minor question. In the scenario of rewriting "g, count(distinct lower(a))" to "group by g, lower(a)", why restrict it to only "lower", "upper", and "cast"? I don’t think any restriction is needed here. If the intention is to express an injective function, then "lower" is not one. If it is an injective function, the rewrite could simply be "group by g, a". Could you please clarify the design rationale behind this?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4133910311
I have one minor question. In the scenario of rewriting
"g, count(distinct lower(a))" to
"group by g, lower(a)", why restrict it to only
"lower",
"upper", and
"cast"? I don’t think any restriction is needed here. If the intention is to express an injective function, then
"lower" is not one. If it is an injective function, the rewrite could simply be
"group by g, a". Could you please clarify the design rationale behind this?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #21088: perf: default multi COUNT(DISTINCT) logical optimizer rewrite
URL: https://github.com/apache/datafusion/pull/21088
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "ydgandhi (via GitHub)" <gi...@apache.org>.
ydgandhi commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4103479506
Thanks for the review. I asked cursor to add a few tests
---
## Tests for `MultiDistinctCountRewrite` (what they cover)
Optimizer unit tests — `datafusion/optimizer/src/multi_distinct_count_rewrite.rs`
| Test | What it asserts |
|------|-----------------|
| `rewrites_two_count_distinct` | `GROUP BY a` + `COUNT(DISTINCT b)`, `COUNT(DISTINCT c)` → inner joins, per-branch null filters on `b`/`c`, `mdc_base` + two `mdc_d` aliases. |
| `rewrites_global_three_count_distinct` | No `GROUP BY`, three `COUNT(DISTINCT …)` → cross/inner join rewrite; **no** `mdc_base` (global-only path). |
| `rewrites_two_count_distinct_with_non_distinct_count` | Grouped BI-style: two distincts + `COUNT(a)` → join rewrite with **`mdc_base`** holding the non-distinct agg. |
| `does_not_rewrite_two_count_distinct_same_column` | Two `COUNT(DISTINCT b)` with different aliases → **no** rewrite (duplicate distinct key). |
| `does_not_rewrite_single_count_distinct` | Only one `COUNT(DISTINCT …)` → **no** rewrite (rule needs ≥2 distincts). |
| `rewrites_three_count_distinct_grouped` | Three grouped `COUNT(DISTINCT …)` on `b`, `c`, `a` → **two** inner joins + `mdc_base`. |
| `rewrites_interleaved_non_distinct_between_distincts` | Order `COUNT(DISTINCT b)`, `COUNT(a)`, `COUNT(DISTINCT c)` → rewrite + `mdc_base` for the middle non-distinct agg (projection order / interleaving). |
| `rewrites_count_distinct_on_cast_exprs` | `COUNT(DISTINCT CAST(b AS Int64))`, same for `c` → rewrite + null filters on the **cast** expressions. |
| `does_not_rewrite_grouping_sets_multi_distinct` | `GROUPING SETS` aggregate with two `COUNT(DISTINCT …)` → **no** rewrite (rule bails on grouping sets). |
| `does_not_rewrite_mixed_agg` | `COUNT(DISTINCT b)` + `COUNT(c)` → **no** rewrite (only **one** `COUNT(DISTINCT …)`; rule requires at least two). |
SQL integration — `datafusion/core/tests/sql/aggregates/multi_distinct_count_rewrite.rs`
| Test | What it asserts |
|------|-----------------|
| `multi_count_distinct_matches_expected_with_nulls` | End-to-end grouped two `COUNT(DISTINCT …)` with **NULLs** in distinct columns; exact sorted batch string vs expected counts. |
| `multi_count_distinct_with_count_star_matches_expected` | `COUNT(*)` plus two `COUNT(DISTINCT …)` per group (BI-style); exact result table. |
| `multi_count_distinct_two_group_keys_matches_expected` | **`GROUP BY g1, g2`** + two distincts; verifies joins line up on **all** group keys and numerics match. |
---
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#issuecomment-4109532126
Yes, I can see it would improve a lot of queries
> Thanks @Dandandan @xiedeyantu for the discussion.
>
> **Join cost / default behavior:** Agreed that joining on grouping keys isn’t free. We’ve added **`datafusion.optimizer.enable_multi_distinct_count_rewrite`** with default **`false`**, so the rewrite is **opt-in** until we have benchmarks that justify turning it on by default. Sessions can enable it explicitly when they want to try the plan shape.
>
> **Benchmarks:** We’re aligned that we should measure **latency and memory** vs baseline across scenarios (e.g. multiple `COUNT(DISTINCT …)` with `GROUP BY`, varying group-key cardinality vs distinct cardinality). We’ll follow up with numbers on the PR.
>
> **GroupsAccumulator / execution-layer improvements:** That work is complementary: better distinct accumulators help **how** each aggregate runs; this rule changes **logical plan shape** when several large distincts share one aggregate. Both can coexist on the roadmap.
>
> **Tests:** We added optimizer + SQL integration coverage (including cases with `lower`/`CAST` inside `COUNT(DISTINCT …)` and a test that the rule is a no-op when the config is off). Happy to iterate on naming and placement with maintainers.
>
> > > I think a rewrite like this might be useful, but I think it can also hurt performance because of the join on grouping keys. So I think it needs to have a config value (off by default) or when enabled some benchmarks showing that it is better in large majority of the cases.
> > > I am also wondering if mostly for memory usage a `GroupsAccumulator` for distinct count / sum might give similar/more improvements.
> >
> >
> > @Dandandan Thank you for the explanation. It’s true that this would add a hash join, but if aggregation can be performed in parallel, there might be advantages in scenarios with two or more COUNT(DISTINCT) operations. I agree to run performance tests across multiple scenarios to evaluate the actual results.
Sounds good - agreed.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on code in PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#discussion_r2994064374
##########
docs/source/user-guide/configs.md:
##########
@@ -144,6 +144,7 @@ The following configuration settings are available:
| datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown | true | When set to true, the optimizer will attempt to push down Aggregate dynamic filters into the file scan phase.
|
| datafusion.optimizer.enable_dynamic_filter_pushdown | true | When set to true attempts to push down dynamic filters generated by operators (TopK, Join & Aggregate) into the file scan phase. For example, for a query such as `SELECT * FROM t ORDER BY timestamp DESC LIMIT 10`, the optimizer will attempt to push down the current top 10 timestamps that the TopK operator references into the file scans. This means that if we already have 10 timestamps in the year 2025 any files that only have timestamps in the year 2024 can be skipped / pruned at various stages in the scan. The config will suppress `enable_join_dynamic_filter_pushdown`, `enable_topk_dynamic_filter_pushdown` & `enable_aggregate_dynamic_filter_pushdown` So if you disable `enable_topk_dynamic_filter_pushdown`, then enable `enable_dynamic_filter_pushdown`, the `enable_topk_dynamic_filter_pushdown` will be overridden.
|
| datafusion.optimizer.filter_null_join_keys | false | When set to true, the optimizer will insert filters before a join between a nullable and non-nullable column to filter out nulls on the nullable side. This filter can add additional overhead when the file format does not fully support predicate push down.
|
+| datafusion.optimizer.enable_multi_distinct_count_rewrite | false | When set to true, the optimizer may rewrite a single aggregate with multiple `COUNT(DISTINCT …)` (with `GROUP BY`) into joins of per-distinct sub-aggregates. This can reduce peak memory but adds join work; default off until benchmarks support enabling broadly.
|
Review Comment:
"This can reduce peak memory"
Could this description be further optimized? Aggregation parallelism does not always reduce peak memory usage.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: default multi COUNT(DISTINCT) logical optimizer rewrite [datafusion]
Posted by "xiedeyantu (via GitHub)" <gi...@apache.org>.
xiedeyantu commented on code in PR #21088:
URL: https://github.com/apache/datafusion/pull/21088#discussion_r2969599296
##########
datafusion/optimizer/src/multi_distinct_count_rewrite.rs:
##########
@@ -0,0 +1,436 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Rewrites a single [`Aggregate`] that contains multiple `COUNT(DISTINCT ...)`
+//! into a join of smaller aggregates so each distinct is computed with one
+//! accumulator set, reducing peak memory for high-cardinality distincts.
+
+use std::sync::Arc;
+
+use crate::optimizer::ApplyOrder;
+use crate::{OptimizerConfig, OptimizerRule};
+
+use datafusion_common::{
+ Column, JoinConstraint, NullEquality, Result, tree_node::Transformed,
+};
+use datafusion_expr::builder::project;
+use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams, ScalarFunction};
+use datafusion_expr::logical_plan::{
+ Aggregate, Join, JoinType, LogicalPlan, SubqueryAlias,
+};
+use datafusion_expr::{Expr, col, lit, logical_plan::LogicalPlanBuilder};
+
+const MAX_DISTINCT_REWRITE_BRANCHES: usize = 8;
+
+/// Optimizer rule: multiple `COUNT(DISTINCT ...)` → join of per-distinct sub-aggregates.
+#[derive(Default, Debug)]
+pub struct MultiDistinctCountRewrite {}
+
+impl MultiDistinctCountRewrite {
+ /// Create a new rule instance.
+ pub fn new() -> Self {
+ Self {}
+ }
+
+ fn is_simple_count_distinct(
+ e: &Expr,
+ ) -> Option<(Arc<datafusion_expr::AggregateUDF>, Expr)> {
+ if let Expr::AggregateFunction(AggregateFunction { func, params }) = e {
+ let AggregateFunctionParams {
+ distinct,
+ args,
+ filter,
+ order_by,
+ ..
+ } = ¶ms;
+ if func.name().eq_ignore_ascii_case("count")
+ && *distinct
+ && args.len() == 1
+ && filter.is_none()
+ && order_by.is_empty()
+ {
+ let arg = args.first().cloned()?;
+ if Self::is_safe_distinct_arg(&arg) {
+ return Some((Arc::clone(func), arg));
+ }
+ }
+ }
+ None
+ }
+
+ fn is_safe_distinct_arg(e: &Expr) -> bool {
+ if e.is_volatile() {
+ return false;
+ }
+ match e {
+ Expr::Column(_) => true,
+ Expr::ScalarFunction(ScalarFunction { func, args }) => {
+ matches!(func.name().to_ascii_lowercase().as_str(), "lower" | "upper")
+ && args.len() == 1
+ && matches!(args.first(), Some(Expr::Column(_)))
+ }
+ Expr::Cast(cast) => matches!(cast.expr.as_ref(), Expr::Column(_)),
+ _ => false,
+ }
+ }
+
+ fn is_simple_group_expr(e: &Expr) -> bool {
+ matches!(e, Expr::Column(_))
+ }
+
+ fn contains_grouping_set(group_expr: &[Expr]) -> bool {
+ group_expr
+ .first()
+ .is_some_and(|e| matches!(e, Expr::GroupingSet(_)))
+ }
+}
+
+impl OptimizerRule for MultiDistinctCountRewrite {
+ fn name(&self) -> &str {
+ "multi_distinct_count_rewrite"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::BottomUp)
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
+ let LogicalPlan::Aggregate(Aggregate {
+ input,
+ aggr_expr,
+ schema,
+ group_expr,
+ ..
+ }) = plan
+ else {
+ return Ok(Transformed::no(plan));
+ };
+
+ if Self::contains_grouping_set(&group_expr) {
+ return Ok(Transformed::no(LogicalPlan::Aggregate(Aggregate::try_new(
+ input, group_expr, aggr_expr,
+ )?)));
+ }
+
+ if !group_expr.iter().all(Self::is_simple_group_expr) {
+ return Ok(Transformed::no(LogicalPlan::Aggregate(Aggregate::try_new(
+ input, group_expr, aggr_expr,
+ )?)));
+ }
+
+ let group_size = group_expr.len();
+ let mut distinct_list: Vec<(Expr, usize, Arc<datafusion_expr::AggregateUDF>)> =
+ vec![];
+ let mut other_list: Vec<(Expr, usize)> = vec![];
+
+ for (i, e) in aggr_expr.iter().enumerate() {
+ if let Some((func, arg)) = Self::is_simple_count_distinct(e) {
Review Comment:
Could you simulate a query like
SELECT g, COUNT(DISTINCT lower(b)), ... GROUP BY g
where the values of b are 'Abc' and 'aBC'?
I understand that cast presents a similar issue. For instance, when casting a double to an int, the original double values might be 1.2 and 1.3.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org