feat: estimate Iceberg distributed task counts by alexanderbianchi · Pull Request #704 · datafusion-contrib/datafusion-distributed

alexanderbianchi · 2026-09-06T00:13:18Z

Summary

Implements the Iceberg desired-task-count handler using the selected snapshot's total-files-size summary, read from the coordinator-local work-unit feed.

Respect an explicitly selected snapshot, otherwise use the current snapshot.
Estimate tasks using the configured bytes per partition and target partitions, with ceiling division and validation of zero divisors.
Treat a table with no current snapshot as zero scan work; decline to estimate when size metadata is missing, invalid, or unavailable on a remote feed.
Add focused tests using the runtime metadata fixtures from refactor: build Iceberg test metadata at runtime #700, including historical snapshots, empty tables, malformed summaries, and decoded remote feeds.

This is a whole-snapshot scan-work estimate, not a post-pruning byte estimate.

Closes #605.

Dependency / scope

Stacked directly on #700's codex/iceberg-runtime-metadata.

Named test-case cases cover current and selected snapshots, empty tables, and missing or invalid file sizes. A separate distributed-execution test follows tests/task_estimator_test.rs: it executes a taxi aggregation through the in-memory workers. The core harness and its existing tests remain available without the integration feature. Distributed planning and worker setup are explicit builder choices through .with_workers(...), requiring integration; CI enables that feature for the distributed test. The construction cleanup belongs to #700, so this PR no longer introduces and then removes create(). Existing plan and result snapshots remain unchanged.

Validation

cargo test -p datafusion-distributed-iceberg --locked — 101 tests passed, including harness-backed tests and the doctest.
cargo test -p datafusion-distributed-iceberg --features integration --locked — 102 tests passed, including distributed execution and the doctest.
cargo clippy -p datafusion-distributed-iceberg --all-targets -- -D warnings — passed, also with --all-features.
cargo check -p datafusion-distributed-iceberg --lib — passed without the integration feature.
cargo fmt --all -- --check — passed.
git diff --check origin/iceberg-0.10 HEAD — passed.

alexanderbianchi · 2026-09-06T21:13:07Z

