You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "danielhumanmod (via GitHub)" <gi...@apache.org> on 2026/01/16 05:54:56 UTC

[PR] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

danielhumanmod opened a new pull request, #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380

   # 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 #1379.
   
    # Rationale for this change
   <!--
    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.  
   -->
   
   Currently, `ShuffleReaderExec` passes batches directly from the upstream `AbortableReceiverStream` to downstream operators. In scenarios with high parallelism or network fragmentation, the shuffle reader can receives small record batches.
   
   This causes performance degradation because:
   1. Downstream operators (aggregates, joins) incur high overhead per function call (`poll_next`).
   2. Vectorized execution benefits are lost on micro-batches.
   
   # What changes are included in this PR?
   <!--
   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.
   -->
   
   1. Core Logic: Introduced `CoalescedShuffleReaderStream` as a stream adapter to buffer and merge small batches using DataFusion's `BatchCoalescer`.
   
   2. Integration: Updated `ShuffleReaderExec::execute` to apply this wrapper, enabling configurable batch sizing via SessionConfig and preserving existing metrics.
   
   # Are there any user-facing changes?
   <!--
   If there are user-facing changes then we may require documentation to be updated before approving the PR.
   -->
   
   No
   
   <!--
   If there are any breaking changes to public APIs, please add the `api change` label.
   -->
   


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "sqlbenchmark (via GitHub)" <gi...@apache.org>.
sqlbenchmark commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3775283812

   ## Ballista TPC-H Benchmark Results
   
   **PR:** #1380 - impl
   **PR Commit:** `cf58271`
   **Base Commit:** `38fca8b` (main)
   **Scale Factor:** SF1
   **Iterations:** 3
   
   ### Results
   
   Could not parse structured results. Raw output:
   
   **Main branch:**
   ```
   ery 12 avg time: 847.1 ms
   Query 13 iteration 0 took 813.5 ms and returned 42 rows
   Query 13 iteration 1 took 915.0 ms and returned 42 rows
   Query 13 iteration 2 took 955.0 ms and returned 42 rows
   Query 13 avg time: 894.5 ms
   Query 14 iteration 0 took 609.0 ms and returned 1 rows
   Query 14 iteration 1 took 552.0 ms and returned 1 rows
   Query 14 iteration 2 took 611.7 ms and returned 1 rows
   Query 14 avg time: 590.9 ms
   Query 15 iteration 0 took 1020.6 ms and returned 0 rows
   Query 15 iteration 1 took 917.2 ms and returned 0 rows
   Query 15 iteration 2 took 917.6 ms and returned 0 rows
   Query 15 avg time: 951.8 ms
   Query 16 iteration 0 took 1222.9 ms and returned 18314 rows
   Query 16 iteration 1 took 1220.9 ms and returned 18314 rows
   Query 16 iteration 2 took 1223.2 ms and returned 18314 rows
   Query 16 avg time: 1222.3 ms
   Query 17 iteration 0 took 754.9 ms and returned 1 rows
   Query 17 iteration 1 took 713.6 ms and returned 1 rows
   Query 17 iteration 2 took 816.5 ms and returned 1 rows
   Query 17 avg time: 761.6 ms
   Query 18 iteration 0 took 1262.5 ms and returned 57 rows
   Query 18 iteration 1 took 1321.6 ms and returned 57 rows
   Query 18 iteration 2 took 1220.5 ms and returned 57 rows
   Query 18 avg time: 1268.2 ms
   Query 19 iteration 0 took 714.9 ms and returned 1 rows
   Query 19 iteration 1 took 715.1 ms and returned 1 rows
   Query 19 iteration 2 took 756.3 ms and returned 1 rows
   Query 19 avg time: 728.8 ms
   Query 20 iteration 0 took 1059.4 ms and returned 186 rows
   Query 20 iteration 1 took 917.3 ms and returned 186 rows
   Query 20 iteration 2 took 957.9 ms and returned 186 rows
   Query 20 avg time: 978.2 ms
   Query 21 iteration 0 took 1769.6 ms and returned 100 rows
   Query 21 iteration 1 took 1731.1 ms and returned 100 rows
   Query 21 iteration 2 took 1625.9 ms and returned 100 rows
   Query 21 avg time: 1708.9 ms
   Query 22 iteration 0 took 918.1 ms and returned 7 rows
   Query 22 iteration 1 took 918.6 ms and returned 7 rows
   Query 22 iteration 2 took 917.2 ms and returned 7 rows
   Query 22 avg time: 917.9 ms
   
   ```
   
   **PR branch:**
   ```
       Finished `release` profile [optimized] target(s) in 0.18s
        Running `target/release/tpch benchmark ballista --host localhost --port 50050 --path /app/data --format parquet --iterations 3`
   error: The following required arguments were not provided:
       --query <query>
   
   USAGE:
       tpch benchmark ballista --batch-size <batch-size> --format <file-format> --host <host> --iterations <iterations> --partitions <partitions> --path <path> --port <port> --query <query>
   
   For more information try --help
   
   ```
   
   ---
   *Automated benchmark run by dfbench*


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3778214894

   > @sqlbenchmark run tpch
   
   There is currently a list of contributors that have permissions to run this. I will open this up to more people once it has been through more testing. It is very experimental at the moment.


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3778130012

    @sqlbenchmark run tpch


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3768100430

   For benhmarking the performance change for small batches, probably it is better to run against Parquet (to make the scan less of a bottleneck) and with a high number of task slots (to make the batch size small).


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm merged PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "sqlbenchmark (via GitHub)" <gi...@apache.org>.
sqlbenchmark commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3778283960

   ## Ballista TPC-H Benchmark Results
   
   **PR:** #1380 - fix format and simplify new
   **PR Commit:** `b7be1f7`
   **Base Commit:** `cd4ead8` (main)
   **Scale Factor:** SF1
   **Iterations:** 3
   
   ### Query Comparison
   
   | Query | Main (ms) | PR (ms) | Change |
   |-------|-----------|---------|--------|
   | Q1 | 816.10 | 817.00 | :white_circle: +0.1% |
   | Q2 | 1351.00 | 1318.40 | :white_circle: -2.4% |
   | Q3 | 1100.30 | 999.00 | :green_circle: -9.2% |
   | Q4 | 827.50 | 827.40 | :white_circle: -0.0% |
   | Q5 | 1655.40 | 1776.20 | :red_circle: +7.3% |
   | Q6 | 401.20 | 408.50 | :white_circle: +1.8% |
   | Q7 | 1798.20 | 1898.70 | :red_circle: +5.6% |
   | Q8 | 2387.20 | 2319.40 | :white_circle: -2.8% |
   | Q9 | 1856.20 | 1912.10 | :white_circle: +3.0% |
   | Q10 | 1268.40 | 1350.10 | :red_circle: +6.4% |
   | Q11 | 1031.60 | 1032.60 | :white_circle: +0.1% |
   | Q12 | 874.90 | 827.60 | :green_circle: -5.4% |
   | Q13 | 948.00 | 929.20 | :white_circle: -2.0% |
   | Q14 | 639.00 | 625.60 | :white_circle: -2.1% |
   | Q15 | 1087.30 | 984.90 | :green_circle: -9.4% |
   | Q16 | 1236.80 | 1222.10 | :white_circle: -1.2% |
   | Q17 | 741.70 | 740.90 | :white_circle: -0.1% |
   | Q18 | 1371.20 | 1417.30 | :white_circle: +3.4% |
   | Q19 | 782.60 | 749.00 | :white_circle: -4.3% |
   | Q20 | 999.10 | 1031.90 | :white_circle: +3.3% |
   | Q21 | 1771.50 | 1762.80 | :white_circle: -0.5% |
   | Q22 | 930.30 | 910.20 | :white_circle: -2.2% |
   
   **Total:** Main=25875.50ms, PR=25860.90ms (-0.1%)
   
   ---
   *Automated benchmark run by dfbench*


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on code in PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#discussion_r2710915797


