Distributed aggregate queries are round-trip bound: push partial aggregates to the remote scanner

[tillrohrmann]

Problem

An aggregate query over sys_vqueue_meta on a customer cluster took 24s. It scanned 6.4M rows, with the heaviest single partition stream carrying 5.5M of them.

Today the query engine only pushes filters to remote nodes. Everything else — projection beyond the pushed schema, aggregation, sorting — happens on the single coordinator node, so every matching row is shipped across the network as Arrow IPC and re-decoded centrally. For SELECT scope, COUNT(*) ... GROUP BY scope that is strictly wasted work: the count could be computed by the partition owner and only one row per group returned.

There are two independent causes:

Cause 1 — the transport is a serial, 128-rows-per-round-trip pull loop

This dominates the 24s.

• The session batch size is 128 rows, set for the whole SessionContext
• It reaches the remote scanner unchanged: PartitionedExecutionPlan::execute reads context.session_config().batch_size() and passes it through ScanPartition::scan_partition into RemoteQueryScannerOpen.batch_size.
• The client loop is strictly serial: it awaits next_batch(), forwards the batch downstream, then loops, never more than one request in flight. The receiver channel depth is 1 (remote_query_scanner_client.rs:163), so there is no read-ahead either.
• The server yields exactly one RecordBatch per Next request, and each batch is its own Arrow IPC stream.

So one partition stream is capped at batch_size / RTT regardless of how fast storage is:

128 rows / 540µs  ≈  237k rows/s
5.5M rows / 128 × 540µs  ≈  23s        # ≈ the observed 24s

A second, independent serialization axis: physical partitions inside one logical partition are scanned sequentially via stream::iter(..).try_flatten() (table_providers.rs:354-376), so whenever #partitions > target_partitions they queue behind each other rather than overlapping.

Cause 2 — we ship rows where aggregates would do

Even with a perfect transport, network and coordinator cost stay O(rows). Aggregate pushdown makes it O(groups) and moves the hash aggregate CPU to the nodes that own the data.

Why DataFusion offers no ready-made hook

DataFusion has no aggregate-pushdown extension point. TableProvider exposes only scan, scan_with_args, supports_filters_pushdown and statistics, and ScanArgs carries projection, filters, limit and preferred ordering — nothing aggregate-shaped. There is no supports_aggregate_pushdown analogue.

The sanctioned route is a custom PhysicalOptimizerRule that rewrites AggregateExec{mode: Partial} sitting above our scan. This is exactly what Ballista, datafusion-distributed and datafusion-federation all do.

The good news is that the plan shape we need to match already exists. CombinePartialFinalAggregate (datafusion-physical-optimizer-54.0.0/src/optimizer.rs:179) only fuses Partial+Final into Single when they are adjacent, and our scan reports Partitioning::UnknownPartitioning(n) (table_providers.rs:250). With n > 1, EnforceDistribution inserts a repartition/coalesce between the two stages, so for any multi-partition scan DataFusion already produces:

AggregateExec(Final / FinalPartitioned)
  CoalescePartitionsExec | RepartitionExec(Hash)
    AggregateExec(Partial)              <-- this is what we want to run remotely
      FilterExec                        <-- always present, because pushdown is Inexact
        PartitionedExecutionPlan        <-- n logical partitions

We also already ship serialized DataFusion physical expressions over the wire for filter pushdown, and datafusion-proto can serialize AggregateExec / AggregateMode / PhysicalGroupBy out of the box, with PhysicalExtensionCodec for our custom leaf.

Make the transport stop being the bottleneck

1. Decouple the remote-scan batch size from the session batch size. The 128 exists to bound memory for interactive SELECT * queries, not for bulk scans; the remote scanner wants something 1–2 orders of magnitude larger. RemoteQueryScannerOpen.batch_size is already a wire field, so this is a plan-time decision.
2. Pipeline the Next RPCs. Keep k requests in flight and raise the receiver depth above 1. The server side is already a queue, drained in order — so this is client-only and needs no wire change.