+        Ok(())
+    }
+
+    async fn run_distributed_query() -> Result<(String, String)> {


Should this be in the harness?

gabotechs

Flushing a first round of comments, will continue tomorrow, but so far looking pretty good, just small things.

gabotechs · 2026-09-07T17:56:39Z

+    let feed = node.feed().inner()?;
+    let metadata = feed.iceberg_table.metadata();
+    let snapshot = match feed.snapshot_id {
+        Some(id) => Some(metadata.snapshot_by_id(id)?),
+        None => metadata.current_snapshot(),
+    };
+    let total_bytes = match snapshot {
+        Some(snapshot) => snapshot
+            .summary()
+            .additional_properties
+            .get("total-files-size")?
+            .parse()
+            .ok()?,
+        None => 0,
+    };


Here, I'd not access directly the guts of the IcebergDataSource, we should have the total byte size available through node.partition_statistics(None). I'd just use that instead (this is what AQE will use anyways for computing CPU cost)

gabotechs · 2026-09-07T18:01:34Z

+        harness
+            .query("SET datafusion.execution.target_partitions = 2")
+            .await?;


It'd be nice to set this in the harness builder, that way the config under which the test runs is all gathered in the same place.

gabotechs · 2026-09-07T18:01:49Z

+        harness
+            .ctx
+            .set_distributed_file_scan_config_bytes_per_partition(1_000_000)?;


Same as above

gabotechs · 2026-09-07T18:24:23Z

+    async fn scan(harness: &IcebergTestHarness) -> Result<Arc<dyn ExecutionPlan>> {
+        // Call the public table provider so empty-table optimization cannot remove the scan.
+        harness
+            .ctx
+            .table_provider("taxi")
+            .await?
+            .scan(&harness.ctx.state(), None, &[], None)
+            .await
+    }
+
+    fn estimate(plan: &Arc<dyn ExecutionPlan>) -> Result<Option<usize>> {
+        let mut config = SessionConfig::new().with_target_partitions(2);


The main point of the harness is not needing to fill test-case code with helper functions. I think we should be able to bring most (if not all) these helper functions to the test harness (probably with a small tweaks).

Make the context private and configure target partitions, scan byte budgets, and column statistics before building the session. Move scan and estimation plumbing into the harness so tests share its configuration. Adapt task-count fixtures to datafusion-contrib#700's empty_taxi_metadata_builder rename. Keep distributed execution opt-in and preserve existing snapshots.

Read partition_statistics(None) rather than duplicating snapshot selection and summary parsing through the work-unit feed. Preserve absent-size fallback and propagate statistics errors.

alexanderbianchi · 2026-09-07T20:22:59Z


 pub struct IcebergTestHarness {
-    pub ctx: SessionContext,
+    ctx: SessionContext,


Make it private to avoid abuse in the future.

alexanderbianchi · 2026-09-07T20:24:53Z

@@ -45,8 +45,10 @@ mod tests {

    #[tokio::test]
    async fn reports_exact_row_count_and_byte_size_for_full_scan_w_col_stats() -> Result<()> {


Changes needed if we are moving ctx private - can be new PR but I think it's a pretty minor change easy to include here.

alexanderbianchi · 2026-09-08T13:40:04Z

+# Upgrading from 4.0.0 to 5.0.0
+
+## Iceberg test harness
+


sorry, agent did this, will delete. not neede d

Replace test-case attributes with seven named Tokio tests sharing the assertion helper. Remove the unused Iceberg test-case dependency and the upgrade note for the unreleased harness.

gabotechs

👍 Nice

gabotechs · 2026-09-08T15:38:51Z

+    config: SessionConfig,
+    distributed_config: DistributedConfig,
+    column_stats_enabled: bool,


Out of experience when dealing in the past with DataFusion config structs, typically the one that is most helpful for ergonomically building using a builder-pattern in SessionStateBuilder.

I can imagine how holding here a reference to SessionStateBuilder rather than SessionConfig or DistributedConfig can be a bit more future proof.

Store SessionStateBuilder directly and expose one configure_session closure instead of per-setting harness methods. Preserve fixture defaults, private context access, and opt-in worker wiring.

Reapplies the reviewed and merged changes from #703 to the intended base, `iceberg-0.10`. #703 was still targeting `codex/iceberg-runtime-metadata` when it merged, after #700 had already merged into `iceberg-0.10`. Consequently, its codec-test improvements landed only on the old topic branch. This PR cherry-picks #703's merge commit (`eafc234`) onto `iceberg-0.10`, without additional code changes. It retains the original 15-line diff in `iceberg/src/codec.rs`: explicit storage properties, including a quote-containing value, and assertions that those properties survive the codec round trip. Validation: - `cargo test -p datafusion-distributed-iceberg --locked --lib roundtrips_data_source_plan` — passed. - `cargo fmt --all -- --check` and `git diff --check` passed. Targets `iceberg-0.10` directly; it does not depend on the separate reapplication of #704.

Reapplies the reviewed and merged changes from #704 to the intended base, `iceberg-0.10`. #704 was still targeting `codex/iceberg-runtime-metadata` when it merged, after #700 had already merged into `iceberg-0.10`. Consequently, the estimator and harness changes landed only on the old topic branch. This PR cherry-picks #704's merge commit (`d3ba554`) onto `iceberg-0.10`. There are no additional code changes; the resulting tree is identical to #704's final head (`12e1512`). Includes the source-statistics-based task estimator, explicit estimation tests, distributed execution coverage, and the private-context harness with `configure_session(...)`. Validation: - `cargo test -p datafusion-distributed-iceberg --features integration --locked --test desired_task_count` — all 9 tests passed. - `cargo fmt --all -- --check` and `git diff --check` passed. - Verified exact tree equality with #704's final head. Targets `iceberg-0.10` directly; it does not depend on the separate reapplication of #703.

alexanderbianchi changed the base branch from iceberg-0.10 to codex/iceberg-runtime-metadata September 6, 2026 00:14

alexanderbianchi force-pushed the codex/iceberg-runtime-metadata branch from 158065b to cdf47f7 Compare September 6, 2026 00:18

alexanderbianchi force-pushed the iceberg/issue-605 branch from d9c0ea8 to b7668ef Compare September 6, 2026 00:25

alexanderbianchi mentioned this pull request Sep 6, 2026

feat: add local TPC-H Iceberg benchmarks #711

Draft

alexanderbianchi commented Sep 6, 2026

View reviewed changes

alexanderbianchi force-pushed the iceberg/issue-605 branch from b90c672 to 1e709b8 Compare September 6, 2026 22:14

alexanderbianchi marked this pull request as ready for review September 6, 2026 23:00

gabotechs reviewed Sep 7, 2026

View reviewed changes

alexanderbianchi added 8 commits September 7, 2026 16:10

feat: estimate Iceberg distributed task counts (datafusion-contrib#605)

42ea1de

fix: harden Iceberg task count estimation

2f1a64c

test: exercise Iceberg task estimates through metadata fixtures

5f4891a

refactor: parameterize Iceberg task estimation tests

f5b95d5

test: verify distributed Iceberg scan task estimation

c4a4645

refactor: move Iceberg worker setup into the test harness

c123569

refactor: use Iceberg source statistics for task estimates

6f6e928

Read partition_statistics(None) rather than duplicating snapshot selection and summary parsing through the work-unit feed. Preserve absent-size fallback and propagate statistics errors.

alexanderbianchi force-pushed the iceberg/issue-605 branch from 1e709b8 to 6f6e928 Compare September 7, 2026 20:19

alexanderbianchi commented Sep 7, 2026

View reviewed changes

alexanderbianchi commented Sep 8, 2026

View reviewed changes

test: use explicit Iceberg task estimation cases

f66606d

Replace test-case attributes with seven named Tokio tests sharing the assertion helper. Remove the unused Iceberg test-case dependency and the upgrade note for the unreleased harness.

gabotechs approved these changes Sep 8, 2026

View reviewed changes

refactor: customize Iceberg sessions through the native builder

12e1512

Store SessionStateBuilder directly and expose one configure_session closure instead of per-setting harness methods. Preserve fixture defaults, private context access, and opt-in worker wiring.

alexanderbianchi merged commit d3ba554 into datafusion-contrib:codex/iceberg-runtime-metadata Sep 8, 2026
32 checks passed

This was referenced Sep 8, 2026

feat: estimate Iceberg distributed task counts #715

Merged

refactor: make Iceberg codec fixture options explicit #716

Merged

Navigation Menu

Uh oh!

feat: estimate Iceberg distributed task counts - #704

feat: estimate Iceberg distributed task counts#704
alexanderbianchi merged 10 commits into
datafusion-contrib:codex/iceberg-runtime-metadatafrom
alexanderbianchi:iceberg/issue-605

alexanderbianchi commented Sep 6, 2026 •

edited

Loading

Uh oh!

alexanderbianchi Sep 6, 2026

Uh oh!

gabotechs left a comment

Uh oh!

gabotechs Sep 7, 2026

Uh oh!

gabotechs Sep 7, 2026

Uh oh!

gabotechs Sep 7, 2026

Uh oh!

Uh oh!

gabotechs Sep 7, 2026

Uh oh!

alexanderbianchi Sep 7, 2026

Uh oh!

alexanderbianchi Sep 7, 2026

Uh oh!

alexanderbianchi Sep 8, 2026

Uh oh!

gabotechs left a comment

Uh oh!

gabotechs Sep 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -45,8 +45,10 @@ mod tests {

		#[tokio::test]
		async fn reports_exact_row_count_and_byte_size_for_full_scan_w_col_stats() -> Result<()> {

Uh oh!

Conversation

alexanderbianchi commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Dependency / scope

Validation

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

alexanderbianchi commented Sep 6, 2026 •

edited

Loading