Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2981051106 > > I agree the existing benchmarks are not perfect > > Not trying to criticize the benchmarks, just saying I'm having trouble using it as a guide. @alamb Note that by default the current branch is using the local counter method @zhuqi-lucas implemented for `YieldStream`. Performance should be identical to `main` in other words. I totally get it -- we'll keep improving the benchmarks / hopefully making them easier to use / more stable over 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: [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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2981029248 🤖: Benchmark completed Details ``` Comparing HEAD and task_budget Benchmark clickbench_extended.json ┏━━┳━┳━┳━━┓ ┃ Query┃HEAD ┃ task_budget ┃ Change ┃ ┡━━╇━╇━╇━━┩ │ QQuery 0 │ 1929.02 ms │ 1971.93 ms │no change │ │ QQuery 1 │ 704.63 ms │ 718.11 ms │no change │ │ QQuery 2 │ 1412.62 ms │ 1368.78 ms │no change │ │ QQuery 3 │ 647.44 ms │ 680.23 ms │ 1.05x slower │ │ QQuery 4 │ 1400.23 ms │ 1358.38 ms │no change │ │ QQuery 5 │ 14995.43 ms │ 15226.83 ms │no change │ │ QQuery 6 │ 1985.14 ms │ 2087.35 ms │ 1.05x slower │ │ QQuery 7 │ 1989.35 ms │ 1969.02 ms │no change │ │ QQuery 8 │ 785.58 ms │ 786.62 ms │no change │ └──┴─┴─┴──┘ ┏┳┓ ┃ Benchmark Summary ┃┃ ┡╇┩ │ Total Time (HEAD) │ 25849.44ms │ │ Total Time (task_budget) │ 26167.24ms │ │ Average Time (HEAD)│ 2872.16ms │ │ Average Time (task_budget) │ 2907.47ms │ │ Queries Faster │ 0 │ │ Queries Slower │ 2 │ │ Queries with No Change │ 7 │ │ Queries with Failure │ 0 │ └┴┘ Benchmark clickbench_partitioned.json ┏━━┳━┳━┳━━━┓ ┃ Query┃HEAD ┃ task_budget ┃Change ┃ ┡━━╇━╇━╇━━━┩ │ QQuery 0 │15.70 ms │15.32 ms │ no change │ │ QQuery 1 │33.40 ms │32.29 ms │ no change │ │ QQuery 2 │78.71 ms │81.60 ms │ no change │ │ QQuery 3 │94.83 ms │96.30 ms │ no change │ │ QQuery 4 │ 590.16 ms │ 586.19 ms │ no change │ │ QQuery 5 │ 839.78 ms │ 884.55 ms │ 1.05x slower │ │ QQuery 6 │23.62 ms │24.22 ms │ no change │ │ QQuery 7 │35.97 ms │35.71 ms │ no change │ │ QQuery 8 │ 845.86 ms │ 836.34 ms │ no change │ │ QQuery 9 │ 1155.51 ms │ 1124.37 ms │ no change │ │ QQuery 10│ 255.94 ms │ 258.01 ms │ no change │ │ QQuery 11│ 283.03 ms │ 285.17 ms │ no change │ │ QQuery 12│ 861.95 ms │ 884.44 ms │ no change │ │ QQuery 13│ 1256.74 ms │ 1274.54 ms │ no change │ │ QQuery 14│ 801.84 ms │ 819.99 ms │ no change │ │ QQuery 15│ 786.28 ms │ 775.94 ms │ no change │ │ QQuery 16│ 1616.39 ms │ 1622.78 ms │ no change │ │ QQuery 17│ 1591.05 ms │ 1593.83 ms │ no change │ │ QQuery 18│ 2857.24 ms │ 2902.54 ms │ no change │ │ QQuery 19│89.43 ms │84.53 ms │ +1.06x faster │ │ QQuery 20│ 1125.85 ms │ 1171.56 ms │ no change │ │ QQuery 21│ 1266.95 ms │ 1307.77 ms │ no change │ │ QQuery 22│ 2089.10 ms │ 2189.53 ms │ no change │ │ QQuery 23│ 7574.01 ms │ 7914.18 ms │ no change │ │ QQuery 24│ 444.06 ms │ 466.35 ms │ 1.05x slower │ │ QQuery 25│ 378.44 ms │ 397.31 ms │ no change │ │ QQuery 26│ 521.01 ms │ 525.65 ms │ no change │ │ QQuery 27│ 1503.49 ms │ 1548.92 ms │ no change │ │ QQuery 28│ 11900.61 ms │ 11913.35 ms │ no change │ │ QQuery 29│ 534.80 ms │ 540.04 ms │ no change │ │ QQuery 30│ 764.78 ms │ 790.34 ms │ no change │ │ QQuery 31│ 812.47 ms │ 825.30 ms │ no change │ │ QQuery 32│ 2541.75 ms │ 2545.00 ms │ no change │ │ QQuery 33│ 3179.10 ms │ 3206.02 ms │ no change │ │ QQuery 34│ 3202.13 ms │ 3242.56 ms │ no change │ │ QQuery 35│ 1219.46 ms │ 1221.75 ms │ no change │ │ QQuery 36│ 121.25 ms │ 126.01 ms │ no change │ │ QQuery 37│54.86 ms │56.17 ms │ no change │ │ QQuery 38│ 119.17 ms │ 125.52 ms │ 1.05x slower │ │ QQuery 39│ 192.11 ms │ 198.37 ms │ no change │ │ QQuery 40│47.73 ms │50.41 ms │ 1.06x slower │ │ QQuery 41│43.85 ms │45.09 ms │ no change │ │ QQuery 42│37.48 ms │38.16 ms │ no change │ └──┴─┴─┴───┘ ┏┳┓ ┃ Benchmark Summary ┃┃ ┡╇┩ │ Total Time (HEAD) │ 53787.93ms │ │ Total Time (task_budget) │ 54664.0
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980961262 🤖 `./gh_compare_branch.sh` [Benchmark Script](https://github.com/alamb/datafusion-benchmarking/blob/main/gh_compare_branch.sh) Running Linux aal-dev 6.11.0-1015-gcp #15~24.04.1-Ubuntu SMP Thu Apr 24 20:41:05 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux Comparing task_budget (a749e8f60cbcdf5542b4b69a0a2f94c7790fab45) to 06631c25316ed2fc20ab8114d0dcc801f353fbad [diff](https://github.com/apache/datafusion/compare/06631c25316ed2fc20ab8114d0dcc801f353fbad..a749e8f60cbcdf5542b4b69a0a2f94c7790fab45) Benchmarks: tpch_mem clickbench_partitioned clickbench_extended Results will be posted here when complete -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980961066 🤖: Benchmark completed Details ``` Comparing HEAD and task_budget Benchmark clickbench_1.json ┏━━┳━┳━┳━━━┓ ┃ Query┃HEAD ┃ task_budget ┃Change ┃ ┡━━╇━╇━╇━━━┩ │ QQuery 0 │51.47 ms │46.15 ms │ +1.12x faster │ │ QQuery 1 │75.49 ms │72.84 ms │ no change │ │ QQuery 2 │ 109.24 ms │ 108.00 ms │ no change │ │ QQuery 3 │ 127.42 ms │ 119.32 ms │ +1.07x faster │ │ QQuery 4 │ 633.00 ms │ 685.11 ms │ 1.08x slower │ │ QQuery 5 │ 853.72 ms │ 863.05 ms │ no change │ │ QQuery 6 │56.74 ms │54.54 ms │ no change │ │ QQuery 7 │80.42 ms │82.07 ms │ no change │ │ QQuery 8 │ 881.04 ms │ 884.51 ms │ no change │ │ QQuery 9 │ 1194.27 ms │ 1202.55 ms │ no change │ │ QQuery 10│ 300.62 ms │ 297.89 ms │ no change │ │ QQuery 11│ 318.31 ms │ 323.44 ms │ no change │ │ QQuery 12│ 879.94 ms │ 888.82 ms │ no change │ │ QQuery 13│ 1245.61 ms │ 1252.15 ms │ no change │ │ QQuery 14│ 803.27 ms │ 814.41 ms │ no change │ │ QQuery 15│ 803.96 ms │ 822.93 ms │ no change │ │ QQuery 16│ 1630.13 ms │ 1686.71 ms │ no change │ │ QQuery 17│ 1584.79 ms │ 1608.29 ms │ no change │ │ QQuery 18│ 2894.48 ms │ 2899.79 ms │ no change │ │ QQuery 19│ 125.96 ms │ 119.63 ms │ +1.05x faster │ │ QQuery 20│ 1140.75 ms │ 1206.46 ms │ 1.06x slower │ │ QQuery 21│ 1337.94 ms │ 1390.68 ms │ no change │ │ QQuery 22│ 2323.74 ms │ 2408.18 ms │ no change │ │ QQuery 23│ 7822.32 ms │ 8103.67 ms │ no change │ │ QQuery 24│ 489.54 ms │ 507.76 ms │ no change │ │ QQuery 25│ 406.25 ms │ 432.43 ms │ 1.06x slower │ │ QQuery 26│ 545.86 ms │ 568.63 ms │ no change │ │ QQuery 27│ 1638.18 ms │ 1683.81 ms │ no change │ │ QQuery 28│ 12379.17 ms │ 12511.87 ms │ no change │ │ QQuery 29│ 565.53 ms │ 581.52 ms │ no change │ │ QQuery 30│ 803.29 ms │ 816.97 ms │ no change │ │ QQuery 31│ 843.85 ms │ 872.76 ms │ no change │ │ QQuery 32│ 2547.32 ms │ 2563.59 ms │ no change │ │ QQuery 33│ 3256.21 ms │ 3307.21 ms │ no change │ │ QQuery 34│ 3266.60 ms │ 3323.18 ms │ no change │ │ QQuery 35│ 1258.28 ms │ 1288.71 ms │ no change │ │ QQuery 36│ 173.64 ms │ 181.16 ms │ no change │ │ QQuery 37│ 101.14 ms │ 101.04 ms │ no change │ │ QQuery 38│ 171.00 ms │ 179.89 ms │ 1.05x slower │ │ QQuery 39│ 252.45 ms │ 266.95 ms │ 1.06x slower │ │ QQuery 40│70.71 ms │88.09 ms │ 1.25x slower │ │ QQuery 41│87.06 ms │86.11 ms │ no change │ │ QQuery 42│78.55 ms │78.27 ms │ no change │ └──┴─┴─┴───┘ ┏┳┓ ┃ Benchmark Summary ┃┃ ┡╇┩ │ Total Time (HEAD) │ 56209.26ms │ │ Total Time (task_budget) │ 57381.15ms │ │ Average Time (HEAD)│ 1307.19ms │ │ Average Time (task_budget) │ 1334.45ms │ │ Queries Faster │ 3 │ │ Queries Slower │ 6 │ │ Queries with No Change │ 34 │ │ Queries with Failure │ 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: [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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2152534664
##
datafusion/physical-plan/src/execution_plan.rs:
##
@@ -912,25 +948,7 @@ impl PlanProperties {
/// 2. CoalescePartitionsExec for collapsing all of the partitions into
one without ordering guarantee
/// 3. SortPreservingMergeExec for collapsing all of the sorted partitions
into one with ordering guarantee
pub fn need_data_exchange(plan: Arc) -> bool {
-if let Some(repartition) = plan.as_any().downcast_ref::()
{
-!matches!(
-repartition.properties().output_partitioning(),
-Partitioning::RoundRobinBatch(_)
-)
-} else if let Some(coalesce) =
plan.as_any().downcast_ref::()
-{
-coalesce.input().output_partitioning().partition_count() > 1
-} else if let Some(sort_preserving_merge) =
-plan.as_any().downcast_ref::()
-{
-sort_preserving_merge
-.input()
-.output_partitioning()
-.partition_count()
-> 1
-} else {
-false
-}
+plan.properties().evaluation_type == EvaluationType::Lazy
Review Comment:
that is certainly a lot nicer and easier to understand
##
datafusion/physical-plan/src/execution_plan.rs:
##
@@ -743,6 +733,38 @@ pub enum EmissionType {
Both,
}
+/// Represents whether an operator's `Stream` has been implemented to actively
cooperate with the
+/// Tokio scheduler or not.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SchedulingType {
+/// The stream generated by [`execute`](ExecutionPlan::execute) does not
actively participate in
+/// cooperative scheduling. This means the implementation of the `Stream`
returned by
+/// [`ExecutionPlan::execute`] does not contain explicit cooperative yield
points.
+Blocking,
+/// The stream generated by [`execute`](ExecutionPlan::execute) actively
participates in
+/// cooperative scheduling by consuming task budget when it was able to
produce a
+/// [`RecordBatch`]. Please refer to the [`coop`](crate::coop) module for
more details.
+Cooperative,
+}
+
+/// Represents how an operator's `Stream` implementation generates
`RecordBatch`es.
+///
+/// Most operators in DataFusion generate `RecordBatch`es when asked to do so
by a call to
+/// `Stream::poll_next`. This is known as demand-driven or lazy evaluation.
+///
+/// Some operators like `Repartition` need to drive `RecordBatch` generation
themselves though. This
+/// is known as data-driven or eager evaluation.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum EvaluationType {
+/// The stream generated by [`execute`](ExecutionPlan::execute) only
generates `RecordBatch`
+/// instances when it is demanded by invoking `Stream::poll_next`.
+Lazy,
+/// The stream generated by [`execute`](ExecutionPlan::execute) eagerly
generates `RecordBatch`
+/// in one or more spawned Tokio tasks. Eager evaluation is only started
the first time
+/// `Stream::poll_next` is called.
Review Comment:
```suggestion
/// `Stream::poll_next` is called. Hash aggregation and HashJoin are
examples of such operators
```
##
datafusion/physical-plan/src/execution_plan.rs:
##
@@ -743,6 +733,38 @@ pub enum EmissionType {
Both,
}
+/// Represents whether an operator's `Stream` has been implemented to actively
cooperate with the
+/// Tokio scheduler or not.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SchedulingType {
+/// The stream generated by [`execute`](ExecutionPlan::execute) does not
actively participate in
+/// cooperative scheduling. This means the implementation of the `Stream`
returned by
+/// [`ExecutionPlan::execute`] does not contain explicit cooperative yield
points.
Review Comment:
I think it would also be helpful here to give an explicitl example of what a
cooperative yield point is. For example
```suggestion
/// [`ExecutionPlan::execute`] does not contain explicit cooperative
yield points such as
/// `await` or [`tokio::task::yield_now`].
```
##
datafusion/core/tests/execution/coop.rs:
##
@@ -0,0 +1,747 @@
+// 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
+
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980925371 > 🤖: Benchmark completed Will see if it is reproducable -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980898342 🤖 `./gh_compare_branch.sh` [Benchmark Script](https://github.com/alamb/datafusion-benchmarking/blob/main/gh_compare_branch.sh) Running Linux aal-dev 6.11.0-1015-gcp #15~24.04.1-Ubuntu SMP Thu Apr 24 20:41:05 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux Comparing task_budget (a749e8f60cbcdf5542b4b69a0a2f94c7790fab45) to 06631c25316ed2fc20ab8114d0dcc801f353fbad [diff](https://github.com/apache/datafusion/compare/06631c25316ed2fc20ab8114d0dcc801f353fbad..a749e8f60cbcdf5542b4b69a0a2f94c7790fab45) Benchmarks: clickbench_1 Results will be posted here when complete -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980898244 🤖: Benchmark completed Details ``` Comparing HEAD and task_budget Benchmark clickbench_extended.json ┏━━┳━┳━┳━━━┓ ┃ Query┃HEAD ┃ task_budget ┃Change ┃ ┡━━╇━╇━╇━━━┩ │ QQuery 0 │ 1938.69 ms │ 1955.98 ms │ no change │ │ QQuery 1 │ 728.67 ms │ 732.71 ms │ no change │ │ QQuery 2 │ 1393.40 ms │ 1380.16 ms │ no change │ │ QQuery 3 │ 654.25 ms │ 667.02 ms │ no change │ │ QQuery 4 │ 1374.93 ms │ 1368.90 ms │ no change │ │ QQuery 5 │ 14843.62 ms │ 14967.16 ms │ no change │ │ QQuery 6 │ 2001.52 ms │ 2082.06 ms │ no change │ │ QQuery 7 │ 1889.83 ms │ 1950.81 ms │ no change │ │ QQuery 8 │ 805.12 ms │ 809.25 ms │ no change │ └──┴─┴─┴───┘ ┏┳┓ ┃ Benchmark Summary ┃┃ ┡╇┩ │ Total Time (HEAD) │ 25630.03ms │ │ Total Time (task_budget) │ 25914.06ms │ │ Average Time (HEAD)│ 2847.78ms │ │ Average Time (task_budget) │ 2879.34ms │ │ Queries Faster │ 0 │ │ Queries Slower │ 0 │ │ Queries with No Change │ 9 │ │ Queries with Failure │ 0 │ └┴┘ Benchmark clickbench_partitioned.json ┏━━┳━┳━┳━━━┓ ┃ Query┃HEAD ┃ task_budget ┃Change ┃ ┡━━╇━╇━╇━━━┩ │ QQuery 0 │15.26 ms │15.20 ms │ no change │ │ QQuery 1 │33.51 ms │32.70 ms │ no change │ │ QQuery 2 │81.78 ms │80.44 ms │ no change │ │ QQuery 3 │ 100.21 ms │96.71 ms │ no change │ │ QQuery 4 │ 578.10 ms │ 585.73 ms │ no change │ │ QQuery 5 │ 838.83 ms │ 887.89 ms │ 1.06x slower │ │ QQuery 6 │24.33 ms │22.32 ms │ +1.09x faster │ │ QQuery 7 │38.56 ms │36.90 ms │ no change │ │ QQuery 8 │ 850.10 ms │ 856.20 ms │ no change │ │ QQuery 9 │ 1151.01 ms │ 1146.42 ms │ no change │ │ QQuery 10│ 252.88 ms │ 263.74 ms │ no change │ │ QQuery 11│ 283.58 ms │ 290.15 ms │ no change │ │ QQuery 12│ 869.70 ms │ 866.73 ms │ no change │ │ QQuery 13│ 1255.69 ms │ 1286.10 ms │ no change │ │ QQuery 14│ 809.83 ms │ 812.02 ms │ no change │ │ QQuery 15│ 777.07 ms │ 786.77 ms │ no change │ │ QQuery 16│ 1613.49 ms │ 1629.44 ms │ no change │ │ QQuery 17│ 1567.97 ms │ 1580.52 ms │ no change │ │ QQuery 18│ 2886.23 ms │ 2891.01 ms │ no change │ │ QQuery 19│87.29 ms │84.38 ms │ no change │ │ QQuery 20│ 1113.26 ms │ 1152.62 ms │ no change │ │ QQuery 21│ 1252.92 ms │ 1328.12 ms │ 1.06x slower │ │ QQuery 22│ 2069.04 ms │ 2168.76 ms │ no change │ │ QQuery 23│ 7539.31 ms │ 7957.48 ms │ 1.06x slower │ │ QQuery 24│ 443.12 ms │ 464.35 ms │ no change │ │ QQuery 25│ 379.12 ms │ 404.81 ms │ 1.07x slower │ │ QQuery 26│ 503.89 ms │ 530.15 ms │ 1.05x slower │ │ QQuery 27│ 1500.52 ms │ 1566.86 ms │ no change │ │ QQuery 28│ 12013.51 ms │ 11817.68 ms │ no change │ │ QQuery 29│ 527.40 ms │ 525.63 ms │ no change │ │ QQuery 30│ 758.72 ms │ 791.22 ms │ no change │ │ QQuery 31│ 811.10 ms │ 829.72 ms │ no change │ │ QQuery 32│ 2485.30 ms │ 2553.71 ms │ no change │ │ QQuery 33│ 3181.91 ms │ 3185.99 ms │ no change │ │ QQuery 34│ 3176.40 ms │ 3209.67 ms │ no change │ │ QQuery 35│ 1228.64 ms │ 1223.48 ms │ no change │ │ QQuery 36│ 120.17 ms │ 127.76 ms │ 1.06x slower │ │ QQuery 37│57.02 ms │56.75 ms │ no change │ │ QQuery 38│ 120.93 ms │ 127.28 ms │ 1.05x slower │ │ QQuery 39│ 191.04 ms │ 202.25 ms │ 1.06x slower │ │ QQuery 40│47.95 ms │47.32 ms │ no change │ │ QQuery 41│45.18 ms │45.53 ms │ no change │ │ QQuery 42│36.63 ms │39.44 ms │ 1.08x slower │ └──┴─┴─┴───┘ ┏┳┓ ┃ Benchmark Summary ┃┃ ┡╇┩ │ Total Time (HEAD) │ 53718.51ms │ │ Total Time (task_budget) │ 54607.96ms │ │ Average Time (HEAD)│
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980783229 > I agree the existing benchmarks are not perfect Not trying to criticize the benchmarks, just saying I'm having trouble using it as a guide. @alamb Note that by default the current branch is using the local counter method @zhuqi-lucas implemented for `YieldStream`. Performance should be identical to `main` in other words. -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980750548 🤖 `./gh_compare_branch.sh` [Benchmark Script](https://github.com/alamb/datafusion-benchmarking/blob/main/gh_compare_branch.sh) Running Linux aal-dev 6.11.0-1015-gcp #15~24.04.1-Ubuntu SMP Thu Apr 24 20:41:05 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux Comparing task_budget (5bfff9798b2815377ef8d40dfaffd26947c51385) to 06631c25316ed2fc20ab8114d0dcc801f353fbad [diff](https://github.com/apache/datafusion/compare/06631c25316ed2fc20ab8114d0dcc801f353fbad..5bfff9798b2815377ef8d40dfaffd26947c51385) Benchmarks: tpch_mem clickbench_partitioned clickbench_extended Results will be posted here when complete -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
alamb commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980748403 > I've been trying to measure the impact of these changes, but for the life of me I can't get useful results out of the benchmark comparison script. The sample below is actually comparing the exact same commit on a mac mini M4 and yet I still see large fluctuations. Not sure what else I can do 🤷♂️. I agree the existing benchmarks are not perfect -- figuring out how to make them better would be good, though as you probably know it is a challenge of benchmarking in general. I'll kick off some runs against this PR with my scripts to see what it shows. -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2152509411
##
datafusion/datasource/src/source.rs:
##
@@ -261,11 +261,7 @@ impl ExecutionPlan for DataSourceExec {
) -> Result {
self.data_source
.open(partition, Arc::clone(&context))
-.map(|stream| wrap_yield_stream(stream, &context,
self.cooperative))
-}
-
-fn with_cooperative_yields(self: Arc) -> Option> {
-self.cooperative.then_some(self)
+.map(|stream| make_cooperative(stream))
Review Comment:
I've gone ahead and made that change and removed the dyn wrapper while I was
at it.
--
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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2152483273
##
datafusion/datasource/src/source.rs:
##
@@ -261,11 +261,7 @@ impl ExecutionPlan for DataSourceExec {
) -> Result {
self.data_source
.open(partition, Arc::clone(&context))
-.map(|stream| wrap_yield_stream(stream, &context,
self.cooperative))
-}
-
-fn with_cooperative_yields(self: Arc) -> Option> {
-self.cooperative.then_some(self)
+.map(|stream| make_cooperative(stream))
Review Comment:
That would make sense indeed. ExecutionPlan fulfills the role of a factory
of Streams. DataSourceExec delegates that to the DataSource itself. So
Datasource is actually the factory. It would make sense to delegate the 'will
you create something cooperative?' query as well.
--
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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
ozankabak commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2152472886
##
datafusion/datasource/src/source.rs:
##
@@ -261,11 +261,7 @@ impl ExecutionPlan for DataSourceExec {
) -> Result {
self.data_source
.open(partition, Arc::clone(&context))
-.map(|stream| wrap_yield_stream(stream, &context,
self.cooperative))
-}
-
-fn with_cooperative_yields(self: Arc) -> Option> {
-self.cooperative.then_some(self)
+.map(|stream| make_cooperative(stream))
Review Comment:
I don't think it is critical performance-wise, but now I wonder if we should
make any additions to the `DataSource` trait to signal built-in cooperative
behavior to `DataSourceExec`?
--
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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2152379108
##
datafusion/datasource/src/source.rs:
##
@@ -261,11 +261,7 @@ impl ExecutionPlan for DataSourceExec {
) -> Result {
self.data_source
.open(partition, Arc::clone(&context))
-.map(|stream| wrap_yield_stream(stream, &context,
self.cooperative))
-}
-
-fn with_cooperative_yields(self: Arc) -> Option> {
-self.cooperative.then_some(self)
+.map(|stream| make_cooperative(stream))
Review Comment:
We could avoid a virtual call here by pushing the cooperative wrapper down
to `FileScanConfig::open`. That does lose the benefit of `DataSourceExec` being
cooperative for any `DataSource`. Is it worth shaving off one virtual call per
record batch?
--
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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2980397594 I've been trying to measure the impact of these changes, but for the life of me I can't get useful results out of the benchmark comparison script. The sample below is actually comparing the exact same commit on a mac mini M4 and yet I still see large fluctuations. Not sure what else I can do 🤷♂️. ``` ┏━━┳┳┳━━━┓ ┃ Query┃ baseline ┃ branch ┃Change ┃ ┡━━╇╇╇━━━┩ │ QQuery 0 │9.71 ms │9.87 ms │ no change │ │ QQuery 1 │ 15.07 ms │ 15.36 ms │ no change │ │ QQuery 2 │ 47.76 ms │ 47.81 ms │ no change │ │ QQuery 3 │ 46.78 ms │ 47.13 ms │ no change │ │ QQuery 4 │ 426.60 ms │ 426.87 ms │ no change │ │ QQuery 5 │ 557.13 ms │ 559.79 ms │ no change │ │ QQuery 6 │ 11.37 ms │ 11.30 ms │ no change │ │ QQuery 7 │ 16.17 ms │ 15.99 ms │ no change │ │ QQuery 8 │ 518.13 ms │ 516.98 ms │ no change │ │ QQuery 9 │ 671.89 ms │ 669.00 ms │ no change │ │ QQuery 10│ 146.75 ms │ 149.95 ms │ no change │ │ QQuery 11│ 162.79 ms │ 164.02 ms │ no change │ │ QQuery 12│ 567.28 ms │ 569.63 ms │ no change │ │ QQuery 13│ 799.84 ms │ 796.18 ms │ no change │ │ QQuery 14│ 542.10 ms │ 535.99 ms │ no change │ │ QQuery 15│ 502.33 ms │ 500.89 ms │ no change │ │ QQuery 16│ 1151.00 ms │ 1150.96 ms │ no change │ │ QQuery 17│ 1127.41 ms │ 1116.08 ms │ no change │ │ QQuery 18│ 2195.19 ms │ 2182.97 ms │ no change │ │ QQuery 19│ 39.96 ms │ 40.64 ms │ no change │ │ QQuery 20│ 858.48 ms │ 859.71 ms │ no change │ │ QQuery 21│ 988.68 ms │ 1048.57 ms │ 1.06x slower │ │ QQuery 22│ 1594.48 ms │ 2153.36 ms │ 1.35x slower │ │ QQuery 23│ 6399.31 ms │ 6905.44 ms │ 1.08x slower │ │ QQuery 24│ 305.36 ms │ 299.88 ms │ no change │ │ QQuery 25│ 274.50 ms │ 276.00 ms │ no change │ │ QQuery 26│ 338.09 ms │ 346.69 ms │ no change │ │ QQuery 27│ 1140.21 ms │ 1213.95 ms │ 1.06x slower │ │ QQuery 28│ 8881.31 ms │ 9061.62 ms │ no change │ │ QQuery 29│ 417.02 ms │ 414.85 ms │ no change │ │ QQuery 30│ 487.07 ms │ 486.73 ms │ no change │ │ QQuery 31│ 507.31 ms │ 508.16 ms │ no change │ │ QQuery 32│ 2154.40 ms │ 1917.22 ms │ +1.12x faster │ │ QQuery 33│ 2498.71 ms │ 2467.38 ms │ no change │ │ QQuery 34│ 2678.02 ms │ 2802.00 ms │ no change │ │ QQuery 35│ 751.17 ms │ 772.78 ms │ no change │ │ QQuery 36│ 61.54 ms │ 62.50 ms │ no change │ │ QQuery 37│ 25.67 ms │ 25.56 ms │ no change │ │ QQuery 38│ 61.80 ms │ 62.39 ms │ no change │ │ QQuery 39│ 101.51 ms │ 102.13 ms │ no change │ │ QQuery 40│ 18.26 ms │ 19.09 ms │ no change │ │ QQuery 41│ 17.70 ms │ 17.77 ms │ no change │ │ QQuery 42│ 14.84 ms │ 15.32 ms │ no 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
ozankabak commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2979894941 I took a quick look at this yesterday and it was looking good -- I will review in more detail in a day or two. Thanks -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2979874048 The current Tokio code is only an approximation of the final code. For the final version we're dependent on a Tokio PR getting merged and released. In order to unblock this PR and allow review to proceed, I've added a couple of feature flags that make the Tokio based code disabled by default. The intention is to allow comparing the two variants from a performance point-of-view and consider the revised cooperative API for merging already. -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2151897561
##
datafusion/core/tests/execution/coop.rs:
##
@@ -0,0 +1,722 @@
+// 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.
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::SortOptions;
+use datafusion::functions_aggregate::sum;
+use datafusion::physical_expr::aggregate::AggregateExprBuilder;
+use datafusion::physical_plan;
+use datafusion::physical_plan::aggregates::{
+AggregateExec, AggregateMode, PhysicalGroupBy,
+};
+use datafusion::physical_plan::execution_plan::Boundedness;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionContext;
+use datafusion_common::{JoinType, ScalarValue};
+use datafusion_execution::TaskContext;
+use datafusion_expr_common::operator::Operator;
+use datafusion_expr_common::operator::Operator::Gt;
+use datafusion_functions_aggregate::min_max;
+use datafusion_physical_expr::expressions::{
+binary, col, lit, BinaryExpr, Column, Literal,
+};
+use datafusion_physical_expr::Partitioning;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::{LexOrdering,
PhysicalSortExpr};
+use datafusion_physical_optimizer::ensure_coop::EnsureCooperative;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
+use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode,
SortMergeJoinExec};
+use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::union::InterleaveExec;
+use futures::StreamExt;
+use parking_lot::RwLock;
+use rstest::rstest;
+use std::error::Error;
+use std::fmt::Formatter;
+use std::ops::Range;
+use std::sync::Arc;
+use std::task::Poll;
+use tokio::runtime::{Handle, Runtime};
+use tokio::select;
+
+#[derive(Debug)]
+struct RangeBatchGenerator {
+schema: SchemaRef,
+value_range: Range,
+boundedness: Boundedness,
+batch_size: usize,
+poll_count: usize,
+}
+
+impl std::fmt::Display for RangeBatchGenerator {
+fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+// Display current counter
+write!(f, "InfiniteGenerator(counter={})", self.poll_count)
+}
+}
+
+impl LazyBatchGenerator for RangeBatchGenerator {
+fn boundedness(&self) -> Boundedness {
+self.boundedness
+}
+
+/// Generate the next RecordBatch.
+fn generate_next_batch(&mut self) ->
datafusion_common::Result> {
+self.poll_count += 1;
+
+let mut builder = Int64Array::builder(self.batch_size);
+for _ in 0..self.batch_size {
+match self.value_range.next() {
+None => break,
+Some(v) => builder.append_value(v),
+}
+}
+let array = builder.finish();
+
+if array.is_empty() {
+return Ok(None);
+}
+
+let batch =
+RecordBatch::try_new(Arc::clone(&self.schema),
vec![Arc::new(array)])?;
+Ok(Some(batch))
+}
+}
+
+fn make_lazy_exec(column_name: &str, pretend_infinite: bool) -> LazyMemoryExec
{
+make_lazy_exec_with_range(column_name, i64::MIN..i64::MAX,
pretend_infinite)
+}
+
+fn make_lazy_exec_with_range(
+column_name: &str,
+range: Range,
+pretend_infinite: bool,
+) -> LazyMemoryExec {
+let schema = Arc::new(Schema::new(vec![Field::new(
+column_name,
+DataType::Int64,
+false,
+)]));
+
+let boundedness = if pretend_infinite {
+Boundedness::Unbounded {
+requires_infinite_memory: false,
+}
+} else {
+Boundedness::Bounded
+};
+
+// Instantiate the generator with the batch and limit
+let gen = RangeBatchGenerator {
+schema: Arc::clone(&schema),
+boundedness,
+value_range: range,
+batch_
Re: [PR] Use Tokio's task budget consistently [datafusion]
zhuqi-lucas commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2151253373
##
datafusion/core/tests/execution/coop.rs:
##
@@ -0,0 +1,722 @@
+// 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.
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::SortOptions;
+use datafusion::functions_aggregate::sum;
+use datafusion::physical_expr::aggregate::AggregateExprBuilder;
+use datafusion::physical_plan;
+use datafusion::physical_plan::aggregates::{
+AggregateExec, AggregateMode, PhysicalGroupBy,
+};
+use datafusion::physical_plan::execution_plan::Boundedness;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionContext;
+use datafusion_common::{JoinType, ScalarValue};
+use datafusion_execution::TaskContext;
+use datafusion_expr_common::operator::Operator;
+use datafusion_expr_common::operator::Operator::Gt;
+use datafusion_functions_aggregate::min_max;
+use datafusion_physical_expr::expressions::{
+binary, col, lit, BinaryExpr, Column, Literal,
+};
+use datafusion_physical_expr::Partitioning;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::{LexOrdering,
PhysicalSortExpr};
+use datafusion_physical_optimizer::ensure_coop::EnsureCooperative;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
+use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode,
SortMergeJoinExec};
+use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::union::InterleaveExec;
+use futures::StreamExt;
+use parking_lot::RwLock;
+use rstest::rstest;
+use std::error::Error;
+use std::fmt::Formatter;
+use std::ops::Range;
+use std::sync::Arc;
+use std::task::Poll;
+use tokio::runtime::{Handle, Runtime};
+use tokio::select;
+
+#[derive(Debug)]
+struct RangeBatchGenerator {
+schema: SchemaRef,
+value_range: Range,
+boundedness: Boundedness,
+batch_size: usize,
+poll_count: usize,
+}
+
+impl std::fmt::Display for RangeBatchGenerator {
+fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+// Display current counter
+write!(f, "InfiniteGenerator(counter={})", self.poll_count)
+}
+}
+
+impl LazyBatchGenerator for RangeBatchGenerator {
+fn boundedness(&self) -> Boundedness {
+self.boundedness
+}
+
+/// Generate the next RecordBatch.
+fn generate_next_batch(&mut self) ->
datafusion_common::Result> {
+self.poll_count += 1;
+
+let mut builder = Int64Array::builder(self.batch_size);
+for _ in 0..self.batch_size {
+match self.value_range.next() {
+None => break,
+Some(v) => builder.append_value(v),
+}
+}
+let array = builder.finish();
+
+if array.is_empty() {
+return Ok(None);
+}
+
+let batch =
+RecordBatch::try_new(Arc::clone(&self.schema),
vec![Arc::new(array)])?;
+Ok(Some(batch))
+}
+}
+
+fn make_lazy_exec(column_name: &str, pretend_infinite: bool) -> LazyMemoryExec
{
+make_lazy_exec_with_range(column_name, i64::MIN..i64::MAX,
pretend_infinite)
+}
+
+fn make_lazy_exec_with_range(
+column_name: &str,
+range: Range,
+pretend_infinite: bool,
+) -> LazyMemoryExec {
+let schema = Arc::new(Schema::new(vec![Field::new(
+column_name,
+DataType::Int64,
+false,
+)]));
+
+let boundedness = if pretend_infinite {
+Boundedness::Unbounded {
+requires_infinite_memory: false,
+}
+} else {
+Boundedness::Bounded
+};
+
+// Instantiate the generator with the batch and limit
+let gen = RangeBatchGenerator {
+schema: Arc::clone(&schema),
+boundedness,
+value_range: range,
+bat
Re: [PR] Use Tokio's task budget consistently [datafusion]
zhuqi-lucas commented on code in PR #16398:
URL: https://github.com/apache/datafusion/pull/16398#discussion_r2151253373
##
datafusion/core/tests/execution/coop.rs:
##
@@ -0,0 +1,722 @@
+// 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.
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow_schema::SortOptions;
+use datafusion::functions_aggregate::sum;
+use datafusion::physical_expr::aggregate::AggregateExprBuilder;
+use datafusion::physical_plan;
+use datafusion::physical_plan::aggregates::{
+AggregateExec, AggregateMode, PhysicalGroupBy,
+};
+use datafusion::physical_plan::execution_plan::Boundedness;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionContext;
+use datafusion_common::{JoinType, ScalarValue};
+use datafusion_execution::TaskContext;
+use datafusion_expr_common::operator::Operator;
+use datafusion_expr_common::operator::Operator::Gt;
+use datafusion_functions_aggregate::min_max;
+use datafusion_physical_expr::expressions::{
+binary, col, lit, BinaryExpr, Column, Literal,
+};
+use datafusion_physical_expr::Partitioning;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::{LexOrdering,
PhysicalSortExpr};
+use datafusion_physical_optimizer::ensure_coop::EnsureCooperative;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
+use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode,
SortMergeJoinExec};
+use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::union::InterleaveExec;
+use futures::StreamExt;
+use parking_lot::RwLock;
+use rstest::rstest;
+use std::error::Error;
+use std::fmt::Formatter;
+use std::ops::Range;
+use std::sync::Arc;
+use std::task::Poll;
+use tokio::runtime::{Handle, Runtime};
+use tokio::select;
+
+#[derive(Debug)]
+struct RangeBatchGenerator {
+schema: SchemaRef,
+value_range: Range,
+boundedness: Boundedness,
+batch_size: usize,
+poll_count: usize,
+}
+
+impl std::fmt::Display for RangeBatchGenerator {
+fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+// Display current counter
+write!(f, "InfiniteGenerator(counter={})", self.poll_count)
+}
+}
+
+impl LazyBatchGenerator for RangeBatchGenerator {
+fn boundedness(&self) -> Boundedness {
+self.boundedness
+}
+
+/// Generate the next RecordBatch.
+fn generate_next_batch(&mut self) ->
datafusion_common::Result> {
+self.poll_count += 1;
+
+let mut builder = Int64Array::builder(self.batch_size);
+for _ in 0..self.batch_size {
+match self.value_range.next() {
+None => break,
+Some(v) => builder.append_value(v),
+}
+}
+let array = builder.finish();
+
+if array.is_empty() {
+return Ok(None);
+}
+
+let batch =
+RecordBatch::try_new(Arc::clone(&self.schema),
vec![Arc::new(array)])?;
+Ok(Some(batch))
+}
+}
+
+fn make_lazy_exec(column_name: &str, pretend_infinite: bool) -> LazyMemoryExec
{
+make_lazy_exec_with_range(column_name, i64::MIN..i64::MAX,
pretend_infinite)
+}
+
+fn make_lazy_exec_with_range(
+column_name: &str,
+range: Range,
+pretend_infinite: bool,
+) -> LazyMemoryExec {
+let schema = Arc::new(Schema::new(vec![Field::new(
+column_name,
+DataType::Int64,
+false,
+)]));
+
+let boundedness = if pretend_infinite {
+Boundedness::Unbounded {
+requires_infinite_memory: false,
+}
+} else {
+Boundedness::Bounded
+};
+
+// Instantiate the generator with the batch and limit
+let gen = RangeBatchGenerator {
+schema: Arc::clone(&schema),
+boundedness,
+value_range: range,
+bat
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2973548782 I've added a commit to this PR that: - Makes the tests more robust. Rather than hanging when there's an issue they will fail. - Removes duplication from the tests - Removes repartitions and coalesces in most tests to ensure we're actually testing the right thing. I'm to get the tests to set up the challenging situations rather than massaging things so that they work. - Merges in some of the extra test cases I still had lying around -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
ozankabak commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2973514966 > Thinking about it some more. The evaluation type is intended to describe how the operator computes record batches itself: lazy on demand, or by driving things itself. I’m kind of trying to refer to the terminology from the volcano paper. That talks about demand-driven and data-driven operators. I had first called this 'drive type' with values 'demand' and 'data', but that felt a bit awkward. Since this is actually a property of how the operator prepares its output, one value per operator is probably fine after all. OK - let's try that way, we can change later if it turns out to be insufficient -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2973036248 Open to suggestions on better names for these properties. -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2973035858 Thinking about it some more. The evaluation type is intended to describe how the operator computes record batches itself: lazy on demand, or by driving things itself. I’m kind of trying to refer to the terminology from the volcano paper. That task about demand-driven and data-driven operators. Since this is actually a property of how the operator prepare its output, one value per operator is probably fine after all. -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
ozankabak commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2972621416 > While I was writing this I started wondering if evaluation type should be a per child thing. In my spawn experiment branch for instance hash join is eager for the build side, but lazy for the probe side. Perhaps it would be best to leave room for that. This is in alignment with what I was thinking, let's do it that way -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
pepijnve commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2970729760 @ozankabak I've pushed the optimizer rule changes I had in mind. This introduces two new execution plan properties that capture the evaluation type (eager vs lazy) and the scheduling type (blocking vs cooperative). With those two combined the tree can be rewritten in a bottom up fashion. Every leaf that is not cooperative gets wrapped as before. Additionally, any eager evaluating nodes (i.e. exchanges) that are not cooperative are wrapped. This should ensure the entire plan participates in cooperative scheduling. The only caveat that remains is dynamic stream creation. Operators that do that need to take the necessary precautions themselves. I already update the spill manager for this in the previous commit. -- 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]
Re: [PR] Use Tokio's task budget consistently [datafusion]
ozankabak commented on PR #16398: URL: https://github.com/apache/datafusion/pull/16398#issuecomment-2970506512 Thanks for the draft -- this is inline with my understanding from your description. I think it will inch us closer to a good, lasting solution (especially after your upstream tokio also PR merges). Feel free to ping me for a more detailed review once you are done with it -- 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]
