Skip to content

Metrics points and tags/dyn filter enabled - #15

Closed
gene-bordegaray wants to merge 99 commits into
mainfrom
metrics-points-and-tags/dyn-filter-enabled
Closed

Metrics points and tags/dyn filter enabled#15
gene-bordegaray wants to merge 99 commits into
mainfrom
metrics-points-and-tags/dyn-filter-enabled

Conversation

@gene-bordegaray

Copy link
Copy Markdown

enable dyn filtering

gabotechs and others added 30 commits October 17, 2025 17:12
This change adds support for displaying a distributed EXPLAIN ANALYZE output with metrics. It updates the TPCH
validation tests to assert the EXPLAIN ANALYZE output for each query. I collected the single node
results [here](https://github.com/datafusion-contrib/datafusion-distributed/blob/js/full-explain-analyze-rebased-comparison/tests/tpch_validation_test.rs) - using the `output_rows` metric values, we can cross check
that the distributed plan metrics are correct.

Implemenation notes:
- Adds `src/explain.rs` to stores the main entrypoint to rendering the output string
- We now re-write the whole plan before rendering using `rewrite_distributed_plan_with_metrics()` which
  is a new public API that takes an executed `DistributedExec` and wraps each node in the appropriate metrics
  - A user can call this method on an executed plan and traverse it to collect metrics from particular nodes
    (This may be hard though because all nodes are wrapped in MetricsWrapperExec...)
- Adds a `show_metrics: bool` param to `display_plan_ascii`
- Significantly refactors TaskMetricsRewriter -> StageMetricsWriter. See comment.

Informs: datafusion-contrib#123

Other follow up work:
- datafusion-contrib#185
- datafusion-contrib#184
- datafusion-contrib#188
- datafusion-contrib#189
- datafusion-contrib#190
…rib#194)

This change makes the explain output more concise by changing the stage formatting from
```
Tasks: t0:[p0,p1,p2,p3,p4,p5] t1:[p0,p1,p2,p3,p4,p5] t2:[p0,p1,p2,p3,p4,p5] t3:[p0,p1,p2,p3,p4,p5]
```
to
```
Tasks: t0:[p0..p5] t1:[p0..p5] t2:[p0..p5] t3:[p0..p5]
```

Closes: datafusion-contrib#192
…sion-contrib#197)

* Refactor distributed planner into its own folder

* Tweak tests so that they use SET statement for configuring tasks.

* make config fields public

* Allow users to call PlanDependantUsize

* Fix typo

* Address conflicts
* Add arrow flight endpoint hooks

* Store multiple hooks
Before this change, `name()` and `as_any()` on `MetricsWrapperExec` did not delegate calls
to the wrapped node. This was a problem because a user could not traverse the plan and
extract metrics from specific nodes.

After this change, those methods are delegated so a user can now extract metrics
from specific nodes.
* remove schema adapter

* garbage collect dictionaries before sending them across the wire

* move into do_get, handle view arrays

* update comments

* Update src/flight_service/do_get.rs

Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com>

---------

Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com>
* allow configuring message sizes

* actually use

* fmt

* use usize::MAX and add helpers to create clients/servers
…fusion-contrib#208)

* Add docs to `start_localhost_context`

* Add test for nullification bug

* Fix typo

* Add test warning
* Struct is public only within crate.

* Declare pub(crate) in mod.rs.
* Add instructions to readme.

* Update comment following Gabriel's feedback.
* Add tests for NetworkCoalesceExec.

* Add standalone NetworkCoalesceExec test and test with n_tasks > 1.
…sion-contrib#218)

* Relax timing constraints.

* Remove unnecessary changes. Add minor fixes.
* Introduce with_distributed_execution method and remove DistributedExt from anything that's not a SessionStateBuilder

* Rework how tasks are assigned to stages

* Use files_per_task and cardinality_effect_factor for choosing the right tasks for the different stages

* Update all tests to use the new task estimation mechanism

* Fix when worker urls is zero or one and add tests

* Move back to `with_distributed_channel_resolver` and let users inject either `DistributedPhysicalOptimizerRule` or their own rules

* Add files_per_task and cardinality_task_sf to examples

* Improve docs

* add impl TaskEstimator for Arc<&dyn TaskEstimator>

* Add task estimator tests

* Explain why we need to count distinct files

* Return an error if task_count == 0

* Factor out apply_scale_factor

* Add task estimator doc comments

* Add more docs to apply_network_boundaries

* Extend files_per_task docs

* Rollback DistributedExt de-implementation from SessionConfig, SessionState and SessionContext
…fusion-contrib#225)

This commit adds two small changes to improve how we parse TTLMapConfigs.

1. It strictly enforces that `tick` is an integer multiple of `ttl`. Previously,
we only recommened that in comments, with the consequence that:

```
true_ttl = desired_ttl - (desired_ttl % tick) <= desired_ttl
```
In the worst case, if `tick` is close to `ttl / 2` then the true `ttl` would
be equivalent to `tick`. For example, if the requested `ttl` was 60 seconds, and the tick
was 31 seconds, we would only allocate a single list of inputs to delete and free them every
tick.

2. There was a small bug in the validation logic that did not correctly validate that tick
was nonzero due to an incorrect application of DeMorgan's law. This commit fixes that to
correctly check that `tick` is nonzero.
…on-contrib#224)

Fixes datafusion-contrib#221. Previously, we encountered a bug at the initial gc tick on
ttl_map where `time` (denoting the number of ticks since map initialization)
starts at 0, but the logic to compute the time to remove the inserted value is

```
free_time = (time - 1) // (number of ticks to ttl)
```

In other words, we maintain a circular buffer of inserts, with a single slot per tick,
and we will delete the item once `buffer.len()` ticks have elapsed (ttl = tick * buffer.len()).

This is all fine, except we implement the logic as:

```
free_time = time.wrapping_sub(1) / buffer.len()
```

which is computing for `u64` `time`:

```
free_time = ((time - 1) mod 2^64) mod buffer.len()
```

when we really just want `free_time = (time - 1) mod buffer.len()`.

Luckily this equality holds as long as ` 0 < time < 2^64`, as for those times `time - 1 mod 2&64 = time`.
This commit changes our behavior to initialize `time` at 1 instead of 0. We don't need to worry about
the overflow case because even if the tick duration was a nanosecond, it would take ~584 years to overflow
64 bits at which point we would surely have other problems besides the momentarily incorrect ttl. For
the underflow case, this should primarily help with unit tests above anything else, as the bug
only happened at `time = 0`.
* Add cdk sample app

* Add bucket and ec2 instances

* Use a static channel resolver in benchmarks

* Add worker.rs

* Use an ec2 api based channel resolver

* Revert "Use a static channel resolver in benchmarks"

This reverts commit 5bdec16

* Refactor ec2 channel resolver

* Add datafusion-bench script

* Fix worker ec2 resolver

* Improve benchmarking script and cdk tests

* Show diff with previous run

* Better logs for the benchmarks

* Improve readme

* Allays drop tables

* Improve Cfn output
* Improve TPCH gen command

* Add comment about default repartition behavior
This change frees up disk space in CI runners. See https://github.com/apache/flink/blob/02d30ace69dc18555a5085eccf70ee884e73a16e/tools/azure-pipelines/free_disk_space.sh
Before
```
Filesystem      Size  Used Avail Use% Mounted on
/dev/root        72G   66G  5.8G  93% /
tmpfs           7.9G   84K  7.9G   1% /dev/shm
tmpfs           3.2G  1.2M  3.2G   1% /run
tmpfs           5.0M     0  5.0M   0% /run/lock
/dev/sda16      881M   62M  758M   8% /boot
/dev/sda15      105M  6.2M   99M   6% /boot/efi
/dev/sdb1        74G  4.1G   66G   6% /mnt
tmpfs           1.6G   12K  1.6G   1% /run/user/1001
```
After
```
Filesystem      Size  Used Avail Use% Mounted on
/dev/root        72G   47G   25G  66% /
tmpfs           7.9G   84K  7.9G   1% /dev/shm
tmpfs           3.2G  1.1M  3.2G   1% /run
tmpfs           5.0M     0  5.0M   0% /run/lock
/dev/sdb16      881M   62M  758M   8% /boot
/dev/sdb15      105M  6.2M   99M   6% /boot/efi
/dev/sda1        74G  4.1G   66G   6% /mnt
tmpfs           1.6G   12K  1.6G   1% /run/user/1001
```
…sion-contrib#232)

* Split tpch tests

* Fix correctness tests

* Fix plan tests
* Remove s3 download step

* Add user commands for Rust tooling
…trib#242)

* Add CoalesceBatchExec in network shuffles

* Add param to CDK benchmarks

* Remove redundant batch coalescing

* Add comment for the future
JSOD11 and others added 29 commits January 21, 2026 21:04
…usion-contrib#256)

* Added join testdata and method to generate parquet from csv.

* Added basic single node join test for comparison.

* Use hive partitioning.

* Results are the same with hive partitioning and Gene's PR.

* Fixed configs and achieved optimal distributed plan.

* Refactoring, adding comments.

* Added check to ensure optimal plan is achieved.

* Update based on Nga's comments.

* Added second test.

* Added third test.

* Refactor.

* Add ORDER BY to queries and switch to snapshot testing.

* Nulls last instead of nulls first.

* Fix tests.
…datafusion-contrib#305)

Add proto serialization support for four additional DataFusion metric types
that were previously filtered out during metrics collection. This enables
these metrics to be transferred across network boundaries in distributed
query execution and displayed in explain analyze output.

Closes datafusion-contrib#237
* Clippy fixes.

* Run clippy for all features.
Closes datafusion-contrib#223 

last one got closed and I can't reopen

Added a separate exexution plan node `BroadcastExec` to handle the
partition caching logic that previously lived in the flight service.

Still get cool speed-ups 😄 

<img width="319" height="888" alt="Screenshot 2026-01-02 at 7 53 39 PM"
src="https://github.com/user-attachments/assets/0757c310-5f1a-4bf7-970d-066abf339e9d"
/>

---------

Co-authored-by: Gabriel Musat Mestre <gabriel.musatmestre@datadoghq.com>
Adds some metrics to the network boundaries so that they can be
displayed in the plans. They look like this:

```
┌───── DistributedExec ── Tasks: t0:[p0] 
│ ProjectionExec: expr=[count(*)@0 as count(*), RainToday@1 as RainToday], metrics=[output_rows=2, elapsed_compute=1.67µs, output_bytes=48.0 B, output_batches=1, expr_0_eval_time=167ns, expr_1_eval_time=84ns]
│   SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST], metrics=[output_rows=2, elapsed_compute=32.29µs, output_bytes=64.0 B, output_batches=1]
│     [Stage 2] => NetworkCoalesceExec: output_partitions=6, input_tasks=2, metrics=[elapsed_compute=1.34ms, bytes_transferred=19.28 K, max_mem_used=1.96 K]
└──────────────────────────────────────────────────
  ┌───── Stage 2 ── Tasks: t0:[p0..p2] t1:[p0..p2] 
  │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true], metrics=[output_rows=2, elapsed_compute=18.71µs, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0]
  │   ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))], metrics=[output_rows=2, elapsed_compute=9.55µs, expr_0_eval_time=421ns, expr_1_eval_time=129ns, expr_2_eval_time=128ns]
  │     AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))], metrics=[output_rows=2, elapsed_compute=208.66µs, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, peak_mem_used=100.9 K, aggregate_arguments_time=21.21µs, aggregation_time=16.34µs, emitting_time=5.17µs, time_calculating_group_ids=30.05µs]
  │       [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=3, metrics=[elapsed_compute=2.08ms, bytes_transferred=19.72 K, max_mem_used=5.19 K]
  └──────────────────────────────────────────────────
    ┌───── Stage 1 ── Tasks: t0:[p0..p5] t1:[p0..p5] t2:[p0..p5] 
    │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=1, metrics=[]
    │   AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))], metrics=[]
    │     PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] , metrics=[]
    │       DataSourceExec: file_groups={3 groups: [[Users/gabriel.musatmestre/github/datafusion-distributed/testdata/weather/result-000000.parquet], [Users/gabriel.musatmestre/github/datafusion-distributed/testdata/weather/result-000001.parquet], [Users/gabriel.musatmestre/github/datafusion-distributed/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet, metrics=[]
    └──────────────────────────────────────────────────
```

It just uses the existing mechanism with very small modifications,
shoutout to @jayshrivastava for the magnificent metrics collection
system.
This PR improves DX for working with benchmarks in this project. It does
several things:

- Allows storing individual query results in individual files, so that
running just a partial set of queries stores some results that can
co-live with previous runs for other queries.
- Adds comparison commands, so that we can compare with previous runs
without necessarily rerunning the benchmarks.
- Stores results per branch (or per engine in the remote benchmarks).
That way, we do not need to rerun the base branch just so that we can
compare to it. This is essential when running big benchmarks in remote
mode, as they take a lot of time, and we don't want to be running them
over an over just so that we can compare them.
- Refactors the Trino deployment to be more reliable and configurable.
- Adds a `built` dependency so that we have compile time information
that can be served by a remote worker (commit hash, build time, etc...)

All of this will make it easier to work with benchmarks in the future.
…b#297)

Added tonic-based observability gRPC codegen for console process.
Updated worker/server wiring to include observability service, updated
console to use client gRPC and render ping status in the TUI (to be
extended to stats and metrics for cluster/workers later).
This change adds a new enum `DistributedMetricsFormat` which is either
`Aggregated` or `PerTask`.

`Aggregated` is the default. When we use
`DisplayableExecutionPlan::with_metrics`, metrics are aggregated by
name. So we leave metrics as is before displaying.

`PerTask` will configure the metrics re-writer to label metrics like
`output_rows` with `__distributed_datafusion__task_id=x` at re-write
time. At display time, we do 3 things to the `MetricSet` of a node:
1. Aggregate by name, task_id
2. Sort by name, task_id
3. Display metrics in the format `output_rows_x` where `x` is the task
id.

If you look at what `DisplayableExecutionPlan::with_metrics` does, it
does the exact same 3 steps above, except it only aggregates by name.
The alternative approach is to rename metrics like in
datafusion-contrib#307. I
felt that this approach was more datafusion-naitive.

I also considered labelling metrics at collection-time, but I figured
it's not worth the effort.

Sample output:

1. `DistributedMetricsFormat::Aggregated`
```
┌───── DistributedExec ── Tasks: t0:[p0]
│ ProjectionExec: expr=[count(*)@0 as count(*), RainToday@1 as RainToday], metrics=[output_rows=2, elapsed_compute=1.92µs, output_bytes=48.0 B, output_batches=1, expr_0_eval_time=208ns, expr_1_eval_time=125ns]
│   SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST], metrics=[output_rows=2, elapsed_compute=10.83µs, output_bytes=64.0 B, output_batches=1]
│     [Stage 2] => NetworkCoalesceExec: output_partitions=6, input_tasks=2, metrics=[]
└──────────────────────────────────────────────────
  ┌───── Stage 2 ── Tasks: t0:[p0..p2] t1:[p0..p2]
  │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true], metrics=[output_rows=2, elapsed_compute=10.80µs, output_bytes=0.0 B, output_batches=0, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0]
  │   ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))], metrics=[output_rows=2, elapsed_compute=5.09µs, output_bytes=32.1 KB, output_batches=2, expr_0_eval_time=462ns,
expr_1_eval_time=254ns, expr_2_eval_time=170ns]
  │     AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))], metrics=[output_rows=2, elapsed_compute=165.50µs, output_bytes=32.1 KB, output_batches=2, spill_count=0, spilled_bytes=0.0 B, spilled_
rows=0, peak_mem_used=100.9 K, aggregate_arguments_time=9.38µs, aggregation_time=9.21µs, emitting_time=5.34µs, time_calculating_group_ids=18.75µs]
  │       [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=3, metrics=[]
  └──────────────────────────────────────────────────
    ┌───── Stage 1 ── Tasks: t0:[p0..p5] t1:[p0..p5] t2:[p0..p5]
    │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=1, metrics=[output_rows=6, elapsed_compute=52.26µs, output_bytes=1152.0 KB, output_batches=6, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, fetch_time=10
.72ms, repartition_time=39.67µs, send_time=13.72µs]
    │   AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))], metrics=[output_rows=6, elapsed_compute=192.54µs, output_bytes=48.1 KB, output_batches=3, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, s
kipped_aggregation_rows=0, peak_mem_used=56.47 K, aggregate_arguments_time=10.42µs, aggregation_time=7.58µs, emitting_time=17.00µs, time_calculating_group_ids=81.88µs, reduction_factor=1.6% (6/366)]
    │     PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] , metrics=[]
    │       DataSourceExec: file_groups={3 groups: [[Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000000.parquet], [Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000001.parq
uet], [Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet, metrics=[output_rows=366, elapsed_compute=3ns, output_bytes=8.2 KB, output_batches=3, file
s_ranges_pruned_statistics=3 total → 3 matched, row_groups_pruned_statistics=3 total → 3 matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_rows_pruned=0 total → 0 matched, batches_split=0, bytes_scanned=3.12 K, fil
e_open_errors=0, file_scan_errors=0, num_predicate_creation_errors=0, predicate_cache_inner_records=0, predicate_cache_records=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, bloom_filter_eval_time=6ns
, metadata_load_time=5.62ms, page_index_eval_time=6ns, row_pushdown_eval_time=6ns, statistics_eval_time=6ns, time_elapsed_opening=5.79ms, time_elapsed_processing=1.51ms, time_elapsed_scanning_total=4.77ms, time_elapsed_scanning_until_da
ta=4.54ms, scan_efficiency_ratio=2.3% (3.12 K/137.4 K)]
    └──────────────────────────────────────────────────
```

2. `DistributedMetricsFormat::PerTask`
```
───── DistributedExec ── Tasks: t0:[p0]
│ ProjectionExec: expr=[count(*)@0 as count(*), RainToday@1 as RainToday], metrics=[output_batches_0=1, output_bytes_0=48, output_rows_0=2, end_timestamp_0=1769118.1 T, start_timestamp_0=1769118.1 T, elapsed_compute_0=1.58µs, expr_0_eva
l_time_0=167ns, expr_1_eval_time_0=83ns]
│   SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST], metrics=[output_batches_0=1, output_bytes_0=64, output_rows_0=2, end_timestamp_0=1769118.1 T, start_timestamp_0=1769118.1 T, elapsed_compute_0=27.04µs]
│     [Stage 2] => NetworkCoalesceExec: output_partitions=6, input_tasks=2, metrics=[]
└──────────────────────────────────────────────────
  ┌───── Stage 2 ── Tasks: t0:[p0..p2] t1:[p0..p2]
  │ SortExec: expr=[count(*)@0 ASC NULLS LAST], preserve_partitioning=[true], metrics=[output_batches_0=0, output_batches_1=0, output_bytes_0=0, output_bytes_1=0, output_rows_0=1, output_rows_1=1, spill_count_0=0, spill_count_1=0, spill
ed_bytes_0=0, spilled_bytes_1=0, spilled_rows_0=0, spilled_rows_1=0, end_timestamp_0=5307354.2 T, end_timestamp_1=5307354.2 T, start_timestamp_0=5307354.2 T, start_timestamp_1=5307354.2 T, elapsed_compute_0=8.42µs, elapsed_compute_1=4.9
2µs]
  │   ProjectionExec: expr=[count(Int64(1))@1 as count(*), RainToday@0 as RainToday, count(Int64(1))@1 as count(Int64(1))], metrics=[output_batches_0=1, output_batches_1=1, output_bytes_0=16.42 K, output_bytes_1=16.42 K, output_rows_0=1
, output_rows_1=1, end_timestamp_0=5307354.2 T, end_timestamp_1=5307354.2 T, start_timestamp_0=5307354.2 T, start_timestamp_1=5307354.2 T, elapsed_compute_0=1.83µs, elapsed_compute_1=1.96µs, expr_0_eval_time_0=168ns, expr_0_eval_time_1=
168ns, expr_1_eval_time_0=43ns, expr_1_eval_time_1=85ns, expr_2_eval_time_0=86ns, expr_2_eval_time_1=85ns]
  │     AggregateExec: mode=FinalPartitioned, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))], metrics=[output_batches_0=1, output_batches_1=1, output_bytes_0=16.42 K, output_bytes_1=16.42 K, output_rows_0=1, output_rows_1=1, spi
ll_count_0=0, spill_count_1=0, spilled_bytes_0=0, spilled_bytes_1=0, spilled_rows_0=0, spilled_rows_1=0, end_timestamp_0=5307354.2 T, end_timestamp_1=5307354.2 T, peak_mem_used_0=50.47 K, peak_mem_used_1=50.47 K, start_timestamp_0=53073
54.2 T, start_timestamp_1=5307354.2 T, aggregate_arguments_time_0=3.17µs, aggregate_arguments_time_1=2.59µs, aggregation_time_0=3.13µs, aggregation_time_1=2.04µs, elapsed_compute_0=93.29µs, elapsed_compute_1=57.67µs, emitting_time_0=2.5
0µs, emitting_time_1=2.04µs, time_calculating_group_ids_0=7.75µs, time_calculating_group_ids_1=4.67µs]
  │       [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=3, metrics=[]
  └──────────────────────────────────────────────────
    ┌───── Stage 1 ── Tasks: t0:[p0..p5] t1:[p0..p5] t2:[p0..p5]
    │ RepartitionExec: partitioning=Hash([RainToday@0], 6), input_partitions=1, metrics=[output_batches_0=2, output_batches_1=2, output_batches_2=2, output_bytes_0=393.2 K, output_bytes_1=393.2 K, output_bytes_2=393.2 K, output_rows_0=2
, output_rows_1=2, output_rows_2=2, spill_count_0=0, spill_count_1=0, spill_count_2=0, spilled_bytes_0=0, spilled_bytes_1=0, spilled_bytes_2=0, spilled_rows_0=0, spilled_rows_1=0, spilled_rows_2=0, end_timestamp_0=10614708.4 T, end_time
stamp_1=10614708.4 T, end_timestamp_2=10614708.4 T, start_timestamp_0=10614708.4 T, start_timestamp_1=10614708.4 T, start_timestamp_2=10614708.4 T, elapsed_compute_0=10.42µs, elapsed_compute_1=11.29µs, elapsed_compute_2=19.63µs, fetch_t
ime_0=3.92ms, fetch_time_1=3.35ms, fetch_time_2=4.08ms, repartition_time_0=7.21µs, repartition_time_1=14.33µs, repartition_time_2=20.21µs, send_time_0=2.63µs, send_time_1=3.63µs, send_time_2=5.46µs]
    │   AggregateExec: mode=Partial, gby=[RainToday@0 as RainToday], aggr=[count(Int64(1))], metrics=[output_batches_0=1, output_batches_1=1, output_batches_2=1, output_bytes_0=16.42 K, output_bytes_1=16.42 K, output_bytes_2=16.42 K, ou
tput_rows_0=2, output_rows_1=2, output_rows_2=2, skipped_aggregation_rows_0=0, skipped_aggregation_rows_1=0, skipped_aggregation_rows_2=0, spill_count_0=0, spill_count_1=0, spill_count_2=0, spilled_bytes_0=0, spilled_bytes_1=0, spilled_
bytes_2=0, spilled_rows_0=0, spilled_rows_1=0, spilled_rows_2=0, end_timestamp_0=1769118.1 T, end_timestamp_1=1769118.1 T, end_timestamp_2=1769118.1 T, peak_mem_used_0=18.82 K, peak_mem_used_1=18.82 K, peak_mem_used_2=18.82 K, reduction
_factor_0=1, reduction_factor_1=1, reduction_factor_2=1, start_timestamp_0=1769118.1 T, start_timestamp_1=1769118.1 T, start_timestamp_2=1769118.1 T, aggregate_arguments_time_0=2.83µs, aggregate_arguments_time_1=5.17µs, aggregate_argume
nts_time_2=7.21µs, aggregation_time_0=2.00µs, aggregation_time_1=3.42µs, aggregation_time_2=5.12µs, elapsed_compute_0=54.71µs, elapsed_compute_1=71.75µs, elapsed_compute_2=85.38µs, emitting_time_0=3.38µs, emitting_time_1=6.83µs, emittin
g_time_2=8.33µs, time_calculating_group_ids_0=25.96µs, time_calculating_group_ids_1=30.79µs, time_calculating_group_ids_2=34.00µs]
    │     PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0] , metrics=[]
    │       DataSourceExec: file_groups={3 groups: [[Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000000.parquet], [Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000001.parq
uet], [Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000002.parquet]]}, projection=[RainToday], file_type=parquet, metrics=[batches_split_0=0, batches_split_1=0, batches_split_2=0, bytes_scanned_0=1.07 K,
bytes_scanned_1=1.01 K, bytes_scanned_2=1.03 K, file_open_errors_0=0, file_open_errors_1=0, file_open_errors_2=0, file_scan_errors_0=0, file_scan_errors_1=0, file_scan_errors_2=0, files_ranges_pruned_statistics_matched_0=1, files_ranges
_pruned_statistics_matched_1=1, files_ranges_pruned_statistics_matched_2=1, num_predicate_creation_errors_0=0, num_predicate_creation_errors_1=0, num_predicate_creation_errors_2=0, output_batches_0=1, output_batches_1=1, output_batches_
2=1, output_bytes_0=2.81 K, output_bytes_1=2.79 K, output_bytes_2=2.80 K, output_rows_0=122, output_rows_1=122, output_rows_2=122, page_index_rows_pruned_matched_0=0, page_index_rows_pruned_matched_1=0, page_index_rows_pruned_matched_2=
0, predicate_cache_inner_records_0=0, predicate_cache_inner_records_1=0, predicate_cache_inner_records_2=0, predicate_cache_records_0=0, predicate_cache_records_1=0, predicate_cache_records_2=0, predicate_evaluation_errors_0=0, predicat
e_evaluation_errors_1=0, predicate_evaluation_errors_2=0, pushdown_rows_matched_0=0, pushdown_rows_matched_1=0, pushdown_rows_matched_2=0, pushdown_rows_pruned_0=0, pushdown_rows_pruned_1=0, pushdown_rows_pruned_2=0, row_groups_pruned_b
loom_filter_matched_0=1, row_groups_pruned_bloom_filter_matched_1=1, row_groups_pruned_bloom_filter_matched_2=1, row_groups_pruned_statistics_matched_0=1, row_groups_pruned_statistics_matched_1=1, row_groups_pruned_statistics_matched_2=
1, end_timestamp_0=1769118.1 T, end_timestamp_1=1769118.1 T, end_timestamp_2=1769118.1 T, scan_efficiency_ratio_0=2, scan_efficiency_ratio_1=2, scan_efficiency_ratio_2=2, start_timestamp_0=1769118.1 T, start_timestamp_1=1769118.1 T, sta
rt_timestamp_2=1769118.1 T, bloom_filter_eval_time_0=2ns, bloom_filter_eval_time_1=2ns, bloom_filter_eval_time_2=2ns, elapsed_compute_0=1ns, elapsed_compute_1=1ns, elapsed_compute_2=1ns, metadata_load_time_0=2.92ms, metadata_load_time_1
=2.77ms, metadata_load_time_2=3.29ms, page_index_eval_time_0=2ns, page_index_eval_time_1=2ns, page_index_eval_time_2=2ns, row_pushdown_eval_time_0=2ns, row_pushdown_eval_time_1=2ns, row_pushdown_eval_time_2=2ns, statistics_eval_time_0=2
ns, statistics_eval_time_1=2ns, statistics_eval_time_2=2ns, time_elapsed_opening_0=2.94ms, time_elapsed_opening_1=2.82ms, time_elapsed_opening_2=3.33ms, time_elapsed_processing_0=374.50µs, time_elapsed_processing_1=432.71µs, time_elapse
d_processing_2=485.12µs, time_elapsed_scanning_total_0=929.29µs, time_elapsed_scanning_total_1=444.71µs, time_elapsed_scanning_total_2=699.75µs, time_elapsed_scanning_until_data_0=866.62µs, time_elapsed_scanning_until_data_1=364.79µs, t
ime_elapsed_scanning_until_data_2=598.04µs]
    └──────────────────────────────────────────────────
```
Closes
datafusion-contrib#304
Adds fast deploys for the benchmarking cluster, this takes the iteration
time from ~3mins to 30 seconds for deploying changes to all ec2 machines
in the cluster.
Fix for merged PR [Add functionality for worker-console communication
datafusion-contrib#297

](datafusion-contrib#297)
where generated protobuf code would end up in
`src/observability/src/generated/` instead of
`src/observability/generated/`.
…atafusion-contrib#314)

Dupe of
datafusion-contrib#281
with minor fixes

- Update test assertions to allow empty metrics for
`PartitionIsolatorExec` nodes
- Fix `assert_metrics_present_in_plan` test helper to handle
`PartitionIsolatorExec` wrapped in `MetricsWrapperExec`
- Fix `count_plan_nodes_up_to_network_boundary` test helper now that
network boundary nodes have metrics

Fixes issue where PartitionIsolatorExec nodes were missing from
displayed plans after metrics rewriting. All metrics collection tests
now pass.

Co-authored-by: Dinesh <dineshmadhavan13@gmail.com>
…contrib#321)

The previous `t3.2xlarge` where "burstable" instances, that are
optimized to be economic at the price of not being able to use all the
CPUs for a sustained amount of time.

When running big benchmarks (TPC-H 100), the latency measurements show
huge difference between runs, not only on DataFusion, but on Trino and
Spark, which doesn't make a good reliable environment for benchmarks.

These new instance types are the bare minimum to run TPC-H 100 on Trino
(very resource hungry) while maintaining good consistency between
benchmark runs.
Before, if a stream was abandoned, we were leaking the `TaskData` until
the TTLMap decides to clean it up due to a timeout.

In this PR:
- We also remove the `StageKey` from the TTLMap in the event of a drop
(not just a graceful completion)
- We are faster at cancelling the background threads that poll data
Just some small tweaks to the git branch names in the benchmarks.
**Problem 1:**
The rpcds benchmark would only let you generate data if you were in the
"./benchmarks" dir. Changed so that you can do it from anywhere

before:
<img width="748" height="420" alt="Screenshot 2026-01-28 at 6 49 44 AM"
src="https://github.com/user-attachments/assets/144c3c98-3082-4d59-bdc2-918a6aee9803"
/>

after:
<img width="755" height="233" alt="Screenshot 2026-01-28 at 6 51 03 AM"
src="https://github.com/user-attachments/assets/6619b59c-1d5a-46d8-8ac6-459e83b88b13"
/>

**Problem 2:**
gen-tpch.sh was fialing to convert to parquet because the csv reader
expected less columns than there actually was. This was because it did
not properly take into account its empty trailing field.

before:
<img width="759" height="622" alt="Screenshot 2026-01-28 at 6 57 03 AM"
src="https://github.com/user-attachments/assets/6ed0f3ea-b5eb-45b5-9985-06b58387b929"
/>

after:
<img width="755" height="514" alt="Screenshot 2026-01-28 at 6 55 32 AM"
src="https://github.com/user-attachments/assets/8ff09fb4-8352-47eb-8a9a-b76f8a685548"
/>

**Problem 3:**
Docs were out of date and added some more info
datafusion-contrib#256
introduced by mistake an upgrade to the arrow library, which seems to be
giving problems for some folks who are not yet prepared for arrow 57.2.0
and need to remain in 57.1.0.

As there's no reason that this library restricts the arrow version to
57.2.0, this PR rolls it back to relax the requirement to 57.1.0
Previously, metrics were displaying like this:
```
  │   ProjectionExec: expr=[Rainfall@0 as MinTemp, RainToday@1 as RainToday]
  │, metrics=[output_rows_2=33, elapsed_compute_2=4.13µs, output_bytes_2=576.0 KB, output_batches_2=3, expr_0_eval_time_2=541ns, expr_1_eval_time_2=291ns]
  │     FilterExec: Rainfall@0 > 5
  │, metrics=[output_rows_2=33, elapsed_compute_2=38.79µs, output_bytes_2=576.0 KB, output_batches_2=3, selectivity_2=14% (33/234)o# Please enter the commit message for your changes. Lines starting
```

This change removes the extra new line
```
  │   ProjectionExec: expr=[Rainfall@0 as MinTemp, RainToday@1 as RainToday], metrics=[output_rows_2=33, elapsed_compute_2=4.75µs, output_bytes_2=576.0 KB, output_batches_2=3, expr_0_ev
al_time_2=542ns, expr_1_eval_time_2=125ns]
  │     FilterExec: Rainfall@0 > 5, metrics=[output_rows_2=33, elapsed_compute_2=45.25µs, output_bytes_2=576.0 KB, output_batches_2=3, selectivity_2=14% (33/234)]
```
(If you are wondering why we only see metrics for task 2 in this
snippet, it's because this stage is the child of a union).

Display without metrics looks fine still
```
  │   ProjectionExec: expr=[Rainfall@0 as MinTemp, RainToday@1 as RainToday]
  │     FilterExec: Rainfall@0 > 5
```
## Summary
This change simplifies the codebase by removing ~600 lines of the custom
TTL Map logic in favor of a production-ready caching solution while
maintaining TTL behavior for task data entries. Moka caches also utilize
lock-free concurrent hash tables preventing deadlocks with certain
currency patterns.

**Key Changes**:
- Remove custom `TTLMap` implementation and associated manual garbage
collection logic
- Switch to moka async `Cache` with equivalent time-to-live semantics
- Update `Worker::task_data_entries` API from sync to async operations

## Benchmark Results
```
=== Comparing tpch_sf1 results from branch 'main' [prev] with 'moka-async-cache' [new] ===
os:        macos
arch:      aarch64
cpu cores: 16
threads:   2 -> 2
workers:   8 -> 8
==========================================================================================
tpch_sf1 q1: prev=  87 ms, new=  84 ms, diff=1.04 faster ✔
tpch_sf1 q2: prev=  63 ms, new=  61 ms, diff=1.03 faster ✔
tpch_sf1 q3: prev= 130 ms, new= 122 ms, diff=1.07 faster ✔
tpch_sf1 q4: prev=  35 ms, new=  35 ms, diff=1.00 slower ✖
tpch_sf1 q5: prev= 171 ms, new= 175 ms, diff=1.02 slower ✖
tpch_sf1 q6: prev=  27 ms, new=  26 ms, diff=1.04 faster ✔
tpch_sf1 q7: prev= 218 ms, new= 218 ms, diff=1.00 slower ✖
tpch_sf1 q8: prev= 156 ms, new= 151 ms, diff=1.03 faster ✔
tpch_sf1 q9: prev= 202 ms, new= 206 ms, diff=1.02 slower ✖
tpch_sf1 q10: prev= 114 ms, new= 115 ms, diff=1.01 slower ✖
tpch_sf1 q11: prev=  44 ms, new=  44 ms, diff=1.00 slower ✖
tpch_sf1 q12: prev=  47 ms, new=  46 ms, diff=1.02 faster ✔
tpch_sf1 q13: prev=  44 ms, new=  41 ms, diff=1.07 faster ✔
tpch_sf1 q14: prev=  34 ms, new=  35 ms, diff=1.03 slower ✖
tpch_sf1 q15: prev=  37 ms, new=  37 ms, diff=1.00 slower ✖
tpch_sf1 q16: prev=  31 ms, new=  31 ms, diff=1.00 slower ✖
tpch_sf1 q17: prev= 162 ms, new= 161 ms, diff=1.01 faster ✔
tpch_sf1 q18: prev= 120 ms, new= 121 ms, diff=1.01 slower ✖
tpch_sf1 q19: prev= 189 ms, new= 189 ms, diff=1.00 slower ✖
tpch_sf1 q20: prev=  70 ms, new=  64 ms, diff=1.09 faster ✔
tpch_sf1 q21: prev= 228 ms, new= 230 ms, diff=1.01 slower ✖
tpch_sf1 q22: prev=  27 ms, new=  28 ms, diff=1.04 slower ✖
```

```
=== Comparing tpcds_sf1 results from branch 'main' [prev] with 'moka-async-cache' [new] ===
os:        macos
arch:      aarch64
cpu cores: 16
threads:   2 -> 2
workers:   8 -> 8
===========================================================================================
tpcds_sf1 q1: Previously failed, and now also failed ❌
tpcds_sf1 q2: prev= 118 ms, new= 116 ms, diff=1.02 faster ✔
tpcds_sf1 q3: prev=  48 ms, new=  54 ms, diff=1.12 slower ✖
tpcds_sf1 q4: prev=1043 ms, new=1076 ms, diff=1.03 slower ✖
tpcds_sf1 q5: prev= 226 ms, new= 228 ms, diff=1.01 slower ✖
tpcds_sf1 q6: prev=3191 ms, new=3147 ms, diff=1.01 faster ✔
tpcds_sf1 q7: prev= 132 ms, new= 132 ms, diff=1.00 slower ✖
tpcds_sf1 q8: prev=  81 ms, new=  82 ms, diff=1.01 slower ✖
tpcds_sf1 q9: Previously failed, and now also failed ❌
tpcds_sf1 q10: prev= 100 ms, new= 101 ms, diff=1.01 slower ✖
tpcds_sf1 q11: prev= 651 ms, new= 652 ms, diff=1.00 slower ✖
tpcds_sf1 q12: prev=  30 ms, new=  30 ms, diff=1.00 slower ✖
tpcds_sf1 q13: prev= 188 ms, new= 187 ms, diff=1.01 faster ✔
tpcds_sf1 q14: prev= 697 ms, new= 694 ms, diff=1.00 faster ✔
tpcds_sf1 q15: prev=  76 ms, new=  75 ms, diff=1.01 faster ✔
tpcds_sf1 q16: prev=  46 ms, new=  42 ms, diff=1.10 faster ✔
tpcds_sf1 q17: Previously failed, and now also failed ❌
tpcds_sf1 q18: prev= 122 ms, new= 141 ms, diff=1.16 slower ✖
tpcds_sf1 q19: prev=  77 ms, new=  78 ms, diff=1.01 slower ✖
tpcds_sf1 q20: prev=  23 ms, new=  24 ms, diff=1.04 slower ✖
tpcds_sf1 q21: prev=  39 ms, new=  36 ms, diff=1.08 faster ✔
tpcds_sf1 q22: prev= 448 ms, new= 451 ms, diff=1.01 slower ✖
tpcds_sf1 q23: prev= 693 ms, new= 693 ms, diff=1.00 slower ✖
tpcds_sf1 q24: Previously failed, and now also failed ❌
tpcds_sf1 q25: prev= 112 ms, new= 116 ms, diff=1.04 slower ✖
tpcds_sf1 q26: prev=  92 ms, new=  82 ms, diff=1.12 faster ✔
tpcds_sf1 q27: Previously failed, and now also failed ❌
tpcds_sf1 q28: prev=  67 ms, new=  69 ms, diff=1.03 slower ✖
tpcds_sf1 q29: prev= 102 ms, new= 121 ms, diff=1.19 slower ✖
tpcds_sf1 q30: Previously failed, and now also failed ❌
tpcds_sf1 q31: prev= 243 ms, new= 227 ms, diff=1.07 faster ✔
tpcds_sf1 q32: prev=  31 ms, new=  31 ms, diff=1.00 slower ✖
tpcds_sf1 q33: prev=  76 ms, new=  78 ms, diff=1.03 slower ✖
tpcds_sf1 q34: prev=  85 ms, new=  86 ms, diff=1.01 slower ✖
tpcds_sf1 q35: prev=  85 ms, new=  86 ms, diff=1.01 slower ✖
tpcds_sf1 q36: Previously failed, and now also failed ❌
tpcds_sf1 q37: prev= 100 ms, new= 100 ms, diff=1.00 slower ✖
tpcds_sf1 q38: prev=  89 ms, new=  85 ms, diff=1.05 faster ✔
tpcds_sf1 q39: prev=  70 ms, new=  71 ms, diff=1.01 slower ✖
tpcds_sf1 q40: prev=  79 ms, new=  79 ms, diff=1.00 slower ✖
tpcds_sf1 q41: prev=  18 ms, new=  18 ms, diff=1.00 slower ✖
tpcds_sf1 q42: prev=  38 ms, new=  39 ms, diff=1.03 slower ✖
tpcds_sf1 q43: prev=  74 ms, new=  74 ms, diff=1.00 slower ✖
tpcds_sf1 q44: prev=  43 ms, new=  44 ms, diff=1.02 slower ✖
tpcds_sf1 q45: prev=  41 ms, new=  40 ms, diff=1.02 faster ✔
tpcds_sf1 q46: prev= 127 ms, new= 129 ms, diff=1.02 slower ✖
tpcds_sf1 q47: prev= 587 ms, new= 589 ms, diff=1.00 slower ✖
tpcds_sf1 q48: prev= 183 ms, new= 183 ms, diff=1.00 slower ✖
tpcds_sf1 q49: prev=  94 ms, new=  95 ms, diff=1.01 slower ✖
tpcds_sf1 q50: Previously failed, and now also failed ❌
tpcds_sf1 q51: prev= 164 ms, new= 175 ms, diff=1.07 slower ✖
tpcds_sf1 q52: prev=  45 ms, new=  49 ms, diff=1.09 slower ✖
tpcds_sf1 q53: prev=  69 ms, new=  69 ms, diff=1.00 slower ✖
tpcds_sf1 q54: Previously failed, and now also failed ❌
tpcds_sf1 q55: prev=  37 ms, new=  37 ms, diff=1.00 slower ✖
tpcds_sf1 q56: prev=  84 ms, new=  83 ms, diff=1.01 faster ✔
tpcds_sf1 q57: prev= 225 ms, new= 228 ms, diff=1.01 slower ✖
tpcds_sf1 q58: prev= 162 ms, new= 164 ms, diff=1.01 slower ✖
tpcds_sf1 q59: prev= 186 ms, new= 192 ms, diff=1.03 slower ✖
tpcds_sf1 q60: prev=  80 ms, new=  80 ms, diff=1.00 slower ✖
tpcds_sf1 q61: prev= 291 ms, new= 285 ms, diff=1.02 faster ✔
tpcds_sf1 q62: prev= 537 ms, new= 486 ms, diff=1.10 faster ✔
tpcds_sf1 q63: prev=  73 ms, new=  70 ms, diff=1.04 faster ✔
tpcds_sf1 q64: prev= 627 ms, new= 596 ms, diff=1.05 faster ✔
tpcds_sf1 q65: prev= 106 ms, new= 106 ms, diff=1.00 slower ✖
tpcds_sf1 q66: prev= 166 ms, new= 164 ms, diff=1.01 faster ✔
tpcds_sf1 q67: prev= 636 ms, new= 676 ms, diff=1.06 slower ✖
tpcds_sf1 q68: prev= 126 ms, new= 127 ms, diff=1.01 slower ✖
tpcds_sf1 q69: prev=  94 ms, new=  94 ms, diff=1.00 slower ✖
tpcds_sf1 q70: Previously failed, and now also failed ❌
tpcds_sf1 q71: prev=  65 ms, new=  64 ms, diff=1.02 faster ✔
tpcds_sf1 q73: prev=  89 ms, new=  89 ms, diff=1.00 slower ✖
tpcds_sf1 q74: prev= 401 ms, new= 393 ms, diff=1.02 faster ✔
tpcds_sf1 q75: prev= 297 ms, new= 263 ms, diff=1.13 faster ✔
tpcds_sf1 q76: prev=  62 ms, new=  69 ms, diff=1.11 slower ✖
tpcds_sf1 q77: prev= 136 ms, new= 135 ms, diff=1.01 faster ✔
tpcds_sf1 q78: prev= 295 ms, new= 307 ms, diff=1.04 slower ✖
tpcds_sf1 q79: prev= 114 ms, new= 112 ms, diff=1.02 faster ✔
tpcds_sf1 q80: prev= 246 ms, new= 243 ms, diff=1.01 faster ✔
tpcds_sf1 q81: prev=  38 ms, new=  36 ms, diff=1.06 faster ✔
tpcds_sf1 q82: prev=  85 ms, new=  86 ms, diff=1.01 slower ✖
tpcds_sf1 q83: prev=  42 ms, new=  42 ms, diff=1.00 slower ✖
tpcds_sf1 q84: prev=  28 ms, new=  27 ms, diff=1.04 faster ✔
tpcds_sf1 q85: prev= 110 ms, new= 108 ms, diff=1.02 faster ✔
tpcds_sf1 q86: prev=  24 ms, new=  25 ms, diff=1.04 slower ✖
tpcds_sf1 q87: prev=  90 ms, new=  87 ms, diff=1.03 faster ✔
tpcds_sf1 q88: prev= 338 ms, new= 341 ms, diff=1.01 slower ✖
tpcds_sf1 q89: prev=  81 ms, new=  81 ms, diff=1.00 slower ✖
tpcds_sf1 q90: prev=  30 ms, new=  30 ms, diff=1.00 slower ✖
tpcds_sf1 q91: prev=  34 ms, new=  33 ms, diff=1.03 faster ✔
tpcds_sf1 q92: prev=  37 ms, new=  37 ms, diff=1.00 slower ✖
tpcds_sf1 q93: prev=  90 ms, new=  91 ms, diff=1.01 slower ✖
tpcds_sf1 q94: prev=  42 ms, new=  39 ms, diff=1.08 faster ✔
tpcds_sf1 q95: prev= 335 ms, new= 334 ms, diff=1.00 faster ✔
tpcds_sf1 q96: prev=  47 ms, new=  47 ms, diff=1.00 slower ✖
tpcds_sf1 q97: prev=  77 ms, new=  77 ms, diff=1.00 slower ✖
tpcds_sf1 q98: prev=  70 ms, new=  71 ms, diff=1.01 slower ✖
tpcds_sf1 q99: prev=5103 ms, new=6484 ms, diff=1.27 slower ❌
```
I think I messed something up in
datafusion-contrib#306
and I see some leaf stages are missing their metrics.

This PR solves it and adds a new test that fails on main but succeeds in
this PR.
…ion-contrib#323)

## Summary

This PR allows `NetworkCoalesceExec` to output multiple consumer tasks.
This helps build the foundation for partial aggregation trees: with
multi-consumer coalesce, we can introduce intermediate “gather” stages
where each task aggregates a subset of upstream tasks, then a final
stage merges/aggregates again.

## What changed
- **`NetworkCoalesceExec` now supports `task_count > 1`**
  - Each consumer task reads a **contiguous group** of upstream tasks.
- Output partitioning is sized using the **maximum group size**, and
consumer tasks with smaller groups return **empty streams** for the
“extra” partitions.
  
## Note
This PR was larger in scope, adding support for Network Coalesce Trees.
Coalesce trees are not that useful on their own, and were removed from
this PR. Future work to add coalesce trees in tandem with partial
aggregation roll-up. Please see the discussion below for more details.

---

## Previous Summary with Network Coalesce Tree (Outdated)

This PR adds support for **hierarchical coalescing across the network**
by allowing `NetworkCoalesceExec` to **output to multiple consumer
tasks** and adjusting the distributed planner to optionally insert a
**coalesce tree** (multiple `NetworkCoalesceExec` layers) when a plan
must ultimately coalesce to a single task.

The key semantic is preserved: the **final** coalesce boundary still
ends in **one task** (matching the intent of `CoalescePartitionsExec` /
`SortPreservingMergeExec`), but fan-in can be reduced via intermediate
coalesce stages.

## What changed

- **`NetworkCoalesceExec` now supports `task_count > 1`**
  - Each consumer task reads a **contiguous group** of upstream tasks.
- Output partitioning is sized using the **maximum group size**, and
consumer tasks with smaller groups return **empty streams** for the
“extra” partitions.
- **Distributed planner can insert a “coalesce tree”**
- When enabled, the planner may insert **intermediate
`NetworkCoalesceExec` stages** to reduce fan-in while still converging
to a single final task.
- The intermediate task count is computed using a **log-based fan-in
heuristic** derived from the worker count (clamped to a maximum of 16).
- **Plan/stage visualization improvements**
- Stage discovery and edge rendering account for coalesce-tree behavior,
so `display_plan_ascii` (and related stage displays) show the expected
stage structure.

## Why this is useful
- **Lower fan-in cost vs a single wide coalesce**
- A coalesce tree can reduce the bottleneck of having one task directly
pull from \(N\) upstream tasks. In practice this reduces connection
pressure and can improve throughput/tail latency by spreading the work
across \(O(\log N)\) coalesce levels instead of one wide fan-in point
(workload/network dependent).
- **Hierarchical distributed aggregation**
- Avoids a single “hot” reducer when many upstream tasks must converge
to one final aggregation stage.
- Enables partial merging closer to producers, reducing bytes
transferred and improving tail latency.
- **Global sorts / `ORDER BY` pipelines**
- Sort-preserving merge patterns often require a final single task;
intermediate coalescing reduces fan-in and connection pressure before
the final merge.
- **Top-K / limit with ordering**
- Intermediate coalesce stages can reduce the number of upstream streams
feeding the final merge/limit stage, improving responsiveness and memory
behavior.
- **Skewed workloads**
- When one stage produces many tasks (or uneven partitions),
hierarchical coalescing reduces single-node bottlenecks and evens out
execution.

## Configuration

This PR adds two distributed planner configuration options:

- **`distributed.coalesce_tree_enabled`** (bool, default: `false`)
- Enables insertion of intermediate `NetworkCoalesceExec` stages (a
coalesce tree) for boundaries that must ultimately coalesce to a single
task.
- **`distributed.coalesce_tree_min_input_tasks`** (usize, default: `8`)
- Minimum number of upstream tasks required before the planner will
insert a coalesce tree.
- If the upstream task count is below this threshold, the planner
inserts only the final single-task `NetworkCoalesceExec`.

---------

Co-authored-by: gkerr_adobe <gkerr@adobe.com>
…b#322)

This change adds `"network_latency_first", "network_latency_min",
"network_latency_max", "network_latency_avg", "network_latency_p99"`
metrics that measures the time between creating a message on
a worker to the client recieving it. This is measured for all messages
received by a task (in all partitions).

I considered tracking it for all partitions (using labels for the
partition id), however, `Time` metrics are aggregated
by summation, not average, so this would not work. We would need a
custom metric for an avg time aggregator and these
are unsupported.

Example plan:
```
───── DistributedExec ── Tasks: t0:[p0]
│ CoalescePartitionsExec, metrics=[output_rows_0=2, elapsed_compute_0=11.08µs, output_bytes_0=256.0 KB, output_batches_0=2]
│   HashJoinExec: mode=CollectLeft, join_type=Left, on=[(RainTomorrow@1, RainTomorrow@1)], projection=[MinTemp@0, MaxTemp@2], metrics=[output_rows_0=2, elapsed_compute_0=139.00µs, output_
bytes_0=256.0 KB, output_batches_0=2, build_input_batches_0=2, build_input_rows_0=2, input_batches_0=3, input_rows_0=2, build_mem_used_0=731, build_time_0=73.75µs, join_time_0=58.71µs, av
g_fanout_0=100% (2/2), probe_hit_rate_0=100% (2/2)]
│     CoalescePartitionsExec, metrics=[output_rows_0=2, elapsed_compute_0=6.88µs, output_bytes_0=640.0 B, output_batches_0=2]
│       [Stage 2] => NetworkCoalesceExec: output_partitions=6, input_tasks=2, metrics=[elapsed_compute_0=476.88µs, bytes_transferred_0=23.73 K, max_mem_used_0=1.54 K, network_latency_avg_
0=1.66ms, network_latency_first_0=1.81ms, network_latency_max_0=1.81ms, network_latency_min_0=1.50ms, network_latency_p99_0=1.70ms]
│     ProjectionExec: expr=[avg(weather.MaxTemp)@1 as MaxTemp, RainTomorrow@0 as RainTomorrow], metrics=[output_rows_0=2, elapsed_compute_0=2.42µs, output_bytes_0=32.1 KB, output_batches_
0=2, expr_0_eval_time_0=293ns, expr_1_eval_time_0=169ns]
│       AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[avg(weather.MaxTemp)], metrics=[output_rows_0=2, elapsed_compute_0=103.08µs, output_bytes_0=32.1
KB, output_batches_0=2, spill_count_0=0, spilled_bytes_0=0.0 B, spilled_rows_0=0, peak_mem_used_0=50.54 K, aggregate_arguments_time_0=6.29µs, aggregation_time_0=7.75µs, emitting_time_0=6.
63µs, time_calculating_group_ids_0=11.08µs]
│         [Stage 3] => NetworkShuffleExec: output_partitions=3, input_tasks=3, metrics=[elapsed_compute_0=1.01ms, bytes_transferred_0=23.26 K, max_mem_used_0=4.74 K, network_latency_avg_0
=14.65ms, network_latency_first_0=14.61ms, network_latency_max_0=16.69ms, network_latency_min_0=13.04ms, network_latency_p99_0=14.59ms]
└──────────────────────────────────────────────────
  ┌───── Stage 2 ── Tasks: t0:[p0..p2] t1:[p0..p2]
  │ ProjectionExec: expr=[avg(weather.MinTemp)@1 as MinTemp, RainTomorrow@0 as RainTomorrow], metrics=[output_rows_0=1, output_rows_1=1, elapsed_compute_0=1.33µs, elapsed_compute_1=1.38µs
, output_bytes_0=16.0 KB, output_bytes_1=16.0 KB, output_batches_0=1, output_batches_1=1, expr_0_eval_time_0=85ns, expr_0_eval_time_1=169ns, expr_1_eval_time_0=44ns, expr_1_eval_time_1=85
ns]
  │   AggregateExec: mode=FinalPartitioned, gby=[RainTomorrow@0 as RainTomorrow], aggr=[avg(weather.MinTemp)], metrics=[output_rows_0=1, output_rows_1=1, elapsed_compute_0=78.09µs, elapse
d_compute_1=78.17µs, output_bytes_0=16.0 KB, output_bytes_1=16.0 KB, output_batches_0=1, output_batches_1=1, spill_count_0=0, spill_count_1=0, spilled_bytes_0=0.0 B, spilled_bytes_1=0.0 B
, spilled_rows_0=0, spilled_rows_1=0, peak_mem_used_0=50.45 K, peak_mem_used_1=50.45 K, aggregate_arguments_time_0=4.38µs, aggregate_arguments_time_1=3.88µs, aggregation_time_0=7.29µs, ag
gregation_time_1=4.63µs, emitting_time_0=3.71µs, emitting_time_1=9.34µs, time_calculating_group_ids_0=7.50µs, time_calculating_group_ids_1=5.71µs]
  │     [Stage 1] => NetworkShuffleExec: output_partitions=3, input_tasks=3, metrics=[elapsed_compute_0=755.83µs, elapsed_compute_1=735.25µs, bytes_transferred_0=16.51 K, bytes_transferre
d_1=10.64 K, max_mem_used_0=3.10 K, max_mem_used_1=3.47 K, network_latency_avg_0=3.05ms, network_latency_avg_1=3.21ms, network_latency_first_0=3.41ms, network_latency_first_1=3.44ms, netw
ork_latency_max_0=3.41ms, network_latency_max_1=3.44ms, network_latency_min_0=2.67ms, network_latency_min_1=2.75ms, network_latency_p99_0=3.30ms, network_latency_p99_1=3.36ms]
  └──────────────────────────────────────────────────
    ┌───── Stage 1 ── Tasks: t0:[p0..p5] t1:[p0..p5] t2:[p0..p5]
    │ RepartitionExec: partitioning=Hash([RainTomorrow@0], 6), input_partitions=1, metrics=[]
    │   AggregateExec: mode=Partial, gby=[RainTomorrow@1 as RainTomorrow], aggr=[avg(weather.MinTemp)], metrics=[]
    │     FilterExec: RainToday@1 = Yes, projection=[MinTemp@0, RainTomorrow@2], metrics=[]
    │       PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0], metrics=[]
    │         DataSourceExec: file_groups={3 groups: [[Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000000.parquet], [Users/jayant.shrivastava/code/datafus
ion-distributed/testdata/weather/result-000001.parquet], [Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000002.parquet]]}, projection=[MinTemp, RainToday, R
ainTomorrow], file_type=parquet, predicate=RainToday@19 = Yes, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= Yes AND Yes <= RainToday_max@1, required_guar
antees=[RainToday in (Yes)], metrics=[]
    └──────────────────────────────────────────────────
 ┌───── Stage 3 ── Tasks: t0:[p0..p2] t1:[p0..p2] t2:[p0..p2]
  │ RepartitionExec: partitioning=Hash([RainTomorrow@0], 3), input_partitions=1, metrics=[output_rows_0=2, output_rows_1=2, output_rows_2=2, elapsed_compute_0=18.50µs, elapsed_compute_1=1
7.50µs, elapsed_compute_2=21.54µs, output_bytes_0=512.0 KB, output_bytes_1=512.0 KB, output_bytes_2=512.0 KB, output_batches_0=2, output_batches_1=2, output_batches_2=2, spill_count_0=0,
spill_count_1=0, spill_count_2=0, spilled_bytes_0=0.0 B, spilled_bytes_1=0.0 B, spilled_bytes_2=0.0 B, spilled_rows_0=0, spilled_rows_1=0, spilled_rows_2=0, fetch_time_0=4.98ms, fetch_tim
e_1=7.02ms, fetch_time_2=5.60ms, repartition_time_0=12.17µs, repartition_time_1=19.62µs, repartition_time_2=19.79µs, send_time_0=4.04µs, send_time_1=5.46µs, send_time_2=6.17µs]
  │   AggregateExec: mode=Partial, gby=[RainTomorrow@1 as RainTomorrow], aggr=[avg(weather.MaxTemp)], metrics=[output_rows_0=2, output_rows_1=2, output_rows_2=2, elapsed_compute_0=57.46µs
, elapsed_compute_1=82.13µs, elapsed_compute_2=75.79µs, output_bytes_0=16.1 KB, output_bytes_1=16.1 KB, output_bytes_2=16.1 KB, output_batches_0=1, output_batches_1=1, output_batches_2=1,
 spill_count_0=0, spill_count_1=0, spill_count_2=0, spilled_bytes_0=0.0 B, spilled_bytes_1=0.0 B, spilled_bytes_2=0.0 B, spilled_rows_0=0, spilled_rows_1=0, spilled_rows_2=0, skipped_aggr
egation_rows_0=0, skipped_aggregation_rows_1=0, skipped_aggregation_rows_2=0, peak_mem_used_0=18.55 K, peak_mem_used_1=18.70 K, peak_mem_used_2=18.67 K, aggregate_arguments_time_0=1.38µs,
 aggregate_arguments_time_1=4.54µs, aggregate_arguments_time_2=3.50µs, aggregation_time_0=6.67µs, aggregation_time_1=9.00µs, aggregation_time_2=8.71µs, emitting_time_0=4.42µs, emitting_ti
me_1=5.79µs, emitting_time_2=5.88µs, time_calculating_group_ids_0=20.38µs, time_calculating_group_ids_1=28.04µs, time_calculating_group_ids_2=27.71µs, reduction_factor_0=2.2% (2/89), redu
ction_factor_1=1.9% (2/107), reduction_factor_2=1.9% (2/104)]
  │     FilterExec: RainToday@1 = No, projection=[MaxTemp@0, RainTomorrow@2], metrics=[output_rows_0=89, output_rows_1=107, output_rows_2=104, elapsed_compute_0=52.08µs, elapsed_compute_1
=44.50µs, elapsed_compute_2=44.67µs, output_bytes_0=192.0 KB, output_bytes_1=192.0 KB, output_bytes_2=192.0 KB, output_batches_0=1, output_batches_1=1, output_batches_2=1, selectivity_0=7
3% (89/122), selectivity_1=88% (107/122), selectivity_2=85% (104/122)]
  │       PartitionIsolatorExec: t0:[p0,__,__] t1:[__,p0,__] t2:[__,__,p0], metrics=[]
  │         DataSourceExec: file_groups={3 groups: [[Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000000.parquet], [Users/jayant.shrivastava/code/datafusio
n-distributed/testdata/weather/result-000001.parquet], [Users/jayant.shrivastava/code/datafusion-distributed/testdata/weather/result-000002.parquet]]}, projection=[MaxTemp, RainToday, Rai
nTomorrow], file_type=parquet, predicate=RainToday@19 = No, pruning_predicate=RainToday_null_count@2 != row_count@3 AND RainToday_min@0 <= No AND No <= RainToday_max@1, required_guarantee
s=[RainToday in (No)], metrics=[output_rows_0=122, output_rows_1=122, output_rows_2=122, elapsed_compute_0=1ns, elapsed_compute_1=1ns, elapsed_compute_2=1ns, output_bytes_0=6.5 KB, output
_bytes_1=6.5 KB, output_bytes_2=6.5 KB, output_batches_0=1, output_batches_1=1, output_batches_2=1, files_ranges_pruned_statistics_0=1 total → 1 matched, files_ranges_pruned_statistics_1=
1 total → 1 matched, files_ranges_pruned_statistics_2=1 total → 1 matched, row_groups_pruned_statistics_0=1 total → 1 matched, row_groups_pruned_statistics_1=1 total → 1 matched, row_grou
ps_pruned_statistics_2=1 total → 1 matched, row_groups_pruned_bloom_filter_0=1 total → 1 matched, row_groups_pruned_bloom_filter_1=1 total → 1 matched, row_groups_pruned_bloom_filter_2=1
total → 1 matched, page_index_rows_pruned_0=122 total → 122 matched, page_index_rows_pruned_1=122 total → 122 matched, page_index_rows_pruned_2=122 total → 122 matched, batches_split_0=0,
 batches_split_1=0, batches_split_2=0, bytes_scanned_0=13.53 K, bytes_scanned_1=13.40 K, bytes_scanned_2=13.44 K, file_open_errors_0=0, file_open_errors_1=0, file_open_errors_2=0, file_sc
an_errors_0=0, file_scan_errors_1=0, file_scan_errors_2=0, num_predicate_creation_errors_0=0, num_predicate_creation_errors_1=0, num_predicate_creation_errors_2=0, predicate_cache_inner_r
ecords_0=0, predicate_cache_inner_records_1=0, predicate_cache_inner_records_2=0, predicate_cache_records_0=0, predicate_cache_records_1=0, predicate_cache_records_2=0, predicate_evaluati
on_errors_0=0, predicate_evaluation_errors_1=0, predicate_evaluation_errors_2=0, pushdown_rows_matched_0=0, pushdown_rows_matched_1=0, pushdown_rows_matched_2=0, pushdown_rows_pruned_0=0,
 pushdown_rows_pruned_1=0, pushdown_rows_pruned_2=0, bloom_filter_eval_time_0=31.00µs, bloom_filter_eval_time_1=29.88µs, bloom_filter_eval_time_2=30.75µs, metadata_load_time_0=3.82ms, met
adata_load_time_1=4.71ms, metadata_load_time_2=3.88ms, page_index_eval_time_0=54.00µs, page_index_eval_time_1=45.21µs, page_index_eval_time_2=46.71µs, row_pushdown_eval_time_0=2ns, row_pu
shdown_eval_time_1=2ns, row_pushdown_eval_time_2=2ns, statistics_eval_time_0=41.00µs, statistics_eval_time_1=33.71µs, statistics_eval_time_2=34.54µs, time_elapsed_opening_0=4.22ms, time_e
lapsed_opening_1=5.13ms, time_elapsed_opening_2=4.19ms, time_elapsed_processing_0=1.57ms, time_elapsed_processing_1=1.53ms, time_elapsed_processing_2=1.44ms, time_elapsed_scanning_total_0
=636.83µs, time_elapsed_scanning_total_1=1.77ms, time_elapsed_scanning_total_2=1.29ms, time_elapsed_scanning_until_data_0=533.79µs, time_elapsed_scanning_until_data_1=1.68ms, time_elapsed
_scanning_until_data_2=1.19ms, scan_efficiency_ratio_0=29% (13.53 K/46.08 K), scan_efficiency_ratio_1=29% (13.40 K/45.46 K), scan_efficiency_ratio_2=29% (13.44 K/45.82 K)]
  └──────────────────────────────────────────────────
```
The `shuffle_batch_size` config parameter was missing some wireup with
the `DistributedExt` trait. This PR adds it.
…-contrib#340)

## Summary

This PR adds support for forwarding arbitrary HTTP headers unchanged to
worker nodes via Arrow Flight requests. This enables propagating context
(such as tracing headers, priority levels, or authentication tokens)
across the distributed cluster.

## Motivation

When running distributed queries, coordinators often receive HTTP
requests containing headers that need to be forwarded to worker nodes.
Examples include:
- Tracing/observability headers (e.g., `x-request-id`, `traceparent`)
- Priority or QoS headers
- Custom application-specific context

Previously, there was no mechanism to pass arbitrary headers through to
workers. The existing `ConfigExtension` system only supports structured
configuration that implements the `ConfigExtension` trait.

## Implementation

### New Module: `passthrough_headers.rs`

A lightweight module that stores headers as a `SessionConfig` extension:

```rust
// Store headers
set_passthrough_headers(&mut config, headers)?;

// Retrieve headers
let headers = get_passthrough_headers(&config);
```

### New Trait Methods on `DistributedExt`

Two new methods following the existing patterns:

```rust
// Builder pattern (returns Result for validation)
fn with_distributed_passthrough_headers(self, headers: HeaderMap) -> Result<Self, DataFusionError>;

// In-place mutation
fn set_distributed_passthrough_headers(&mut self, headers: HeaderMap) -> Result<(), DataFusionError>;
```

Implemented for all four types: `SessionConfig`, `SessionStateBuilder`,
`SessionState`, and `SessionContext`.

### Integration with Header Propagation

Passthrough headers are merged into outgoing worker requests in
`get_config_extension_propagation_headers()`, alongside the existing
config extension headers.

### Safety Validation

To prevent accidental override of internal headers, any header starting
with the reserved prefix `x-datafusion-distributed-config-` is rejected
with a configuration error:

```
Configuration error: Passthrough header 'x-datafusion-distributed-config-foo'
uses reserved prefix 'x-datafusion-distributed-config-'.
This prefix is reserved for internal config extension propagation.
```

## Usage Example

```rust
use datafusion::execution::SessionStateBuilder;
use datafusion_distributed::DistributedExt;
use http::HeaderMap;

// Coordinator captures headers from incoming request
let mut passthrough = HeaderMap::new();
if let Some(v) = incoming_request.headers().get("x-request-id") {
    passthrough.insert("x-request-id", v.clone());
}
if let Some(v) = incoming_request.headers().get("x-priority") {
    passthrough.insert("x-priority", v.clone());
}

// Headers will be forwarded to all worker nodes
let state = SessionStateBuilder::new()
    .with_distributed_passthrough_headers(passthrough)?
    .build();
```

## Testing

Added 7 new tests:

| Test | Description |
|------|-------------|
| `test_set_and_get_passthrough_headers` | Basic round-trip
functionality |
| `test_get_passthrough_headers_empty` | Empty/unset case returns empty
map |
| `test_overwrite_passthrough_headers` | Setting headers replaces
previous |
| `test_rejects_reserved_prefix` | Validates reserved prefix rejection |
| `test_accepts_similar_but_different_prefix` | Similar prefixes that
don't match are allowed |
| `test_passthrough_headers_included_in_propagation` | Integration:
headers appear in propagation |
| `test_passthrough_headers_without_config_extensions` | Works
standalone without config extensions |

## Checklist

- [x] Code compiles without warnings
- [x] All existing tests pass (159 tests)
- [x] New tests added and passing (7 tests)
- [x] Clippy passes
- [x] Documentation added to trait methods
- [x] Follows existing code patterns and conventions

---------

Co-authored-by: Gabriel Musat Mestre <gabriel.musatmestre@datadoghq.com>
## Context and Problem

closes datafusion-contrib#319 

### 1. Panics in tpcds
There have been panics after that look like:
```text
thread 'tokio-runtime-worker' panicked at /home/runner/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/datafusion-physical-plan-52.0.0/src/repartition/mod.rs:972:22: partition not used yet
```
This was happening because some operators (like TopK) cancel once
satisfied. When used with `BroadcastExec` this would cause issues with
the caching mechanism: since the consumer task was responsible for
populating the cache, if it was cancelled, the cache would not be
populated although it should be. Thus, the first consumer task would
execute partition, begin populating cache, abort it when cancelled,
second task would execute the same partition because cache not
populated.

### 2. Streaming Cache 
There was follow-up work noted to make the cache in broadcast
"streaming".
Prior to this, the cache was simple: the first task would collect all
batches from the input and populate the cache. This is not optimal as
there can be multiple subscribers that want the same partition data
concurrently. Requiring all other consumers to wait until all batches
have been fetched obviously causes slow-down as consumers can receive
partial data as its received.

## Solution
This PR solves these two problems by:
1. Cache population in separate task
- put the cache population in its own task and producer push batches
into the entry as it receives them.

2. Introduce streaming cache via `CacheEntry` and `EntryState`
- Struct that contains state and buffered data for new subscribers to
replay data that has been received and wait for new data to be pulled.

## Results / Benches
Some queries that are using broadcast see speedups 😄 

<img width="325" height="950" alt="Screenshot 2026-01-27 at 6 28 13 PM"
src="https://github.com/user-attachments/assets/f6754664-fab8-41e0-9a57-c61636901f75"
/>

and we definitively fixed the panic error on re-executing the same
partition. here is a before and after on a query that had this panic:

**Before:** we see the panic about partitions being executed multiple
times
<img width="1428" height="144" alt="Screenshot 2026-01-27 at 6 28 56 PM"
src="https://github.com/user-attachments/assets/3d5a8efd-2b0c-485e-b0d8-2253feccbad2"
/>

**After:** we no longer see it but see the arrow error as explained
below
<img width="1430" height="107" alt="Screenshot 2026-01-27 at 6 29 18 PM"
src="https://github.com/user-attachments/assets/1e5af5f9-13b1-4b9b-817b-667b3dafe78d"
/>

## IMPORTANT NOTE - DataFusion Regression
Some tpcds queries panic with:
```text
thread 'tokio-runtime-worker' panicked at .../h2-0.4.13/...:
state=Closed(ScheduledLibraryReset(CANCEL))
```
due to upstream error:
```text
Arrow error: Invalid argument error: column types must match schema types, 
expected Dictionary(UInt16, Utf8) but found Utf8 at column index 1
```
I have confirmed that this is an upstream DataFusion issue (which I may
or may not have time to do tonight 😅) and not related to broadcast, but
it unveiled it. I proved this by reproducing this error on single-node
DataFusion on the same data:

Prerequisite: have the tpcds data downloaded at some arbitrary location
1. Go to datafusion repo
2. run: `datafusion-cli`
3. use sql:
```sql
CREATE EXTERNAL TABLE store STORED AS PARQUET LOCATION 'path_to_tpcds_data/store';
SELECT s_county, s_state FROM store LIMIT 1;
```
4. You will see:
```text
Error: Arrow error: Invalid argument error: column types must match schema types, expected Dictionary(UInt16, Utf8) but found Utf8 at column index 1
```

When initially debugging this, I didn't see it in my single-node (I was
using 50.2.0). Then I realized we upgraded in distributed thus I
upgraded to (52.1.0) and the error occurs.

---------

Co-authored-by: Gabriel Musat Mestre <gabriel.musatmestre@datadoghq.com>
This PR adds several improvements to the benchmarks:
- Claude Skills: adds some Claude Code skills for dealing with remote
benchmarks. Typically these benchmarks take a while to execute, and
analyzing the results can be a bit time consuming. This PR allows Claude
to autonomously do all the tedious tasks of running benchmarks, waiting
for the result, analyzing them, performing small changes, and redeploy
for benchmark again.
- Adds ballista back: adds ballista benchmarks back so that we have
another engine to compare to. This time, ballista is added in a cargo
crate detached from the workspace, so users do not need to care about
importing it. All the ballista code is pretty much what was removed from
datafusion-contrib#292
- Stores the execution plans with metrics attached to the benchmarks
run. This way, Claude can run benchmarks and analyze the execution plans
along with the results in order to identify bottlenecks and improvement
points.
…contrib#341)

Adds a test reproducing the issue described in
[datafusion-contrib#187](datafusion-contrib#187).

Since the test fails by design, it is marked as `#[ignore]`.

The underlying issue is that metrics are attached to the last
`FlightData` message of the last partition stream on the worker side (in
`do_get.rs`). When a `LIMIT` is present, the client-side network
boundary node (e.g. `NetworkShuffleExec`) drops the stream once it has
enough rows, before the worker finishes all partitions. The final
message carrying the metrics is never sent or never consumed, so metrics
for that stage are lost.

The test runs a `GROUP BY ... LIMIT 1` query against the `flights_1m`
dataset (1M rows). The large dataset ensures the worker is still
producing data when the LIMIT causes the client to drop the stream.

Note: This don't actually solve the issue, just showcases that it is
really an issue!
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.