Repository: apache/datafusion
Pull request #15367: FIX : some benchmarks are failing
Author: login: getChan
Labels: name: core
description: Core DataFusion crate
Created: 2025-03-23T16:24:08Z | Merged: 2025-03-31T18:28:33Z

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->

- Closes #15213 .

## Rationale for this change

It is not certain, but it seems that plan creation and `collect()` should share the same runtime.. It is presumed that the issue occurred because RepartitionExec lazily polls within the runtime.
- https://github.com/apache/datafusion/blob/main/datafusion/physical-plan/src/repartition/mod.rs#L165-L177
- https://github.com/apache/datafusion/pull/10009

I will add more details if I find anything additional.

<!--
 Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
 Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
-->

## What changes are included in this PR?
- Move the `Runtime::new()` into`bench_function`
- Ensure that plan creation (`ctx.sql()`) and `collect()` share the same runtime.
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->

## Are these changes tested?
yes. below test are succeded
- `cargo bench -p datafusion --bench topk_aggregate`
- `cargo bench -p datafusion --bench distinct_query_sql`
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
3. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->

## Are there any user-facing changes?
No.
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->


Linked issue #123: Implement physical plan for EXISTS subquery
*Note*: migrated from original JIRA: https://issues.apache.org/jira/browse/ARROW-10819

The TPC-H queries include use of the EXISTS which is used to test for the existence of any record in a subquery. For example:

```sql
and *exists* (
    select
        *
    from
        lineitem
    where
        l_orderkey = o_orderkey
        and l_commitdate < l_receiptdate
)
```
Linked issue comments: user: login: alamb
created_at: 2021-04-26T13:23:43Z
body: Comment from Andy Grove(andygrove) @ 2020-12-31T19:42:45.132+0000:
<pre>The example given here is a correlated subquery that can be translated into a join.

Here is a random stackoverflow discussion on this for reference (I have not reviewed it)

https://stackoverflow.com/questions/1772609/procedurally-transform-subquery-into-join</pre>
author_association: CONTRIBUTOR
user: login: alamb
created_at: 2022-10-21T19:21:30Z
body: In case anyone is curious -- we support correlated versions of these queries (via a join) but if there is no correlation (not super useful) we do not

```sql
❯ create table foo as select * from (values (1), (2), (NULL)) as sql
;
0 rows in set. Query took 0.022 seconds.
3 rows in set. Query took 0.007 seconds.
❯ create table bar as select * from (values (1), (2), (NULL)) as sql;
0 rows in set. Query took 0.000 seconds.
❯ select * from foo where exists (select column1 from bar);
NotImplemented("Physical plan does not support logical expression EXISTS (<subquery>)")
❯ select * from foo where exists (select column1 from bar where foo.column1 = bar.column1);
+---------+
| column1 |
+---------+
| 2       |
| 1       |
+---------+
```
author_association: CONTRIBUTOR
user: login: logan-keede
created_at: 2025-02-09T19:53:56Z
body: 
> In case anyone is curious -- we support correlated versions of these queries (via a join) but if there is no correlation (not super useful) we do not
> 
> ❯ create table foo as select * from (values (1), (2), (NULL)) as sql
> ;
> 0 rows in set. Query took 0.022 seconds.
> 3 rows in set. Query took 0.007 seconds.
> ❯ create table bar as select * from (values (1), (2), (NULL)) as sql;
> 0 rows in set. Query took 0.000 seconds.
> ❯ select * from foo where exists (select column1 from bar);
> NotImplemented("Physical plan does not support logical expression EXISTS (<subquery>)")
> ❯ select * from foo where exists (select column1 from bar where foo.column1 = bar.column1);
> +---------+
> | column1 |
> +---------+
> | 2       |
> | 1       |
> +---------+

```sql

> explain select * from foo where exists (select column1 from bar);
+---------------+-----------------------------------------------------+
| plan_type     | plan                                                |
+---------------+-----------------------------------------------------+
| logical_plan  | LeftSemi Join:                                      |
|               |   TableScan: foo projection=[column1]               |
|               |   SubqueryAlias: __correlated_sq_1                  |
|               |     TableScan: bar projection=[]                    |
| physical_plan | NestedLoopJoinExec: join_type=RightSemi             |
|               |   DataSourceExec: partitions=1, partition_sizes=[1] |
|               |   DataSourceExec: partitions=1, partition_sizes=[1] |
|               |                                                     |
+---------------+-----------------------------------------------------+
2 row(s) fetched. 
Elapsed 0.007 seconds.

> select * from foo where exists (select column1 from bar);
+---------+
| column1 |
+---------+
| 1       |
| 2       |
| NULL    |
+---------+
3 row(s) fetched. 
Elapsed 0.006 seconds.
```