##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -197,11 +201,22 @@ impl ExecutionPlan for ShuffleReaderExec {
             prefer_flight,
         );
 
-        let result = RecordBatchStreamAdapter::new(
-            Arc::new(self.schema.as_ref().clone()),
+        let input_stream = Box::pin(RecordBatchStreamAdapter::new(
+            self.schema.clone(),
             response_receiver.try_flatten(),
-        );
-        Ok(Box::pin(result))
+        ));
+
+        Ok(Box::pin(CoalescedShuffleReaderStream {
+            schema: self.schema.clone(),
+            input: input_stream,
+            coalescer: LimitedBatchCoalescer::new(

Review Comment:
   > what is the reason using `LimitedBatchCoalescer` instead of `BatchCoalescer` as there is no fetch limit?
   > 
   > would there be a case where we can use limit 🤔?
   
   Regarding `LimitedBatchCoalescer`, I initially chose it because it offers better encapsulation for the stream state machine compared to the raw `BatchCoalescer`. Specifically, it manages the finished state and provides PushBatchStatus, which makes the poll_next implementation cleaner and safer, even without a limit.



-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3775485250

   Perhaps @danielhumanmod could rebase the PR


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3778217645

   @sqlbenchmark run tpch -s 10 -i 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


Re: [PR] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3758575639

   But one wired thing is the performance seems not improve based on benchmark test result, I used the `Running the Ballista Benchmarks` setting
   
   After
   ```
   st --port 60000 --query 10 --path ./data --format tbl --iterations 3`
   Running benchmarks with the following options: BallistaBenchmarkOpt { query: 10, debug: false, expected_results: None, iterations: 3, batch_size: 8192, path: "./data", file_format: "tbl", partitions: 2, host: Some("localhost"), port: Some(60000), output_path: None }
   Running benchmark with queries 10:
    ["select\n    c_custkey,\n    c_name,\n    sum(l_extendedprice * (1 - l_discount)) as revenue,\n    c_acctbal,\n    n_name,\n    c_address,\n    c_phone,\n    c_comment\nfrom\n    customer,\n    orders,\n    lineitem,\n    nation\nwhere\n        c_custkey = o_custkey\n  and l_orderkey = o_orderkey\n  and o_orderdate >= date '1993-10-01'\n  and o_orderdate < date '1994-01-01'\n  and l_returnflag = 'R'\n  and c_nationkey = n_nationkey\ngroup by\n    c_custkey,\n    c_name,\n    c_acctbal,\n    c_phone,\n    n_name,\n    c_address,\n    c_comment\norder by\n    revenue desc\nlimit 20"]
   Query 10 iteration 0 took 2533.0 ms and returned 20 rows
   Query 10 iteration 1 took 1853.6 ms and returned 20 rows
   Query 10 iteration 2 took 1958.3 ms and returned 20 rows
   Query 10 avg time: 2114.96 ms
   ```
   Before 
   ```
   st --port 60000 --query 10 --path ./data --format tbl --iterations 3`
   Running benchmarks with the following options: BallistaBenchmarkOpt { query: 10, debug: false, expected_results: None, iterations: 3, batch_size: 8192, path: "./data", file_format: "tbl", partitions: 2, host: Some("localhost"), port: Some(60000), output_path: None }
   Running benchmark with queries 10:
    ["select\n    c_custkey,\n    c_name,\n    sum(l_extendedprice * (1 - l_discount)) as revenue,\n    c_acctbal,\n    n_name,\n    c_address,\n    c_phone,\n    c_comment\nfrom\n    customer,\n    orders,\n    lineitem,\n    nation\nwhere\n        c_custkey = o_custkey\n  and l_orderkey = o_orderkey\n  and o_orderdate >= date '1993-10-01'\n  and o_orderdate < date '1994-01-01'\n  and l_returnflag = 'R'\n  and c_nationkey = n_nationkey\ngroup by\n    c_custkey,\n    c_name,\n    c_acctbal,\n    c_phone,\n    n_name,\n    c_address,\n    c_comment\norder by\n    revenue desc\nlimit 20"]
   Query 10 iteration 0 took 2739.5 ms and returned 20 rows
   Query 10 iteration 1 took 1855.9 ms and returned 20 rows
   Query 10 iteration 2 took 1850.1 ms and returned 20 rows
   Query 10 avg time: 2148.49 ms
   ```


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3775084170

   @sqlbenchmark run tpch


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on code in PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#discussion_r2710911652


##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -162,16 +165,17 @@ impl ExecutionPlan for ShuffleReaderExec {
         let max_message_size = config.ballista_grpc_client_max_message_size();
         let force_remote_read = config.ballista_shuffle_reader_force_remote_read();
         let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight();
+        let batch_size = config.batch_size();
 
         if force_remote_read {
             debug!(
-                "All shuffle partitions will be read as remote partitions! To disable this behavior set: `{}=false`",
-                crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ
-            );
+            "All shuffle partitions will be read as remote partitions! To disable this behavior set: `{}=false`",
+            crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ
+        );
         }
 
         log::debug!(
-            "ShuffleReaderExec::execute({task_id}) max_request_num: {max_request_num}, max_message_size: {max_message_size}"
+            "ShuffleReaderExec::execute({task_id}) max_request_num: {max_request_num}, max_message_size: {max_message_size}, batch_size: {batch_size}"

Review Comment:
   make sense, will remove



-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3765215122

   @andygrove if you need verification for benchmarking framework this PR is a good candidate 


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3775430603

   @sqlbenchmark run tpch


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3776126245

   @sqlbenchmark run tpch


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "sqlbenchmark (via GitHub)" <gi...@apache.org>.
sqlbenchmark commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3778181057

   ## Ballista TPC-H Benchmark Results
   
   **PR:** #1380 - fix format and simplify new
   **PR Commit:** `b7be1f7`
   **Base Commit:** `91558a5` (main)
   **Scale Factor:** SF1
   **Iterations:** 3
   
   ### Query Comparison
   
   | Query | Main (ms) | PR (ms) | Change |
   |-------|-----------|---------|--------|
   | Q1 | 817.00 | 783.90 | :white_circle: -4.1% |
   | Q2 | 1319.70 | 1354.50 | :white_circle: +2.6% |
   | Q3 | 991.60 | 1101.70 | :red_circle: +11.1% |
   | Q4 | 831.00 | 877.30 | :red_circle: +5.6% |
   | Q5 | 1745.30 | 1737.40 | :white_circle: -0.5% |
   | Q6 | 401.50 | 401.70 | :white_circle: +0.0% |
   | Q7 | 1825.40 | 1825.70 | :white_circle: +0.0% |
   | Q8 | 2298.90 | 2153.00 | :green_circle: -6.3% |
   | Q9 | 1927.00 | 1940.60 | :white_circle: +0.7% |
   | Q10 | 1237.30 | 1350.00 | :red_circle: +9.1% |
   | Q11 | 1045.90 | 1025.90 | :white_circle: -1.9% |
   | Q12 | 863.00 | 841.40 | :white_circle: -2.5% |
   | Q13 | 917.10 | 928.00 | :white_circle: +1.2% |
   | Q14 | 617.60 | 618.20 | :white_circle: +0.1% |
   | Q15 | 964.10 | 958.40 | :white_circle: -0.6% |
   | Q16 | 1230.30 | 1251.10 | :white_circle: +1.7% |
   | Q17 | 721.30 | 686.10 | :white_circle: -4.9% |
   | Q18 | 1270.00 | 1282.10 | :white_circle: +1.0% |
   | Q19 | 695.90 | 723.10 | :white_circle: +3.9% |
   | Q20 | 978.60 | 1027.20 | :white_circle: +5.0% |
   | Q21 | 1703.00 | 1770.10 | :white_circle: +3.9% |
   | Q22 | 877.40 | 911.90 | :white_circle: +3.9% |
   
   **Total:** Main=25278.90ms, PR=25549.30ms (+1.1%)
   
   ---
   *Automated benchmark run by dfbench*


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3776125973

   Rebased, try benchmark test again


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on code in PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#discussion_r2710922716


##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -1017,10 +1102,177 @@ mod tests {
         .unwrap()
     }
 
+    fn create_custom_test_batch(rows: usize) -> RecordBatch {
+        let schema = create_test_schema();
+
+        // 1. Create number column (0, 1, 2, ..., rows-1)
+        let number_vec: Vec<u32> = (0..rows as u32).collect();
+        let number_array = UInt32Array::from(number_vec);
+
+        // 2. Create string column ("s0", "s1", ..., "s{rows-1}")
+        // Just to fill data, the content is not important
+        let string_vec: Vec<String> = (0..rows).map(|i| format!("s{}", i)).collect();
+        let string_array = StringArray::from(string_vec);
+
+        RecordBatch::try_new(schema, vec![Arc::new(number_array), Arc::new(string_array)])
+            .unwrap()
+    }
+
     fn create_test_schema() -> SchemaRef {
         Arc::new(Schema::new(vec![
             Field::new("number", DataType::UInt32, true),
             Field::new("str", DataType::Utf8, true),
         ]))
     }
+
+    use datafusion::physical_plan::memory::MemoryStream;
+
+    #[tokio::test]
+    async fn test_coalesce_stream_logic() -> Result<()> {
+        // 1. Create test data - 10 small batches, each with 3 rows
+        let schema = create_test_schema();
+        let small_batch = create_test_batch();
+        let batches = vec![small_batch.clone(); 10];
+
+        // 2. Create mock upstream stream (Input Stream)
+        let input_stream = MemoryStream::try_new(batches, schema.clone(), None)?;
+        let input_stream = Box::pin(input_stream) as SendableRecordBatchStream;
+
+        // 3. Configure Coalescer: target batch size to 10 rows
+        let target_batch_size = 10;
+
+        // 4. Manually build the CoalescedShuffleReaderStream
+        let coalesced_stream = CoalescedShuffleReaderStream {

Review Comment:
   Good point, will add, 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: 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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3776637497

   @sqlbenchmark run tpch


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "sqlbenchmark (via GitHub)" <gi...@apache.org>.
sqlbenchmark commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3778302197

   ## Ballista TPC-H Benchmark Results
   
   **PR:** #1380 - fix format and simplify new
   **PR Commit:** `b7be1f7`
   **Base Commit:** `cd4ead8` (main)
   **Scale Factor:** SF10
   **Iterations:** 3
   
   ### Query Comparison
   
   | Query | Main (ms) | PR (ms) | Change |
   |-------|-----------|---------|--------|
   | Q1 | 2621.60 | 2655.50 | :white_circle: +1.3% |
   | Q2 | 1914.20 | 1877.40 | :white_circle: -1.9% |
   | Q3 | 2074.00 | 2112.60 | :white_circle: +1.9% |
   | Q4 | 1248.90 | 1301.00 | :white_circle: +4.2% |
   | Q5 | 4046.90 | 4199.60 | :white_circle: +3.8% |
   | Q6 | 943.00 | 996.90 | :red_circle: +5.7% |
   | Q7 | 5601.90 | 5519.40 | :white_circle: -1.5% |
   | Q8 | 5155.20 | 5128.40 | :white_circle: -0.5% |
   | Q9 | 6406.80 | 6414.00 | :white_circle: +0.1% |
   | Q10 | 2480.70 | 2408.60 | :white_circle: -2.9% |
   | Q11 | 1484.70 | 1418.60 | :white_circle: -4.5% |
   | Q12 | 1843.90 | 1769.50 | :white_circle: -4.0% |
   | Q13 | 2444.90 | 2211.40 | :green_circle: -9.6% |
   | Q14 | 1166.50 | 1065.30 | :green_circle: -8.7% |
   | Q15 | 1601.80 | 1596.00 | :white_circle: -0.4% |
   | Q16 | 1318.30 | 1338.40 | :white_circle: +1.5% |
   | Q17 | 4877.10 | 4296.30 | :green_circle: -11.9% |
   | Q18 | 7092.60 | 6496.30 | :green_circle: -8.4% |
   | Q19 | 1866.20 | 1865.00 | :white_circle: -0.1% |
   | Q20 | 1831.60 | 1742.50 | :white_circle: -4.9% |
   | Q21 | 6209.10 | 5693.70 | :green_circle: -8.3% |
   | Q22 | 1087.00 | 1149.00 | :red_circle: +5.7% |
   
   **Total:** Main=65316.90ms, PR=63255.40ms (-3.2%)
   
   ---
   *Automated benchmark run by dfbench*


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3758307153

   cc @milenkovicm 


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3760573957

   thanks @danielhumanmod will have a look asap 


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#discussion_r2702347954


##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -162,16 +165,17 @@ impl ExecutionPlan for ShuffleReaderExec {
         let max_message_size = config.ballista_grpc_client_max_message_size();
         let force_remote_read = config.ballista_shuffle_reader_force_remote_read();
         let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight();
+        let batch_size = config.batch_size();
 
         if force_remote_read {
             debug!(
-                "All shuffle partitions will be read as remote partitions! To disable this behavior set: `{}=false`",
-                crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ
-            );
+            "All shuffle partitions will be read as remote partitions! To disable this behavior set: `{}=false`",
+            crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ
+        );
         }
 
         log::debug!(
-            "ShuffleReaderExec::execute({task_id}) max_request_num: {max_request_num}, max_message_size: {max_message_size}"
+            "ShuffleReaderExec::execute({task_id}) max_request_num: {max_request_num}, max_message_size: {max_message_size}, batch_size: {batch_size}"

Review Comment:
   do we need batch size here? its part of configuration, we can find out if needed 



##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -197,11 +201,22 @@ impl ExecutionPlan for ShuffleReaderExec {
             prefer_flight,
         );
 
-        let result = RecordBatchStreamAdapter::new(
-            Arc::new(self.schema.as_ref().clone()),
+        let input_stream = Box::pin(RecordBatchStreamAdapter::new(
+            self.schema.clone(),
             response_receiver.try_flatten(),
-        );
-        Ok(Box::pin(result))
+        ));
+
+        Ok(Box::pin(CoalescedShuffleReaderStream {
+            schema: self.schema.clone(),
+            input: input_stream,
+            coalescer: LimitedBatchCoalescer::new(

Review Comment:
   maybe we could add limit to ShuffleReader, defaulting to None, and add `with_limit` option, so if we find a use-case we can set it ? 



##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -197,11 +201,22 @@ impl ExecutionPlan for ShuffleReaderExec {
             prefer_flight,
         );
 
-        let result = RecordBatchStreamAdapter::new(
-            Arc::new(self.schema.as_ref().clone()),
+        let input_stream = Box::pin(RecordBatchStreamAdapter::new(
+            self.schema.clone(),
             response_receiver.try_flatten(),
-        );
-        Ok(Box::pin(result))
+        ));
+
+        Ok(Box::pin(CoalescedShuffleReaderStream {
+            schema: self.schema.clone(),
+            input: input_stream,
+            coalescer: LimitedBatchCoalescer::new(

Review Comment:
   what is the reason using `LimitedBatchCoalescer` instead of `BatchCoalescer` as there is no fetch limit?
   
   would there be a case where we can use limit 🤔?



##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -1017,10 +1102,177 @@ mod tests {
         .unwrap()
     }
 
+    fn create_custom_test_batch(rows: usize) -> RecordBatch {
+        let schema = create_test_schema();
+
+        // 1. Create number column (0, 1, 2, ..., rows-1)
+        let number_vec: Vec<u32> = (0..rows as u32).collect();
+        let number_array = UInt32Array::from(number_vec);
+
+        // 2. Create string column ("s0", "s1", ..., "s{rows-1}")
+        // Just to fill data, the content is not important
+        let string_vec: Vec<String> = (0..rows).map(|i| format!("s{}", i)).collect();
+        let string_array = StringArray::from(string_vec);
+
+        RecordBatch::try_new(schema, vec![Arc::new(number_array), Arc::new(string_array)])
+            .unwrap()
+    }
+
     fn create_test_schema() -> SchemaRef {
         Arc::new(Schema::new(vec![
             Field::new("number", DataType::UInt32, true),
             Field::new("str", DataType::Utf8, true),
         ]))
     }
+
+    use datafusion::physical_plan::memory::MemoryStream;
+
+    #[tokio::test]
+    async fn test_coalesce_stream_logic() -> Result<()> {
+        // 1. Create test data - 10 small batches, each with 3 rows
+        let schema = create_test_schema();
+        let small_batch = create_test_batch();
+        let batches = vec![small_batch.clone(); 10];
+
+        // 2. Create mock upstream stream (Input Stream)
+        let input_stream = MemoryStream::try_new(batches, schema.clone(), None)?;
+        let input_stream = Box::pin(input_stream) as SendableRecordBatchStream;
+
+        // 3. Configure Coalescer: target batch size to 10 rows
+        let target_batch_size = 10;
+
+        // 4. Manually build the CoalescedShuffleReaderStream
+        let coalesced_stream = CoalescedShuffleReaderStream {

Review Comment:
   i believe adding 'new' method would make sense 



-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3775448540

   @milenkovicm the benchmark script failed because this PR does not contain the changes from https://github.com/apache/datafusion-ballista/pull/1391


-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "danielhumanmod (via GitHub)" <gi...@apache.org>.
danielhumanmod commented on code in PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#discussion_r2710920809


##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -197,11 +201,22 @@ impl ExecutionPlan for ShuffleReaderExec {
             prefer_flight,
         );
 
-        let result = RecordBatchStreamAdapter::new(
-            Arc::new(self.schema.as_ref().clone()),
+        let input_stream = Box::pin(RecordBatchStreamAdapter::new(
+            self.schema.clone(),
             response_receiver.try_flatten(),
-        );
-        Ok(Box::pin(result))
+        ));
+
+        Ok(Box::pin(CoalescedShuffleReaderStream {
+            schema: self.schema.clone(),
+            input: input_stream,
+            coalescer: LimitedBatchCoalescer::new(

Review Comment:
   > maybe we could add limit to ShuffleReader, defaulting to None, and add `with_limit` option, so if we find a use-case we can set it ?
   
   Sounds like a good point. One use case I can think of is optimizations like Limit Pushdown to the shuffle read stage



-- 
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] feat: Add batch coalescing ability to shuffle reader exec [datafusion-ballista]

Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1380:
URL: https://github.com/apache/datafusion-ballista/pull/1380#issuecomment-3776307151

   @sqlbenchmark run tpch


-- 
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