Ideally we can make the batch sizes bound by memory and not by number of rows as this makes the memory footprint more controllable.

Expected effect: large constant-factor speedup, but network and coordinator cost stay O(rows).

Aggregate pushdown

The core of both variants is the same: a Restate PhysicalOptimizerRule that matches AggregateExec{Partial} → [FilterExec →] PartitionedExecutionPlan and folds the aggregate (and the filter) into the leaf, so the leaf's output becomes pre-aggregated rows and the coordinator keeps only the final merge. The variants differ in what crosses the wire.

Ship a restricted, Restate-owned aggregate spec

Push only aggregates that are merge-stable in their own domain, i.e. where the partial result is an ordinary row of the same type: COUNT → merged with SUM, SUM, MIN, MAX, with GROUP BY restricted to a whitelist of column types. AVG is either derived on the coordinator from pushed SUM/COUNT or simply not pushed.

• Pro: the wire format is ours and stable — no DataFusion internals on the wire, no version coupling, no
  capability negotiation beyond "does the peer know tag 9". Small correctness surface, and it covers the
  motivating query and the COUNT(*) GROUP BY status example in docs/dev/datafusion.md.
• Con: manual mapping from DataFusion's aggregate UDFs to our spec; no coverage for percentiles,
  approx-distinct, etc. Those fall back to row streaming.

Alternative option: Ship a serialized AggregateExec(Partial) sub-plan

Add an Option<...> field to RemoteQueryScannerOpen (next free bilrost tag is 9, crates/types/src/net/remote_query_scanner.rs:37-65) carrying the aggregate encoded with datafusion-proto. The server rebuilds the Partial aggregate over its local scan stream and returns partial-state batches; the coordinator's Final/FinalPartitioned stage merges them.

• Pro: maximum coverage (anything DataFusion can aggregate) and maximum reuse — no accumulator logic of our own.
• Con — the sharp edge: a Partial aggregate's output schema is expr.state_fields(), i.e. per-UDF accumulator state, not values. That makes a DataFusion-internal, version-dependent structure part of our wire contract. A DataFusion upgrade becomes a protocol change, and mixed-version clusters during a rolling upgrade must not attempt it.
• Con — no free fallback. Filter pushdown is safe against old peers precisely because we advertise it as Inexact / PushedDown::No: a peer that ignores the predicate still returns correct results, just slower, and the retained FilterExec catches everything. Aggregate pushdown has no such property — a peer that ignores the new field returns rows where the coordinator expects partial state. There is no scanner-level protocol version negotiation today and bilrost silently drops unknown tags, so this needs an explicit capability check with a row-streaming fallback.
• Correctness traps to enumerate: DISTINCT aggregates, ordered aggregates (ORDER BY inside the aggregate), per-aggregate FILTER clauses, AggregateMode::Single (must be split back, or the rule must run before CombinePartialFinalAggregate), and absorbing the FilterExec that Inexact pushdown always leaves behind — which is safe because the server applies the predicate exactly but must be made explicit.

Prior art and references

• docs/dev/datafusion.md — "Aggregation Optimization" (L283) already sketches the plan-fragment approach;
  "Remote Scanner Bottleneck" (L300) already names the batch-size/pipelining problem.
• 2ce0c4377 — Support static filter pushdown through remote scanner (Support static filter pushdown through remote scanner #3795): the template for a new
  RemoteQueryScannerOpen field plus expression serialization, and the effort yardstick.
• ae6f783d9 — Support dynamic filters via remote scanner (Support dynamic filters via remote scanner #4077): the template for mid-scan plan updates
  (RemoteQueryScannerNext.next_predicate + DynamicFilterPhysicalExpr).
• 43bef12c6 — Make scanner open idempotent (Make scanner open idempotent #4813): the two-release wire-evolution pattern.
• datafusion-contrib/datafusion-distributed · Ballista