I get this on `datafusion 45.0.0`, perhaps this issue can be closed or at least updated if there is a case where we still don't support EXISTS in physical plan.

author_association: CONTRIBUTOR
user: login: alamb
created_at: 2025-02-10T15:04:31Z
body: Looks good -- thanks for checking @logan-keede  -- we can open a new issue if we find another hole
author_association: CONTRIBUTOR

Diff:
diff --git a/datafusion/core/benches/distinct_query_sql.rs b/datafusion/core/benches/distinct_query_sql.rs
index c242798a56f00..ccc6a0e74652c 100644
--- a/datafusion/core/benches/distinct_query_sql.rs
+++ b/datafusion/core/benches/distinct_query_sql.rs
@@ -118,12 +118,9 @@ async fn distinct_with_limit(
     Ok(())
 }
 
-fn run(plan: Arc<dyn ExecutionPlan>, ctx: Arc<TaskContext>) {
-    let rt = Runtime::new().unwrap();
-    criterion::black_box(
-        rt.block_on(async { distinct_with_limit(plan.clone(), ctx.clone()).await }),
-    )
-    .unwrap();
+fn run(rt: &Runtime, plan: Arc<dyn ExecutionPlan>, ctx: Arc<TaskContext>) {
+    criterion::black_box(rt.block_on(distinct_with_limit(plan.clone(), ctx.clone())))
+        .unwrap();
 }
 
 pub async fn create_context_sampled_data(
@@ -145,58 +142,47 @@ pub async fn create_context_sampled_data(
 
 fn criterion_benchmark_limited_distinct_sampled(c: &mut Criterion) {
     let rt = Runtime::new().unwrap();
-
     let limit = 10;
     let partitions = 100;
     let samples = 100_000;
     let sql =
         format!("select DISTINCT trace_id from traces group by trace_id limit {limit};");
-
-    let distinct_trace_id_100_partitions_100_000_samples_limit_100 = rt.block_on(async {
-        create_context_sampled_data(sql.as_str(), partitions, samples)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
-        |b| b.iter(|| run(distinct_trace_id_100_partitions_100_000_samples_limit_100.0.clone(),
-                                   distinct_trace_id_100_partitions_100_000_samples_limit_100.1.clone())),
+        |b| b.iter(|| {
+            let (plan, ctx) = rt.block_on(
+                create_context_sampled_data(sql.as_str(), partitions, samples)
+            ).unwrap();
+            run(&rt, plan.clone(), ctx.clone())
+        }),
     );
 
     let partitions = 10;
     let samples = 1_000_000;
     let sql =
         format!("select DISTINCT trace_id from traces group by trace_id limit {limit};");
-
-    let distinct_trace_id_10_partitions_1_000_000_samples_limit_10 = rt.block_on(async {
-        create_context_sampled_data(sql.as_str(), partitions, samples)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
-        |b| b.iter(|| run(distinct_trace_id_10_partitions_1_000_000_samples_limit_10.0.clone(),
-                                   distinct_trace_id_10_partitions_1_000_000_samples_limit_10.1.clone())),
+        |b| b.iter(|| {
+            let (plan, ctx) = rt.block_on(
+                create_context_sampled_data(sql.as_str(), partitions, samples)
+            ).unwrap();
+            run(&rt, plan.clone(), ctx.clone())
+        }),
     );
 
     let partitions = 1;
     let samples = 10_000_000;
     let sql =
         format!("select DISTINCT trace_id from traces group by trace_id limit {limit};");
-
-    let rt = Runtime::new().unwrap();
-    let distinct_trace_id_1_partition_10_000_000_samples_limit_10 = rt.block_on(async {
-        create_context_sampled_data(sql.as_str(), partitions, samples)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
-        |b| b.iter(|| run(distinct_trace_id_1_partition_10_000_000_samples_limit_10.0.clone(),
-                                   distinct_trace_id_1_partition_10_000_000_samples_limit_10.1.clone())),
+        |b| b.iter(|| {
+            let (plan, ctx) = rt.block_on(
+                create_context_sampled_data(sql.as_str(), partitions, samples)
+            ).unwrap();
+            run(&rt, plan.clone(), ctx.clone())
+        }),
     );
 }
 
diff --git a/datafusion/core/benches/topk_aggregate.rs b/datafusion/core/benches/topk_aggregate.rs
index 922cbd2b42292..777d586b344c4 100644
--- a/datafusion/core/benches/topk_aggregate.rs
+++ b/datafusion/core/benches/topk_aggregate.rs
@@ -55,8 +55,7 @@ async fn create_context(
     Ok((physical_plan, ctx.task_ctx()))
 }
 
-fn run(plan: Arc<dyn ExecutionPlan>, ctx: Arc<TaskContext>, asc: bool) {
-    let rt = Runtime::new().unwrap();
+fn run(rt: &Runtime, plan: Arc<dyn ExecutionPlan>, ctx: Arc<TaskContext>, asc: bool) {
     criterion::black_box(
         rt.block_on(async { aggregate(plan.clone(), ctx.clone(), asc).await }),
     )
@@ -99,40 +98,37 @@ async fn aggregate(
 }
 
 fn criterion_benchmark(c: &mut Criterion) {
+    let rt = Runtime::new().unwrap();
     let limit = 10;
     let partitions = 10;
     let samples = 1_000_000;
 
-    let rt = Runtime::new().unwrap();
-    let topk_real = rt.block_on(async {
-        create_context(limit, partitions, samples, false, true)
-            .await
-            .unwrap()
-    });
-    let topk_asc = rt.block_on(async {
-        create_context(limit, partitions, samples, true, true)
-            .await
-            .unwrap()
-    });
-    let real = rt.block_on(async {
-        create_context(limit, partitions, samples, false, false)
-            .await
-            .unwrap()
-    });
-    let asc = rt.block_on(async {
-        create_context(limit, partitions, samples, true, false)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("aggregate {} time-series rows", partitions * samples).as_str(),
-        |b| b.iter(|| run(real.0.clone(), real.1.clone(), false)),
+        |b| {
+            b.iter(|| {
+                let real = rt.block_on(async {
+                    create_context(limit, partitions, samples, false, false)
+                        .await
+                        .unwrap()
+                });
+                run(&rt, real.0.clone(), real.1.clone(), false)
+            })
+        },
     );
 
     c.bench_function(
         format!("aggregate {} worst-case rows", partitions * samples).as_str(),
-        |b| b.iter(|| run(asc.0.clone(), asc.1.clone(), true)),
+        |b| {
+            b.iter(|| {
+                let asc = rt.block_on(async {
+                    create_context(limit, partitions, samples, true, false)
+                        .await
+                        .unwrap()
+                });
+                run(&rt, asc.0.clone(), asc.1.clone(), true)
+            })
+        },
     );
 
     c.bench_function(
@@ -141,7 +137,16 @@ fn criterion_benchmark(c: &mut Criterion) {
             partitions * samples
         )
         .as_str(),
-        |b| b.iter(|| run(topk_real.0.clone(), topk_real.1.clone(), false)),
+        |b| {
+            b.iter(|| {
+                let topk_real = rt.block_on(async {
+                    create_context(limit, partitions, samples, false, true)
+                        .await
+                        .unwrap()
+                });
+                run(&rt, topk_real.0.clone(), topk_real.1.clone(), false)
+            })
+        },
     );
 
     c.bench_function(
@@ -150,7 +155,16 @@ fn criterion_benchmark(c: &mut Criterion) {
             partitions * samples
         )
         .as_str(),
-        |b| b.iter(|| run(topk_asc.0.clone(), topk_asc.1.clone(), true)),
+        |b| {
+            b.iter(|| {
+                let topk_asc = rt.block_on(async {
+                    create_context(limit, partitions, samples, true, true)
+                        .await
+                        .unwrap()
+                });
+                run(&rt, topk_asc.0.clone(), topk_asc.1.clone(), true)
+            })
+        },
     );
 }
 


Reviews:
user: login: alamb
body: Thanks for this @getChan 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T21:00:02Z
user: login: getChan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T16:17:19Z
user: login: Omega359
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T17:12:55Z
user: login: alamb
body: Thank you @getChan and @Omega359 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T18:27:41Z

Review comments:
diff_hunk: @@ -144,59 +141,50 @@ pub async fn create_context_sampled_data(
 }
 
 fn criterion_benchmark_limited_distinct_sampled(c: &mut Criterion) {
-    let rt = Runtime::new().unwrap();
-
     let limit = 10;
     let partitions = 100;
     let samples = 100_000;
     let sql =
         format!("select DISTINCT trace_id from traces group by trace_id limit {limit};");
-
-    let distinct_trace_id_100_partitions_100_000_samples_limit_100 = rt.block_on(async {
-        create_context_sampled_data(sql.as_str(), partitions, samples)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
-        |b| b.iter(|| run(distinct_trace_id_100_partitions_100_000_samples_limit_100.0.clone(),
-                                   distinct_trace_id_100_partitions_100_000_samples_limit_100.1.clone())),
+        |b| b.iter(|| {
+            let rt = Runtime::new().unwrap();
path: datafusion/core/benches/distinct_query_sql.rs
user: login: alamb
body: I think this means that the benchmark will include the time to create each tokio runtime (with a bunch of threads, etc)

To avoid this I think you can create the runtime once, and then use it for each tieration:

```rist
            let rt = Runtime::new().unwrap();
    c.bench_function(
        format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
        |b| b.iter(|| {
```
created_at: 2025-03-25T20:59:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -144,59 +141,50 @@ pub async fn create_context_sampled_data(
 }
 
 fn criterion_benchmark_limited_distinct_sampled(c: &mut Criterion) {
-    let rt = Runtime::new().unwrap();
-
     let limit = 10;
     let partitions = 100;
     let samples = 100_000;
     let sql =
         format!("select DISTINCT trace_id from traces group by trace_id limit {limit};");
-
-    let distinct_trace_id_100_partitions_100_000_samples_limit_100 = rt.block_on(async {
-        create_context_sampled_data(sql.as_str(), partitions, samples)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
-        |b| b.iter(|| run(distinct_trace_id_100_partitions_100_000_samples_limit_100.0.clone(),
-                                   distinct_trace_id_100_partitions_100_000_samples_limit_100.1.clone())),
+        |b| b.iter(|| {
+            let rt = Runtime::new().unwrap();
path: datafusion/core/benches/distinct_query_sql.rs
user: login: getChan
body: Got it. Changed the runtime to be shared between loops.

However, it seems that other benchmark codes still have runtime creation within the iteration.
created_at: 2025-03-26T16:17:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -144,59 +141,50 @@ pub async fn create_context_sampled_data(
 }
 
 fn criterion_benchmark_limited_distinct_sampled(c: &mut Criterion) {
-    let rt = Runtime::new().unwrap();
-
     let limit = 10;
     let partitions = 100;
     let samples = 100_000;
     let sql =
         format!("select DISTINCT trace_id from traces group by trace_id limit {limit};");
-
-    let distinct_trace_id_100_partitions_100_000_samples_limit_100 = rt.block_on(async {
-        create_context_sampled_data(sql.as_str(), partitions, samples)
-            .await
-            .unwrap()
-    });
-
     c.bench_function(
         format!("distinct query with {} partitions and {} samples per partition with limit {}", partitions, samples, limit).as_str(),
-        |b| b.iter(|| run(distinct_trace_id_100_partitions_100_000_samples_limit_100.0.clone(),
-                                   distinct_trace_id_100_partitions_100_000_samples_limit_100.1.clone())),
+        |b| b.iter(|| {
+            let rt = Runtime::new().unwrap();
path: datafusion/core/benches/distinct_query_sql.rs
user: login: Omega359
body: Filed as https://github.com/apache/datafusion/issues/15507
created_at: 2025-03-31T17:12:54Z
author_association: CONTRIBUTOR

PR comments:
user: login: 2010YOUY01
created_at: 2025-03-30T08:56:13Z
body: Thank you for the fix! I noticed there are two other tests panicked on the same line of source code, is this fix still applicable?
1. Running sqllogictest with sqlite check: https://github.com/apache/datafusion/actions/runs/14152508679/job/39647603410
2. `custom_datasource` example: https://github.com/apache/datafusion/actions/runs/14152369014/job/39647355093
author_association: CONTRIBUTOR
user: login: getChan
created_at: 2025-03-30T12:00:17Z
body: > Thank you for the fix! I noticed there are two other tests panicked on the same line of source code, is this fix still applicable?
> 
> 1. Running sqllogictest with sqlite check: https://github.com/apache/datafusion/actions/runs/14152508679/job/39647603410
> 2. `custom_datasource` example: https://github.com/apache/datafusion/actions/runs/14152369014/job/39647355093

@2010YOUY01 
This PR is to fix the bench test code. The `partition not used yet` error can occur depending on the execution environment of each test (it is not likely an error in repartitionExec). Therefore, this PR does not solve the issue you mentioned.

author_association: CONTRIBUTOR
user: login: alamb
created_at: 2025-03-31T18:28:41Z
body: Let's get this in and keep iterating on the benchmarks
author_association: CONTRIBUTOR