fix: size the AQE broadcast decision by bytes, not row count

[andygrove]

Which issue does this PR close?

Closes #2081.

Waiting for #2086 to merge before marking this as ready for review.

Rationale for this change

DynamicJoinSelectionExec::supports_collect_by_thresholds decides whether to collect a join's build side into a CollectLeft broadcast. When total_byte_size is known it compares that against hash_join_single_partition_threshold (10 MB), which is the intent. When total_byte_size is Absent it instead compares the row count against hash_join_single_partition_threshold_rows (1,000,000), and a row count says nothing about size. A build side of up to a million arbitrarily wide rows is replicated to every probe task without the byte threshold ever applying.

Unknown total_byte_size is the common case, not an edge case:

1. Every join discards it. DataFusion's estimate_join_statistics hardcodes total_byte_size: Precision::Absent, keeping only an estimated row count.
2. A single variable-width column loses it permanently. Statistics::calculate_total_byte_size sums DataType::primitive_width() across the schema; that is None for Utf8/LargeUtf8/Utf8View/Binary (and Boolean), so it falls back to total_byte_size.to_inexact() — and to_inexact() on Absent is still Absent.

In TPC-H that covers most dimension-side join results, which carry names, addresses, or comments.

Observed on TPC-H SF1000 (2 executors x 8 cores / 56 GiB, target_partitions=32, prefer_hash_join=false, adaptive planner on) with RUST_LOG=info,ballista_scheduler::state::aqe=debug. Q2 alone:

to_actual_join - plan_id: 1, decision: CollectLeft left: (Inexact(799140) | Absent), right: (Exact(10000000) | Absent)
to_actual_join - plan_id: 7, decision: CollectLeft left: (Inexact(639312) | Absent), right: (Inexact(159969600) | Inexact(3839270400))

Both are broadcast on the row rule alone, at 799,140 and 639,312 rows of unknown size. They are part/supplier-derived rows carrying string columns, so the payload replicated to all 32 tasks is plausibly hundreds of MB rather than the 10 MB the threshold is meant to cap.

The same run shows how easily the size is lost — two decisions over the same 25-row nation table:

• Exact(25) | Exact(400) where the projection is all-numeric, so bytes are recomputed as rows x width.
• Exact(25) | Absent where the schema carries a string column, so the recompute gives up.

What changes are included in this PR?

1. Size the decision by bytes (dynamic_join.rs). When total_byte_size is absent, estimate it and hold it to the same byte threshold. Each column contributes the best figure available: its own ColumnStatistics::byte_size (a total for that column's output, already scaled for filters and limits), else its fixed width times the row count, else a default width mirroring Spark's StringType(20)/BinaryType(100) defaults. Boolean is handled explicitly since primitive_width() reports it as None. An estimate that overflows declines the broadcast rather than wrapping to a small number.

The row threshold is retained as a ceiling, which gives the change a useful property: it can only ever reject a broadcast the row rule would have allowed, never introduce a new one. The worst case is a shuffle where we previously broadcast.

For reference, Spark drives broadcast selection from spark.sql.autoBroadcastJoinThreshold in bytes and estimates a per-row size from the schema (EstimationUtils.getSizePerRow) when it must; it has no row-count threshold standing in for size.

2. Make this class of decision testable without a cluster (test/stats_table.rs, test/broadcast_thresholds.rs). This bug was only visible on a deployed cluster, which is the more serious problem: the decision is a deterministic function of statistics, yet nothing could test it.

The existing tests build MemTables of ~40 Int32 rows, whose statistics are small and always exact, and toggle the decision by zeroing the threshold. Nothing in the suite can express "800,000 rows of unknown size over a string column" — the shape the bug lives in — and make_ctx uses a bare SessionConfig::new(), so the tests never exercise the thresholds Ballista actually ships.

StatsTable declares its statistics and holds no rows, so that shape is one line. Its scan reports the declared figures and cannot be executed, which is all the planner tests need. It deliberately does not recompute total_byte_size on projection, since an unknown size is the case the fixtures exist to express.

The new tests run under SessionConfig::new_with_ballista(), so they exercise the 10 MB / 1,000,000-row thresholds a deployment ships with, and one test pins those defaults directly — a test that sets its own thresholds would pass even if the shipped default were zero, which it was until recently.

Are these changes tested?

Yes.

• wide_rows_of_unknown_size_are_not_broadcast fails on the previous rule and passes with this one — it is a regression test for the reported bug, and it runs in well under a second.
• The other tests pass either way by design: they guard against the estimate rejecting broadcasts it should allow (a 25-row dimension of unknown size, 100k narrow rows, and a known size under the threshold must all still broadcast; a known 64 MB side must not).
• Unit tests on the decision function cover the byte/row/absent/overflow paths directly.
• 267 scheduler tests pass; cargo clippy --all-targets -- -D warnings and cargo fmt --check are clean.

Performance: no measured benefit, and I want to be explicit about that.

Measured on TPC-H SF1000 against an adaptive-planner baseline of 6320.7 s (2 executors x 8 cores / 56 GiB, node-local data, 1 iteration). The run was stopped after 6 queries because the result was not interpretable:

Query	Baseline	This PR
q1 (no joins)	36.2	41.6
q2 (the only query whose plan this changes)	63.6	66.9
q3	255.0	240.3
q4	53.9	60.3
q5	700.5	646.3
q6 (no joins)	12.7	11.9

The two join-free queries moved by +15% and -6%, so the noise floor on this cluster is wider than any effect here. q5 improved 8%, but this PR does not change q5's plan — its unknown-size broadcasts are a 25-row nation and a 1-row aggregate, which pass the estimate — so that improvement cannot be credited to this change. The one query whose plan does change got marginally slower.

So this is offered as a correctness/safety fix, not a performance improvement: it removes an unbounded broadcast that the byte threshold is supposed to cap. It is reasonable to expect it to matter more at wider schemas or higher partition counts, but that is an expectation, not a measurement.

The decision-level effect is verified. On Q2 at SF1000 this PR turns the 799,140-row unknown-size CollectLeft into Partitioned while keeping every genuinely small broadcast (nation at 25 rows, region at 5, 1-row aggregates).

Are there any user-facing changes?

No API or configuration changes. Under the adaptive planner, a build side whose size is unknown and estimated above hash_join_single_partition_threshold is now repartitioned instead of broadcast. hash_join_single_partition_threshold_rows keeps its meaning as a ceiling.

[milenkovicm]

One note, while trying AQE on TPCDS Q72 for some reason bite size is not available in all joins (not sure why) hence AQE makes wrong decisions.

[andygrove]

One note, while trying AQE on TPCDS Q72 for some reason bite size is not available in all joins (not sure why) hence AQE makes wrong decisions.

Yes, that's exactly what I am seeing. I am getting up to speed with the statistics code now.