From a3c5712a5191d97231eb43c852107855e41519f4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 11:44:08 -0400 Subject: [PATCH 01/14] Report physical partitioning, and stop two panics escaping as panics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Groundwork for a multi-library distributed-execution example. Each item here is something that example needs and cannot get today. `ExecutionPlan.output_partitioning` is new. `partition_count` already existed but discards everything except the count, so a driver deciding how to split work across workers could not tell hash-distributed output from merely counted output, nor read the hash keys. It returns a `PhysicalPartitioning`, named to keep it distinct from `datafusion.expr.Partitioning` — that one is the logical partitioning `repartition_by_hash` takes as a request, this one is what a built plan does. Physical expressions have no Python representation, so the hash keys are returned in their displayed form. `SessionContext.execute` now bounds-checks the partition index. The plan's leaves index their partition vector directly, so an out-of-range index reached `MemorySourceConfig` and panicked; the panic was caught as a tokio `JoinError` and arrived as `index out of bounds: the len is 2 but the index is 5`, naming neither the plan nor the index the caller passed. `SessionConfig.set` no longer routes through `SessionConfig::set_str`, which unwraps. An unknown namespace — `datafusion.runtime.*`, or a config extension not yet installed — aborted with a `PanicException`, which derives from `BaseException` and so escapes `except Exception`. `information_schema. df_settings` lists keys in both categories, so replaying settings onto a worker hit this first. Two docstrings on `ExecutionPlan` claimed that a table registered from record batches cannot be serialized. That is true of `LogicalPlan`, whose `try_encode_table_provider` has no arm for one, and false of the physical layer, which inlines the batches: verified by decoding on a context sharing nothing with the encoder and executing. A test pins it, since it is what lets a worker run a plan the driver encoded. Also documents, rather than fixes, the `ForeignExecutionPlan` arm in the example provider's physical codec. It claims every other library's nodes, which the extension guide tells authors not to do — but it is load-bearing: `EnsureCooperative` runs during a foreign planner's `create_physical_plan` and hands the library back a `ForeignExecutionPlan` wrapping the host's `CooperativeExec`, which has no reachable `try_to_proto`. Narrowing the arm makes 31 of the 51 tests in the query-planner example fail, all on that node. The comment now says so, and says a planner that controls its own physical optimizer rules needs no such arm. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 26 +++- crates/core/src/lib.rs | 1 + crates/core/src/physical_plan.rs | 77 ++++++++++ docs/source/user-guide/upgrade-guides.md | 35 +++++ .../src/physical_extension_codec.rs | 21 ++- python/datafusion/__init__.py | 9 +- python/datafusion/plan.py | 131 +++++++++++++++++- python/tests/test_plans.py | 82 ++++++++++- 8 files changed, 369 insertions(+), 13 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index c711a62dc..8aa0504fa 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -44,6 +44,7 @@ use datafusion::execution::options::{ArrowReadOptions, ReadOptions}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::execution::{FunctionRegistry, TaskContextProvider}; +use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::prelude::{ AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, }; @@ -193,8 +194,17 @@ impl PySessionConfig { Self::from(self.config.clone().with_parquet_pruning(enabled)) } - fn set(&self, key: &str, value: &str) -> Self { - Self::from(self.config.clone().set_str(key, value)) + /// Set a config option by key. + /// + /// Not routed through `SessionConfig::set_str`, which unwraps the result: + /// an unknown namespace -- `datafusion.runtime.*`, or a config extension + /// that has not been installed yet -- would abort as a `PanicException` + /// rather than raise. `information_schema.df_settings` lists keys in both + /// of those categories, so replaying it is otherwise unsafe. + fn set(&self, key: &str, value: &str) -> PyDataFusionResult { + let mut config = self.config.clone(); + config.options_mut().set(key, value)?; + Ok(Self::from(config)) } pub fn with_extension(&self, extension: Bound) -> PyResult { @@ -1412,6 +1422,18 @@ impl PySessionContext { ) -> PyDataFusionResult { let ctx: TaskContext = TaskContext::from(&self.ctx.state()); let plan = plan.plan.clone(); + // Checked here because the leaves index their partitions directly: a + // `MemorySourceConfig` panics with a bare `index out of bounds`, which + // surfaces as a `JoinError::Panic` naming neither the plan nor the + // partition the caller asked for. + let partition_count = plan.output_partitioning().partition_count(); + if part >= partition_count { + return Err(PyValueError::new_err(format!( + "Partition index {part} is out of range for a plan with \ + {partition_count} partition(s)" + )) + .into()); + } let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?; Ok(PyRecordBatchStream::new(stream)) } diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 7f0f9cb39..492d57643 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -94,6 +94,7 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index 594655a60..ddb344a96 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -17,6 +17,7 @@ use std::sync::Arc; +use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; use datafusion_proto::physical_plan::AsExecutionPlan; use prost::Message; @@ -126,6 +127,82 @@ impl PyExecutionPlan { pub fn partition_count(&self) -> usize { self.plan.output_partitioning().partition_count() } + + #[getter] + pub fn output_partitioning(&self) -> PyPhysicalPartitioning { + self.plan.output_partitioning().clone().into() + } +} + +/// How a physical plan's output rows are spread across its partitions. +/// +/// Distinct from `datafusion.expr.Partitioning`, which is the *logical* +/// partitioning `DataFrame.repartition` takes as a request. This one reports +/// what a built plan actually does. +#[pyclass( + from_py_object, + frozen, + name = "PhysicalPartitioning", + module = "datafusion", + subclass +)] +#[derive(Debug, Clone)] +pub struct PyPhysicalPartitioning { + partitioning: Partitioning, +} + +#[pymethods] +impl PyPhysicalPartitioning { + /// Which partitioning scheme this is. + /// + /// One of `RoundRobinBatch`, `Hash`, `Range`, or `UnknownPartitioning`. + /// `UnknownPartitioning` is what a plan reports when it knows how many + /// partitions it has but nothing about how rows are distributed between + /// them, which is the common case for a file scan. + #[getter] + pub fn scheme(&self) -> &'static str { + match self.partitioning { + Partitioning::RoundRobinBatch(_) => "RoundRobinBatch", + Partitioning::Hash(_, _) => "Hash", + Partitioning::Range(_) => "Range", + Partitioning::UnknownPartitioning(_) => "UnknownPartitioning", + } + } + + #[getter] + pub fn partition_count(&self) -> usize { + self.partitioning.partition_count() + } + + /// The expressions rows are hashed on, or `None` for other schemes. + /// + /// These are physical expressions, which have no Python representation, so + /// they are returned in their displayed form. + #[getter] + pub fn hash_expressions(&self) -> Option> { + match &self.partitioning { + Partitioning::Hash(exprs, _) => { + Some(exprs.iter().map(|expr| format!("{expr}")).collect()) + } + _ => None, + } + } + + fn __repr__(&self) -> String { + format!("{}", self.partitioning) + } +} + +impl From for PyPhysicalPartitioning { + fn from(partitioning: Partitioning) -> Self { + Self { partitioning } + } +} + +impl From for Partitioning { + fn from(partitioning: PyPhysicalPartitioning) -> Self { + partitioning.partitioning + } } impl From for Arc { diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index f98590b0a..d71f93707 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -169,6 +169,41 @@ installed produces the same bytes as before, as do functions encoded by name. Regenerate any plan you serialized with an earlier release and stored for later use, if it was produced by a session with an extension codec installed. +### Physical plans report their partitioning scheme + +{py:attr}`~datafusion.ExecutionPlan.output_partitioning` is new, and reports +what {py:attr}`~datafusion.ExecutionPlan.partition_count` leaves out: whether +the rows in those partitions are hash-distributed on known keys, spread +round-robin, or merely counted. It returns a +{py:class}`~datafusion.PhysicalPartitioning`. + +```python +partitioning = df.execution_plan().output_partitioning +partitioning.scheme # 'Hash', 'RoundRobinBatch', 'Range', 'UnknownPartitioning' +partitioning.partition_count +partitioning.hash_expressions # display strings, or None +``` + +This is additive; `partition_count` keeps working and agrees with +`output_partitioning.partition_count`. Note that +{py:class}`datafusion.expr.Partitioning` is a different type: that one is the +*logical* partitioning {py:meth}`~datafusion.DataFrame.repartition_by_hash` +takes as a request, while this one is what a built plan actually does. + +### Two error paths that used to abort the interpreter + +{py:meth}`~datafusion.SessionContext.execute` now raises `ValueError` for a +partition index that is out of range. Previously the plan's leaves indexed +their partitions directly and the resulting Rust panic surfaced as an error +naming neither the plan nor the index. + +{py:meth}`~datafusion.SessionConfig.set` now raises for a key whose namespace +does not exist -- `datafusion.runtime.*`, or a config extension that has not +been installed yet. Previously it aborted with a `PanicException`, which +derives from `BaseException` and so escaped `except Exception`. This matters +when replaying settings read back from `information_schema.df_settings`, which +lists keys in both of those categories. + ### Changes to the `datafusion-python-util` crate Extension libraries written in Rust usually depend on the diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index f9e96382e..e8fd19697 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -122,9 +122,24 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - // The provider owns DataSourceExec. A ForeignExecutionPlan can wrap a - // host-added execution decorator around that scan; retaining the opaque - // wrapper preserves its original library identity without downcasting it. + // `DataSourceExec` is this library's own node. The `ForeignExecutionPlan` + // arm is a workaround, not a pattern to copy, and it is load-bearing: + // a host physical optimizer rule that runs during a foreign planner's + // `create_physical_plan` -- `EnsureCooperative` always does -- hands the + // library back a `ForeignExecutionPlan` wrapping the host's + // `CooperativeExec`. That type has no reachable `try_to_proto`, so + // nothing can encode it natively and `FFI_QueryPlanner` must serialize + // the plan it returns. Claiming it here is what lets those plans + // round-trip at all. + // + // The cost is that this codec also claims every *other* library's + // nodes, since that is the type any node arrives as once it has crossed + // the boundary -- see `extension_codec_order`. Narrowing this to + // `DataSourceExec` alone makes 31 tests in + // `datafusion-ffi-query-planner-example` fail with the error above. + // + // A library whose planner controls its own physical optimizer rules + // never sees a foreign node and needs no such arm. if node.is::() || node.is::() { self.counters .encode_execution_plan diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 3696d92a8..1b44f8a73 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -100,7 +100,13 @@ ) from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions -from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet +from .plan import ( + ExecutionPlan, + LogicalPlan, + Metric, + MetricsSet, + PhysicalPartitioning, +) from .record_batch import RecordBatch, RecordBatchStream from .user_defined import ( Accumulator, @@ -133,6 +139,7 @@ "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", + "PhysicalPartitioning", "QueryPlannerExportable", "RecordBatch", "RecordBatchStream", diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 8d03bae2c..49b2035d5 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -34,6 +34,7 @@ "LogicalPlan", "Metric", "MetricsSet", + "PhysicalPartitioning", ] @@ -178,17 +179,54 @@ def __repr__(self) -> str: @property def partition_count(self) -> int: - """Returns the number of partitions in the physical plan.""" + """Returns the number of partitions in the physical plan. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.execution_plan().partition_count + 1 + """ return self._raw_plan.partition_count + @property + def output_partitioning(self) -> PhysicalPartitioning: + """Returns how this plan's output rows are spread across its partitions. + + Where :py:attr:`partition_count` gives only the number of partitions, + this also reports the scheme, so a caller executing partitions + separately can tell whether they are hash-distributed on known keys or + merely counted. See :ref:`distributed_query_engines`. + + Examples: + >>> import pyarrow as pa + >>> from datafusion import SessionConfig, SessionContext + >>> ctx = SessionContext(SessionConfig().with_target_partitions(4)) + >>> ctx.register_record_batches("t", [ + ... [pa.record_batch({"a": [1, 2, 3]})], + ... [pa.record_batch({"a": [4, 5, 6]})], + ... ]) + >>> ctx.sql("select a from t").execution_plan().output_partitioning + UnknownPartitioning(2) + + A group-by redistributes rows, so the plan reports the keys: + + >>> grouped = ctx.sql("select a, count(*) from t group by a") + >>> partitioning = grouped.execution_plan().output_partitioning + >>> partitioning.scheme + 'Hash' + >>> partitioning.partition_count + 4 + """ + return PhysicalPartitioning(self._raw_plan.output_partitioning) + @staticmethod def from_bytes(ctx: SessionContext, data: bytes) -> ExecutionPlan: """Create an ExecutionPlan from serialized protobuf bytes. Decoding routes through the codecs installed on ``ctx`` with :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`. - Tables created in memory from record batches are currently not - supported. Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and positional, and there is no fallback to a worker or global context. @@ -204,8 +242,10 @@ def to_bytes(self, ctx: SessionContext | None = None) -> bytes: When ``ctx`` is supplied, encoding routes through the codecs installed on it with :py:meth:`~datafusion.SessionContext.with_physical_extension_codec`. - Tables created in memory from record batches are currently not - supported. + + Unlike :py:meth:`LogicalPlan.to_bytes`, a plan reading a table + registered from record batches does round-trip: the batches travel + inside the encoded scan. Round-tripping through this method and :py:meth:`from_bytes` is how an extension library checks that its own codec claimed its nodes, @@ -288,6 +328,87 @@ def _walk(node: ExecutionPlan) -> None: return result +class PhysicalPartitioning: + """How a physical plan's output rows are spread across its partitions. + + Returned by :py:attr:`ExecutionPlan.output_partitioning`. This is the + partitioning a built plan *has*, which is different from + :py:class:`datafusion.expr.Partitioning` — the partitioning + :py:meth:`~datafusion.DataFrame.repartition_by_hash` *asks* for. + """ + + def __init__(self, partitioning: df_internal.PhysicalPartitioning) -> None: + """This constructor should not be called by the end user.""" + self._raw_partitioning = partitioning + + @property + def scheme(self) -> str: + """Which partitioning scheme this is. + + One of ``"RoundRobinBatch"``, ``"Hash"``, ``"Range"``, or + ``"UnknownPartitioning"``. A plan reports ``"UnknownPartitioning"`` + when it knows how many partitions it has but nothing about how rows + are distributed between them, which is the usual case for a file scan. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.execution_plan().output_partitioning.scheme + 'UnknownPartitioning' + """ + return self._raw_partitioning.scheme + + @property + def partition_count(self) -> int: + """The number of partitions. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> df.execution_plan().output_partitioning.partition_count + 1 + """ + return self._raw_partitioning.partition_count + + @property + def hash_expressions(self) -> list[str] | None: + """The expressions rows are hashed on, or ``None`` for other schemes. + + Physical expressions have no Python representation, so these are + returned in their displayed form. + + Examples: + >>> import pyarrow as pa + >>> from datafusion import SessionConfig, SessionContext + >>> ctx = SessionContext(SessionConfig().with_target_partitions(4)) + >>> ctx.register_record_batches("t", [ + ... [pa.record_batch({"a": [1, 2, 3]})], + ... [pa.record_batch({"a": [4, 5, 6]})], + ... ]) + >>> scan = ctx.sql("select a from t").execution_plan() + >>> scan.output_partitioning.hash_expressions is None + True + >>> grouped = ctx.sql("select a, count(*) from t group by a") + >>> grouped.execution_plan().output_partitioning.hash_expressions + ['a@0'] + """ + return self._raw_partitioning.hash_expressions + + def __repr__(self) -> str: + """Print a string representation of the partitioning. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> df = ctx.from_pydict({"a": [1, 2, 3]}) + >>> repr(df.execution_plan().output_partitioning) + 'UnknownPartitioning(1)' + """ + return self._raw_partitioning.__repr__() + + class MetricsSet: """A set of metrics for a single execution plan operator. diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 0145d123e..95b5a77f4 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -24,14 +24,15 @@ LogicalPlan, Metric, MetricsSet, + SessionConfig, SessionContext, col, udf, ) -# Note: We must use CSV because memory tables are currently not supported for -# conversion to/from protobuf. +# Note: CSV because a *logical* plan cannot carry a memory table. The physical +# layer can — see `test_execution_plan_over_memory_batches_round_trips`. @pytest.fixture def df(): ctx = SessionContext() @@ -95,6 +96,83 @@ def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None: assert df.collect() == df_round_trip.collect() +def test_execution_plan_over_memory_batches_round_trips() -> None: + """A physical plan reading record batches decodes on an unrelated session. + + Only the *logical* layer cannot carry a memory table: its + `try_encode_table_provider` has no arm for one. The physical scan inlines + the batches, so it needs neither a shared session nor an extension codec — + which is what lets a worker process execute a plan the driver encoded. + """ + ctx = SessionContext() + ctx.register_record_batches( + "t", + [[pa.record_batch({"a": [1, 2, 3]})], [pa.record_batch({"a": [4, 5, 6]})]], + ) + plan_bytes = ctx.sql("select a from t").execution_plan().to_bytes(ctx) + + # A session that shares nothing with the encoder: no codecs, no tables. + fresh = SessionContext() + decoded = ExecutionPlan.from_bytes(fresh, plan_bytes) + rows = sum( + batch.to_pyarrow().num_rows + for partition in range(decoded.partition_count) + for batch in fresh.execute(decoded, partition) + ) + assert rows == 6 + + +def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None: + """`output_partitioning` distinguishes hash-distributed output from counted.""" + ctx = SessionContext(SessionConfig().with_target_partitions(4)) + ctx.register_record_batches( + "t", + [[pa.record_batch({"a": [1, 2, 3]})], [pa.record_batch({"a": [4, 5, 6]})]], + ) + + scan = ctx.sql("select a from t").execution_plan() + assert scan.output_partitioning.scheme == "UnknownPartitioning" + assert scan.output_partitioning.hash_expressions is None + # Agrees with the count-only accessor it supplements. + assert scan.output_partitioning.partition_count == scan.partition_count + + grouped = ctx.sql("select a, count(*) from t group by a").execution_plan() + partitioning = grouped.output_partitioning + assert partitioning.scheme == "Hash" + assert partitioning.hash_expressions == ["a@0"] + assert partitioning.partition_count == 4 + assert repr(partitioning) == "Hash([a@0], 4)" + + +def test_execute_rejects_an_out_of_range_partition() -> None: + """An out-of-range partition index raises instead of panicking. + + The leaves index their partition vector directly, so without this check a + bad index surfaces as a `JoinError::Panic` carrying `index out of bounds` + and naming neither the plan nor the index requested. + """ + ctx = SessionContext() + ctx.register_record_batches("t", [[pa.record_batch({"a": [1, 2, 3]})]]) + plan = ctx.sql("select a from t").execution_plan() + assert plan.partition_count == 1 + + with pytest.raises(ValueError, match="Partition index 5 is out of range"): + ctx.execute(plan, 5) + + +def test_session_config_set_rejects_an_unknown_namespace() -> None: + """A bad config key raises rather than aborting through a Rust panic. + + `datafusion.runtime.*` appears in `information_schema.df_settings` but has + no `ConfigOptions` namespace, so it is the key a naive "read the settings + back and replay them on the worker" loop hits first. + """ + with pytest.raises(Exception, match="runtime") as excinfo: + SessionConfig().set("datafusion.runtime.memory_limit", "unlimited") + # A panic would arrive as BaseException, escaping `except Exception`. + assert isinstance(excinfo.value, Exception) + + def test_installing_a_physical_codec_preserves_strict_mode() -> None: """Installing a physical extension codec must not re-enable inlining. From 4c097c70a1d89f72f2236974b27bcc24920cf2ca Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 09:35:45 -0400 Subject: [PATCH 02/14] Name the execute partition index for what it is, trim the upgrade guide `SessionContext.execute` took its second argument as `partitions`, which reads as a count when it is a single partition index. Rename it to `partition` and give the method a real docstring with a doctest covering both a full sweep over `partition_count` and the out-of-range `ValueError`. Every call site in the repo, docs, and examples passes it positionally, so add a short note to the upgrade guide for anyone passing it by keyword. Drop two changelog-shaped sections from the upgrade guide. `output_partitioning` is additive and `SessionContext.execute` / `SessionConfig.set` only trade a panic for a raise, so neither asks the reader to change anything. The `with_extensions` recommendation goes for the same reason; it is advice, and the extension guide already carries it under `extension_bundles`. Also remove the comment above the partition range check. It explained the check by way of a `MemorySourceConfig` panic, which reads as a complaint about DataFusion's leaves. The index arrives straight from Python and no planner has seen it, so validating it needs no more justification than any other argument check at the boundary. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 4 --- docs/source/user-guide/upgrade-guides.md | 46 ++++-------------------- python/datafusion/context.py | 39 ++++++++++++++++++-- python/tests/test_plans.py | 7 +--- 4 files changed, 43 insertions(+), 53 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 8aa0504fa..940512321 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1422,10 +1422,6 @@ impl PySessionContext { ) -> PyDataFusionResult { let ctx: TaskContext = TaskContext::from(&self.ctx.state()); let plan = plan.plan.clone(); - // Checked here because the leaves index their partitions directly: a - // `MemorySourceConfig` panics with a bare `index out of bounds`, which - // surfaces as a `JoinError::Panic` naming neither the plan nor the - // partition the caller asked for. let partition_count = plan.output_partitioning().partition_count(); if part >= partition_count { return Err(PyValueError::new_err(format!( diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index d71f93707..d60653a21 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -96,16 +96,6 @@ way `add_physical_optimizer_rule` does and returns nothing — the query planner lives in `SessionState`, so it belongs to the session rather than to a particular handle on it. See {ref}`extension_planners` for the full protocol. -If a library ships codecs *and* a planner, prefer -`SessionContext.with_extensions(bundle)` over installing each piece by hand. It -installs every codec before it binds any planner, so a planner cannot end up -carrying a chain that a later `with_logical_extension_codec` call has grown. -The library exposes a bundle object implementing -`__datafusion_session_components__` for its codecs and -`__datafusion_session_planner__` for its planner — the latter is handed the -planner installed so far, so several libraries that each ship one nest instead -of displacing each other. See {ref}`extension_bundles`. - (extension_version_mismatch)= ### Mismatched extension libraries now fail loudly @@ -169,41 +159,17 @@ installed produces the same bytes as before, as do functions encoded by name. Regenerate any plan you serialized with an earlier release and stored for later use, if it was produced by a session with an extension codec installed. -### Physical plans report their partitioning scheme +### `SessionContext.execute` renamed its second parameter -{py:attr}`~datafusion.ExecutionPlan.output_partitioning` is new, and reports -what {py:attr}`~datafusion.ExecutionPlan.partition_count` leaves out: whether -the rows in those partitions are hash-distributed on known keys, spread -round-robin, or merely counted. It returns a -{py:class}`~datafusion.PhysicalPartitioning`. +The parameter is a single partition index, not a count, and is now named +`partition` rather than `partitions`. Positional calls are unaffected; update +any call passing it by keyword. ```python -partitioning = df.execution_plan().output_partitioning -partitioning.scheme # 'Hash', 'RoundRobinBatch', 'Range', 'UnknownPartitioning' -partitioning.partition_count -partitioning.hash_expressions # display strings, or None +ctx.execute(plan, partitions=0) # before +ctx.execute(plan, partition=0) # after ``` -This is additive; `partition_count` keeps working and agrees with -`output_partitioning.partition_count`. Note that -{py:class}`datafusion.expr.Partitioning` is a different type: that one is the -*logical* partitioning {py:meth}`~datafusion.DataFrame.repartition_by_hash` -takes as a request, while this one is what a built plan actually does. - -### Two error paths that used to abort the interpreter - -{py:meth}`~datafusion.SessionContext.execute` now raises `ValueError` for a -partition index that is out of range. Previously the plan's leaves indexed -their partitions directly and the resulting Rust panic surfaced as an error -naming neither the plan nor the index. - -{py:meth}`~datafusion.SessionConfig.set` now raises for a key whose namespace -does not exist -- `datafusion.runtime.*`, or a config extension that has not -been installed yet. Previously it aborted with a `PanicException`, which -derives from `BaseException` and so escaped `except Exception`. This matters -when replaying settings read back from `information_schema.df_settings`, which -lists keys in both of those categories. - ### Changes to the `datafusion-python-util` crate Extension libraries written in Rust usually depend on the diff --git a/python/datafusion/context.py b/python/datafusion/context.py index bbf08e84e..2f9db53a6 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2351,9 +2351,42 @@ def read_table( """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table.""" return DataFrame(self.ctx.read_table(table)) - def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream: - """Execute the ``plan`` and return the results.""" - return RecordBatchStream(self.ctx.execute(plan._raw_plan, partitions)) + def execute(self, plan: ExecutionPlan, partition: int) -> RecordBatchStream: + """Execute a single partition of ``plan`` and stream its batches. + + Args: + plan: The physical plan to execute. + partition: Index of the partition to execute, in + ``range(plan.partition_count)``. + + Returns: + A stream over the record batches that partition produces. + + Raises: + ValueError: If ``partition`` is not a valid index for ``plan``. + + Example usage: + + >>> import pyarrow as pa + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx.register_record_batches( + ... "t", [[pa.record_batch({"a": [1, 2]})], [pa.record_batch({"a": [3]})]] + ... ) + >>> plan = ctx.sql("select a from t").execution_plan() + >>> plan.partition_count + 2 + >>> sum( + ... batch.to_pyarrow().num_rows + ... for p in range(plan.partition_count) + ... for batch in ctx.execute(plan, p) + ... ) + 3 + >>> ctx.execute(plan, 2) + Traceback (most recent call last): + ValueError: Partition index 2 is out of range for a plan with 2 partition(s) + """ + return RecordBatchStream(self.ctx.execute(plan._raw_plan, partition)) @staticmethod def _convert_file_sort_order( diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 95b5a77f4..65ba3a2ac 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -145,12 +145,7 @@ def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None: def test_execute_rejects_an_out_of_range_partition() -> None: - """An out-of-range partition index raises instead of panicking. - - The leaves index their partition vector directly, so without this check a - bad index surfaces as a `JoinError::Panic` carrying `index out of bounds` - and naming neither the plan nor the index requested. - """ + """An out-of-range partition index raises instead of panicking.""" ctx = SessionContext() ctx.register_record_batches("t", [[pa.record_batch({"a": [1, 2, 3]})]]) plan = ctx.sql("select a from t").execution_plan() From 846fe6e14c781718447a88b9264fcbd92e47db00 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 09:55:01 -0400 Subject: [PATCH 03/14] Link the upstream FFI issue, stop teaching execute by counter-example The `ForeignExecutionPlan` arm's comment explained the symptom but left the reader no way to find out whether the workaround is still needed. Name the upstream umbrella issue, apache/datafusion#25152, and the cascade behind it: `FFI_PlanProperties` carries no `scheduling_type`, so `EnsureCooperative` reads every foreign leaf as non-cooperative and wraps it, and the resulting `ForeignExecutionPlan` then cannot serialize itself. Fixing either half retires the arm. Drop the out-of-range call from `SessionContext.execute`'s doctest. A docstring example shows a reader how to use the method, and this one put a wrong call in front of them; the message it asserted is already pinned by `test_execute_rejects_an_out_of_range_partition`. The `Raises:` section is the right home for that behaviour, so complete it: a negative or oversized index raises `OverflowError` from the `usize` conversion, not the `ValueError` the bounds check produces. Co-Authored-By: Claude Opus 5 (1M context) --- .../datafusion-ffi-example/src/physical_extension_codec.rs | 7 +++++++ python/datafusion/context.py | 5 ++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index e8fd19697..810e090de 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -140,6 +140,13 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { // // A library whose planner controls its own physical optimizer rules // never sees a foreign node and needs no such arm. + // + // Both halves are upstream defects, tracked together in + // https://github.com/apache/datafusion/issues/25152: `FFI_PlanProperties` + // carries no `scheduling_type`, so `EnsureCooperative` reads every + // foreign leaf as non-cooperative and wraps it, and the resulting + // `ForeignExecutionPlan` then has no way to serialize itself. Fixing + // either one retires this arm. if node.is::() || node.is::() { self.counters .encode_execution_plan diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 2f9db53a6..57669c556 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2364,6 +2364,8 @@ def execute(self, plan: ExecutionPlan, partition: int) -> RecordBatchStream: Raises: ValueError: If ``partition`` is not a valid index for ``plan``. + OverflowError: If ``partition`` is negative, or too large to fit a + platform-sized unsigned integer. Example usage: @@ -2382,9 +2384,6 @@ def execute(self, plan: ExecutionPlan, partition: int) -> RecordBatchStream: ... for batch in ctx.execute(plan, p) ... ) 3 - >>> ctx.execute(plan, 2) - Traceback (most recent call last): - ValueError: Partition index 2 is out of range for a plan with 2 partition(s) """ return RecordBatchStream(self.ctx.execute(plan._raw_plan, partition)) From ebf4bc036b266f59db44264e8e7b076415378ec7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:03:07 -0400 Subject: [PATCH 04/14] Document what SessionConfig.set actually does The docstring was wrong in three ways at once, and the panic fix earlier in this branch is what made it worth opening: it promised "a new SessionConfig object" when the method mutates in place and returns self, it mis-indented the `Args` entries so Sphinx rendered them as body text rather than a field list, and it had no `Raises` at all -- so the one behaviour this branch changed, an unknown key raising instead of aborting the interpreter, was documented only in the Rust source that no user reads. Give it a truthful `Returns`, a `Raises` covering both an unknown key and an unparsable value, and a doctest that reads the option back out through `information_schema.df_settings`. The `datafusion.runtime.*` trap goes to `configuration.md`, which is where a reader is when they need it: those keys appear in `df_settings` but are not settable, so replaying that table verbatim onto a worker fails on the first such row. The docstring states it in one sentence and points at the guide. That page also now records that the `with_*` methods modify in place, which is what makes the chained style it already demonstrates work. Note the same false `Returns` line appears on every other `with_*` method on this class; correcting those is a separate sweep, not this branch's business. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/user-guide/configuration.md | 23 ++++++++++++++++++ python/datafusion/context.py | 31 +++++++++++++++++++++---- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/docs/source/user-guide/configuration.md b/docs/source/user-guide/configuration.md index d1c5c9b44..7eaf7cc13 100644 --- a/docs/source/user-guide/configuration.md +++ b/docs/source/user-guide/configuration.md @@ -49,6 +49,29 @@ ctx = SessionContext(config, runtime) print(ctx) ``` +## Setting options by key + +The `with_*` methods cover the common options, but any option DataFusion declares can be +set by its fully qualified key with {py:meth}`~datafusion.SessionConfig.set`. The value is +always a string, and is parsed according to the type the option declares, so an unknown key +or an unparsable value raises rather than being silently ignored: + +```python +config = SessionConfig().set("datafusion.execution.batch_size", "1024") +``` + +Every method above modifies the config in place and returns it, which is what makes the +chained style work — the object you started with is the object you end up passing to +`SessionContext`. + +One trap is worth knowing about if you read settings back out of a session and replay them +somewhere else, such as onto a worker process or into a test fixture. With +`with_information_schema(True)`, the `information_schema.df_settings` table lists the +`datafusion.runtime.*` keys alongside the rest, but those come from the runtime environment +rather than from `ConfigOptions` and cannot be set with `set`. Feeding that table's rows +back in verbatim will fail on the first such row. Configure the runtime through +`RuntimeEnvBuilder` instead, and skip the `datafusion.runtime.` prefix when replaying. + ## Maximizing CPU Usage DataFusion uses partitions to parallelize work. For small queries the diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 57669c556..ec55c3fd0 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -341,14 +341,37 @@ def with_parquet_pruning(self, enabled: bool = True) -> SessionConfig: return self def set(self, key: str, value: str) -> SessionConfig: - """Set a configuration option. + """Set a configuration option by its fully qualified key. + + Not every key that ``information_schema.df_settings`` lists can be set + here: the ``datafusion.runtime.*`` entries come from the runtime + environment rather than from the session config. See + :ref:`configuration`. Args: - key: Option key. - value: Option value. + key: Option key including its namespace, such as + ``datafusion.execution.batch_size``. + value: Option value as a string, parsed according to the type the + option declares. Returns: - A new :py:class:`SessionConfig` object with the updated setting. + This :py:class:`SessionConfig`, modified in place, so that calls + chain. + + Raises: + Exception: If ``key`` names no known option, or if ``value`` does + not parse as that option's declared type. + + Example usage: + + >>> from datafusion import SessionConfig, SessionContext + >>> config = SessionConfig().set("datafusion.execution.batch_size", "1024") + >>> ctx = SessionContext(config.with_information_schema(True)) + >>> ctx.sql( + ... "select value from information_schema.df_settings" + ... " where name = 'datafusion.execution.batch_size'" + ... ).collect()[0]["value"][0] + """ self.config_internal = self.config_internal.set(key, value) return self From d30079aaae49ee895e0c31d0705de4b1c82c5def Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:07:16 -0400 Subject: [PATCH 05/14] Raise ValueError from SessionConfig.set, not a bare Exception Trading a `PanicException` for an untyped `Exception` only half-fixed the problem: a caller replaying `information_schema.df_settings` still could not catch the failure without swallowing every other error this crate raises, and was left matching on message text. A rejected key or an unparsable value is an argument error, so it should raise `ValueError` -- the same thing an out-of-range partition index gets from `execute` two methods away. Map through `from_datafusion_error`, which already produces `PyValueError` and which `SessionContext.sql` already uses in this file, rather than propagating a `PyDataFusionError` and taking its blanket conversion. That conversion is deliberately untouched: reclassifying every error out of this crate is a much wider change with its own compatibility story. The test can now assert something. It previously ended in `assert isinstance(excinfo.value, Exception)` under a `pytest.raises(Exception, ...)` that had already proven exactly that, so the line could never fail. Assert the type instead, and add a case for a known key with a value of the wrong type, which reaches the same path by a route a settings-replay loop is just as likely to take. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 12 ++++++++++-- python/datafusion/context.py | 2 +- python/tests/test_plans.py | 13 ++++++++++--- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 940512321..b14640d72 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -201,9 +201,17 @@ impl PySessionConfig { /// that has not been installed yet -- would abort as a `PanicException` /// rather than raise. `information_schema.df_settings` lists keys in both /// of those categories, so replaying it is otherwise unsafe. - fn set(&self, key: &str, value: &str) -> PyDataFusionResult { + /// + /// Mapped with `from_datafusion_error` rather than propagated as a + /// `PyDataFusionError`, whose blanket conversion yields a bare `Exception`. + /// A rejected key or value is an argument error, so it raises `ValueError` + /// the way an out-of-range partition index does in `execute`. + fn set(&self, key: &str, value: &str) -> PyResult { let mut config = self.config.clone(); - config.options_mut().set(key, value)?; + config + .options_mut() + .set(key, value) + .map_err(from_datafusion_error)?; Ok(Self::from(config)) } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index ec55c3fd0..020ea06eb 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -359,7 +359,7 @@ def set(self, key: str, value: str) -> SessionConfig: chain. Raises: - Exception: If ``key`` names no known option, or if ``value`` does + ValueError: If ``key`` names no known option, or if ``value`` does not parse as that option's declared type. Example usage: diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 65ba3a2ac..226c65125 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -162,10 +162,17 @@ def test_session_config_set_rejects_an_unknown_namespace() -> None: no `ConfigOptions` namespace, so it is the key a naive "read the settings back and replay them on the worker" loop hits first. """ - with pytest.raises(Exception, match="runtime") as excinfo: + # `ValueError`, not a bare `Exception`: a panic would arrive as + # `PanicException`, which derives from `BaseException` and so would not be + # caught here at all. + with pytest.raises(ValueError, match="runtime"): SessionConfig().set("datafusion.runtime.memory_limit", "unlimited") - # A panic would arrive as BaseException, escaping `except Exception`. - assert isinstance(excinfo.value, Exception) + + +def test_session_config_set_rejects_an_unparsable_value() -> None: + """A well-known key with a value of the wrong type raises too.""" + with pytest.raises(ValueError, match="batch_size"): + SessionConfig().set("datafusion.execution.batch_size", "not_an_int") def test_installing_a_physical_codec_preserves_strict_mode() -> None: From 9714734a3f62ec37ef0a900109ee8c7fcb193835 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:13:01 -0400 Subject: [PATCH 06/14] Point output_partitioning at a page that discusses partitioning The docstring referred the reader to `distributed_query_engines`, which is wrong twice over: that page's premise is that you do *not* partition by hand because the engine does it for you, and it documents work that is not yet usable from datafusion-python. A reader following the link to find out what a scheme means landed on a status page for a different road. Nothing under docs/source/ discussed plan partitioning from the caller's side, so give the claim a home on the page that already tells readers to call `repartition` and `repartition_by_hash` to keep their cores busy -- and that, until now, gave them no way to check whether it worked. The new section is worth more than a pointer. `repartition_by_hash(col("a"), num=8)` followed by an aggregation reports `Hash([a@0], 16)`: the optimizer discarded the requested repartition and inserted its own at `target_partitions`, so neither the scheme nor the count is what was asked for. Drop the aggregation and the repartition vanishes entirely, leaving `UnknownPartitioning` over the source's partition count. Both were verified against a Parquet source, and both are invisible without this accessor, which is the concrete form of the have-versus-ask distinction the class docstring asserts. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/user-guide/configuration.md | 37 +++++++++++++++++++++++++ python/datafusion/plan.py | 3 +- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/docs/source/user-guide/configuration.md b/docs/source/user-guide/configuration.md index 7eaf7cc13..398a874c0 100644 --- a/docs/source/user-guide/configuration.md +++ b/docs/source/user-guide/configuration.md @@ -119,6 +119,43 @@ df = df.repartition_by_hash(col("a"), num=16) result = df.collect() ``` +(checking_partitioning)= + +### Checking what the plan actually does + +`repartition` and `repartition_by_hash` are requests, not instructions. The optimizer is +free to drop a repartition nothing downstream needs, to collapse partitions again for an +operator that requires a single stream, or to substitute a repartition of its own sized by +`target_partitions`. So the number you passed is not necessarily the number you get. + +{py:attr}`~datafusion.ExecutionPlan.output_partitioning` reports what the built plan does, +as opposed to what was asked of it: + +```python +from datafusion import SessionConfig, SessionContext, col, functions as f + +config = SessionConfig().with_target_partitions(16) +ctx = SessionContext(config) + +df = ctx.read_parquet("data.parquet").repartition_by_hash(col("a"), num=8) +plan = df.aggregate([col("a")], [f.sum(col("b"))]).execution_plan() + +partitioning = plan.output_partitioning +print(partitioning.scheme) # 'Hash' +print(partitioning.partition_count) # 16 -- target_partitions, not the 8 requested +print(partitioning.hash_expressions) # ['a@0'] +``` + +The request for eight partitions did not survive: the optimizer inserted its own hash +repartition at `target_partitions` instead. Had the aggregation been left off, the +repartition would have been removed altogether and the plan would report +`UnknownPartitioning` over the source's own partition count. + +`UnknownPartitioning` means the plan knows how many partitions it has but nothing about how +rows are distributed across them, which is the ordinary case for a file scan. +{py:attr}`~datafusion.ExecutionPlan.partition_count` gives the same count on its own when +the scheme does not matter. + ### Benchmark Example The repository includes a benchmark script that demonstrates how to maximize CPU usage diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 49b2035d5..4e4ab3956 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -197,7 +197,8 @@ def output_partitioning(self) -> PhysicalPartitioning: Where :py:attr:`partition_count` gives only the number of partitions, this also reports the scheme, so a caller executing partitions separately can tell whether they are hash-distributed on known keys or - merely counted. See :ref:`distributed_query_engines`. + merely counted. A plan does not necessarily partition the way it was + asked to; see :ref:`checking_partitioning`. Examples: >>> import pyarrow as pa From baf828fa815f5ae58694a40402cd5e926cc2d33a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:18:26 -0400 Subject: [PATCH 07/14] Drop PhysicalPartitioning's unused inbound conversion `From for Partitioning` had no callers. The type is a read-only report of what a built plan does, nothing accepts a partitioning as an argument, and a Rust caller wanting the `Partitioning` reads it off the plan, so there is no inbound direction to support. Removing the impl alone left a deprecation warning: pyo3 auto-derives `FromPyObject` for a `#[pyclass]` that implements `Clone` and now wants the choice made explicitly. Omission is not the way to decline it, so say `skip_from_py_object`, which is what the example crate's providers already use. Clippy is clean again, and `--all-targets` keeps it that way. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/physical_plan.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index ddb344a96..5b03652be 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -139,8 +139,11 @@ impl PyExecutionPlan { /// Distinct from `datafusion.expr.Partitioning`, which is the *logical* /// partitioning `DataFrame.repartition` takes as a request. This one reports /// what a built plan actually does. +// `skip_from_py_object` because this is a read-only report: nothing accepts a +// partitioning as an argument, so there is no inbound direction to support. +// Rust callers that need the `Partitioning` read it off the plan instead. #[pyclass( - from_py_object, + skip_from_py_object, frozen, name = "PhysicalPartitioning", module = "datafusion", @@ -199,12 +202,6 @@ impl From for PyPhysicalPartitioning { } } -impl From for Partitioning { - fn from(partitioning: PyPhysicalPartitioning) -> Self { - partitioning.partitioning - } -} - impl From for Arc { fn from(plan: PyExecutionPlan) -> Arc { plan.plan.clone() From fd0486392c5ff2ed80814982a4709d71c0437e49 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:25:08 -0400 Subject: [PATCH 08/14] Cover RoundRobinBatch, stop claiming a plan can report Range The `scheme` docstring listed four values as though a plan could report any of them. Two of the four needed opposite corrections. `RoundRobinBatch` is reachable, and now tested. It is easy to miss because it never survives at the root: the optimizer inserts one only above a source with fewer partitions than `target_partitions` and CPU work above it to parallelize, so a single-file Parquet scan under a filter and a grouped aggregate produces `RepartitionExec: partitioning=RoundRobinBatch(8)` three levels down. The test walks `children` and asserts all three schemes in that tree, which also pins that the accessor reads each node's own partitioning rather than the root's. `Range` is the other way: it is in the upstream enum, but nothing constructs one in a physical plan. `RangePartitioning` says optimizer and execution support is deliberately unimplemented, per apache/datafusion#22395. Say so, and say why the match arm exists anyway -- it keeps this getter compiling when that support lands, and dropping it would make the match non-exhaustive. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/physical_plan.rs | 10 ++++++++-- python/datafusion/plan.py | 9 +++++---- python/tests/test_plans.py | 27 +++++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 6 deletions(-) diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index 5b03652be..f22e59fa6 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -158,10 +158,16 @@ pub struct PyPhysicalPartitioning { impl PyPhysicalPartitioning { /// Which partitioning scheme this is. /// - /// One of `RoundRobinBatch`, `Hash`, `Range`, or `UnknownPartitioning`. /// `UnknownPartitioning` is what a plan reports when it knows how many /// partitions it has but nothing about how rows are distributed between - /// them, which is the common case for a file scan. + /// them, which is the common case for a file scan. `RoundRobinBatch` and + /// `Hash` come from a `RepartitionExec`. + /// + /// `Range` is in the upstream enum but no plan reports it yet: optimizer + /// and execution support is deliberately unimplemented, per + /// . The arm is here so + /// this getter keeps compiling when that lands, not because it is + /// reachable today. #[getter] pub fn scheme(&self) -> &'static str { match self.partitioning { diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 4e4ab3956..39ec801ac 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -346,10 +346,11 @@ def __init__(self, partitioning: df_internal.PhysicalPartitioning) -> None: def scheme(self) -> str: """Which partitioning scheme this is. - One of ``"RoundRobinBatch"``, ``"Hash"``, ``"Range"``, or - ``"UnknownPartitioning"``. A plan reports ``"UnknownPartitioning"`` - when it knows how many partitions it has but nothing about how rows - are distributed between them, which is the usual case for a file scan. + ``"UnknownPartitioning"`` means the plan knows how many partitions it + has but nothing about how rows are distributed between them, which is + the usual case for a file scan. ``"RoundRobinBatch"`` and ``"Hash"`` + come from a repartition the optimizer inserted. ``"Range"`` is defined + upstream but no plan reports it yet. Examples: >>> from datafusion import SessionContext diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 226c65125..ae44700bb 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -18,6 +18,7 @@ import datetime import pyarrow as pa +import pyarrow.parquet as pq import pytest from datafusion import ( ExecutionPlan, @@ -144,6 +145,32 @@ def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None: assert repr(partitioning) == "Hash([a@0], 4)" +def test_output_partitioning_reports_round_robin(tmp_path) -> None: + """A round-robin repartition reports `RoundRobinBatch`. + + The optimizer only inserts one above a source with fewer partitions than + `target_partitions` and CPU work above it to parallelize, and it never + survives at the root, so reach it by walking `children`. + """ + path = tmp_path / "rr.parquet" + pq.write_table(pa.table({"a": list(range(2000)), "b": [1] * 2000}), path) + + ctx = SessionContext(SessionConfig().with_target_partitions(8)) + ctx.register_parquet("t", str(path)) + plan = ctx.sql("select a, sum(b) from t where a > 5 group by a").execution_plan() + + schemes = set() + stack = [plan] + while stack: + node = stack.pop() + schemes.add(node.output_partitioning.scheme) + stack.extend(node.children()) + + # The single-file scan, the round-robin above it, and the hash repartition + # for the grouping are all in one tree. + assert schemes == {"UnknownPartitioning", "RoundRobinBatch", "Hash"} + + def test_execute_rejects_an_out_of_range_partition() -> None: """An out-of-range partition index raises instead of panicking.""" ctx = SessionContext() From 0ddfbc43a63358deb56f8c11585d12e0b11aa882 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:29:30 -0400 Subject: [PATCH 09/14] Tighten three small things around output_partitioning `scheme` returns one of exactly four strings, so annotate it `Literal` rather than `str`. The promise is safe to make: the Rust getter matches exhaustively over `Partitioning`, so a new upstream variant is a compile error here before it can be a lie in the type. Bind `output_partitioning` once in the scan half of its test. It was read three times, and each read clones the `Partitioning` and builds a fresh wrapper. Validate the partition index in `execute` before building the `TaskContext`, not after. Nothing observable changes; the rejected call just stops doing setup work it is about to throw away. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 2 +- python/datafusion/plan.py | 6 ++++-- python/tests/test_plans.py | 7 ++++--- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index b14640d72..dc57d8f37 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1428,7 +1428,6 @@ impl PySessionContext { part: usize, py: Python, ) -> PyDataFusionResult { - let ctx: TaskContext = TaskContext::from(&self.ctx.state()); let plan = plan.plan.clone(); let partition_count = plan.output_partitioning().partition_count(); if part >= partition_count { @@ -1438,6 +1437,7 @@ impl PySessionContext { )) .into()); } + let ctx: TaskContext = TaskContext::from(&self.ctx.state()); let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?; Ok(PyRecordBatchStream::new(stream)) } diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 39ec801ac..72fcbb932 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -20,7 +20,7 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import datafusion._internal as df_internal @@ -343,7 +343,9 @@ def __init__(self, partitioning: df_internal.PhysicalPartitioning) -> None: self._raw_partitioning = partitioning @property - def scheme(self) -> str: + def scheme( + self, + ) -> Literal["RoundRobinBatch", "Hash", "Range", "UnknownPartitioning"]: """Which partitioning scheme this is. ``"UnknownPartitioning"`` means the plan knows how many partitions it diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index ae44700bb..07072a1f0 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -132,10 +132,11 @@ def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None: ) scan = ctx.sql("select a from t").execution_plan() - assert scan.output_partitioning.scheme == "UnknownPartitioning" - assert scan.output_partitioning.hash_expressions is None + scanned = scan.output_partitioning + assert scanned.scheme == "UnknownPartitioning" + assert scanned.hash_expressions is None # Agrees with the count-only accessor it supplements. - assert scan.output_partitioning.partition_count == scan.partition_count + assert scanned.partition_count == scan.partition_count grouped = ctx.sql("select a, count(*) from t group by a").execution_plan() partitioning = grouped.output_partitioning From 2de97dd7df6448e58d39d5e1851fa595963f906a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:55:47 -0400 Subject: [PATCH 10/14] Stop SessionConfig's constructor aborting on an unknown key `SessionConfig.set` was fixed to raise instead of panicking, but the constructor still routed its dictionary through `SessionConfig::set`, which forwards to `set_str` and unwraps. So `SessionConfig({"datafusion.runtime.memory_limit": "unlimited"})` aborted as a `PanicException` -- a `BaseException`, past any `except Exception` -- while the same key passed to `set` raised `ValueError`. The constructor is the likelier of the two to meet a bad key: it takes a `dict[str, str]`, which is the shape a replayed set of settings arrives in, and `information_schema.df_settings` lists eight `datafusion.runtime.*` rows that cannot be set from a session config at all. Route each entry through `options_mut().set` and map with `from_datafusion_error`, so both routes reject the same keys the same way. The `ScalarValue::Utf8` wrapper went with it: the parameter is already `HashMap`, and upstream only called `to_string()` back on it. Which entry a dictionary with several bad keys reports is unspecified, since `HashMap` iteration order is arbitrary. Documented rather than sorted -- a caller fixes the reported key and runs again. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 17 ++++++++-- docs/source/user-guide/configuration.md | 19 ++++++++++-- python/datafusion/context.py | 24 ++++++++++++++- python/tests/test_context.py | 41 +++++++++++++++++++++++++ 4 files changed, 94 insertions(+), 7 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index dc57d8f37..7b97e2479 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -121,17 +121,28 @@ impl From for PySessionConfig { #[pymethods] impl PySessionConfig { + /// Build a config, optionally applying options by key. + /// + /// Each entry goes through the same fallible path as [`Self::set`] rather + /// than `SessionConfig::set`, which forwards to `set_str` and unwraps: an + /// unknown namespace would abort as a `PanicException` before any of the + /// remaining entries were applied. Replaying a settings dictionary is the + /// reason this constructor takes one, and + /// `information_schema.df_settings` lists keys it cannot accept. #[pyo3(signature = (config_options=None))] #[new] - fn new(config_options: Option>) -> Self { + fn new(config_options: Option>) -> PyResult { let mut config = SessionConfig::new(); if let Some(hash_map) = config_options { for (k, v) in &hash_map { - config = config.set(k, &ScalarValue::Utf8(Some(v.clone()))); + config + .options_mut() + .set(k, v) + .map_err(from_datafusion_error)?; } } - Self { config } + Ok(Self { config }) } fn with_create_default_catalog_and_schema(&self, enabled: bool) -> Self { diff --git a/docs/source/user-guide/configuration.md b/docs/source/user-guide/configuration.md index 398a874c0..158f559c1 100644 --- a/docs/source/user-guide/configuration.md +++ b/docs/source/user-guide/configuration.md @@ -64,13 +64,26 @@ Every method above modifies the config in place and returns it, which is what ma chained style work — the object you started with is the object you end up passing to `SessionContext`. +A whole dictionary of options can be applied at once by passing it to the +{py:class}`~datafusion.SessionConfig` constructor, which is the shape a replayed set of +settings usually arrives in: + +```python +config = SessionConfig({"datafusion.execution.batch_size": "1024"}) +``` + +Both routes reject the same keys, so which one you use does not change what is accepted. The +constructor applies its entries in an unspecified order, so a dictionary with more than one +bad key does not report a predictable one first. + One trap is worth knowing about if you read settings back out of a session and replay them somewhere else, such as onto a worker process or into a test fixture. With `with_information_schema(True)`, the `information_schema.df_settings` table lists the `datafusion.runtime.*` keys alongside the rest, but those come from the runtime environment -rather than from `ConfigOptions` and cannot be set with `set`. Feeding that table's rows -back in verbatim will fail on the first such row. Configure the runtime through -`RuntimeEnvBuilder` instead, and skip the `datafusion.runtime.` prefix when replaying. +rather than from `ConfigOptions` and cannot be set this way. Feeding that table's rows back +in verbatim will fail on the first such row, whichever route you use. Configure the runtime +through `RuntimeEnvBuilder` instead, and skip the `datafusion.runtime.` prefix when +replaying. ## Maximizing CPU Usage diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 020ea06eb..94de3782b 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -167,8 +167,30 @@ class SessionConfig: def __init__(self, config_options: dict[str, str] | None = None) -> None: """Create a new :py:class:`SessionConfig` with the given configuration options. + Each entry is applied as though passed to :py:meth:`set`, so the same + keys are rejected. See :ref:`configuration`. + Args: - config_options: Configuration options. + config_options: Options to apply, keyed by fully qualified name. + + Raises: + ValueError: If a key names no known option, or a value does not + parse as that option's declared type. Which of several bad + entries is reported is not defined. + + Example usage: + + >>> from datafusion import SessionConfig + >>> ctx = SessionContext(SessionConfig()) + >>> config = SessionConfig( + ... config_options={"datafusion.execution.batch_size": "1024"} + ... ) + >>> ctx = SessionContext(config.with_information_schema(True)) + >>> ctx.sql( + ... "select value from information_schema.df_settings" + ... " where name = 'datafusion.execution.batch_size'" + ... ).collect()[0]["value"][0] + """ self.config_internal = SessionConfigInternal(config_options) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 1a15e4a54..d8bf9cc79 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -103,6 +103,47 @@ def test_create_context_with_all_valid_args(): ctx.catalog("datafusion") +def test_session_config_constructor_applies_options(): + """A dict passed to the constructor reaches the session's options.""" + config = SessionConfig( + { + "datafusion.execution.batch_size": "1024", + "datafusion.execution.target_partitions": "3", + } + ) + ctx = SessionContext(config.with_information_schema(True)) + + settings = ctx.sql( + "select name, value from information_schema.df_settings" + " where name in ('datafusion.execution.batch_size'," + " 'datafusion.execution.target_partitions')" + ).to_pydict() + + assert dict(zip(settings["name"], settings["value"], strict=True)) == { + "datafusion.execution.batch_size": "1024", + "datafusion.execution.target_partitions": "3", + } + + +def test_session_config_constructor_rejects_an_unknown_namespace(): + """A bad key in the constructor's dict raises rather than panicking. + + The same defect as `SessionConfig.set` had, reached through the argument + that a replayed `information_schema.df_settings` dictionary arrives in. + `ValueError`, not a bare `Exception`: a panic would arrive as + `PanicException`, which derives from `BaseException` and so would not be + caught here at all. + """ + with pytest.raises(ValueError, match="runtime"): + SessionConfig({"datafusion.runtime.memory_limit": "unlimited"}) + + +def test_session_config_constructor_rejects_an_unparsable_value(): + """A well-known key with a value of the wrong type raises too.""" + with pytest.raises(ValueError, match="batch_size"): + SessionConfig({"datafusion.execution.batch_size": "not_an_int"}) + + def test_register_record_batches(ctx): # create a RecordBatch and register it as memtable batch = pa.RecordBatch.from_arrays( From 079223af2b1c3c7b16be5466ed27ac7791547af4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 10:58:00 -0400 Subject: [PATCH 11/14] Say where Range partitioning actually comes from Two docstrings claimed no plan reports `Partitioning::Range` because upstream support was unimplemented. Neither half holds for DataFusion 55: `repartition/mod.rs` routes range partitioning through `RangeExpr`, and `physical_planner.rs` has a test asserting a planned `RepartitionExec` reports it. The comment's reasoning was also inverted -- the match arm compiles today because the variant exists today, not in anticipation of it landing. What is true is narrower: nothing in this package's own API asks for one. `repartition` requests round-robin, `repartition_by_hash` requests hash, and SQL has no range-repartition syntax. But a plan need not have been built here. `datafusion-proto` encodes and decodes physical range partitioning and `datafusion-ffi` carries it in both directions, so `ExecutionPlan.from_bytes` can return a plan reporting it, as can an extension library's query planner. State that where each reader is: the reachability in the guide beside the rest of the scheme discussion, one sentence and a `:ref:` in the docstring. Also note that `hash_expressions` is `None` for `Range`, which leaves the ordering and split points reachable only through `repr`. No test: reaching `Range` from Python needs a plan built in Rust, and a Rust test here would never run in CI. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/physical_plan.rs | 15 ++++++++++----- docs/source/user-guide/configuration.md | 15 +++++++++++++++ python/datafusion/plan.py | 8 +++++--- 3 files changed, 30 insertions(+), 8 deletions(-) diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index f22e59fa6..87386ee85 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -163,11 +163,13 @@ impl PyPhysicalPartitioning { /// them, which is the common case for a file scan. `RoundRobinBatch` and /// `Hash` come from a `RepartitionExec`. /// - /// `Range` is in the upstream enum but no plan reports it yet: optimizer - /// and execution support is deliberately unimplemented, per - /// . The arm is here so - /// this getter keeps compiling when that lands, not because it is - /// reachable today. + /// `Range` is implemented upstream and reaches this getter, but never from + /// a plan this package built: `DataFrame.repartition` requests round-robin, + /// `repartition_by_hash` requests hash, and SQL has no range-repartition + /// syntax. It arrives on a plan built elsewhere -- decoded by + /// `ExecutionPlan.from_bytes`, or returned by an extension library's query + /// planner -- since `datafusion-proto` and `datafusion-ffi` both carry + /// `Partitioning::Range` faithfully. #[getter] pub fn scheme(&self) -> &'static str { match self.partitioning { @@ -187,6 +189,9 @@ impl PyPhysicalPartitioning { /// /// These are physical expressions, which have no Python representation, so /// they are returned in their displayed form. + /// + /// `None` for `Range` too, whose ordering and split points this class does + /// not expose yet. #[getter] pub fn hash_expressions(&self) -> Option> { match &self.partitioning { diff --git a/docs/source/user-guide/configuration.md b/docs/source/user-guide/configuration.md index 158f559c1..32d13e0bc 100644 --- a/docs/source/user-guide/configuration.md +++ b/docs/source/user-guide/configuration.md @@ -169,6 +169,21 @@ rows are distributed across them, which is the ordinary case for a file scan. {py:attr}`~datafusion.ExecutionPlan.partition_count` gives the same count on its own when the scheme does not matter. +Four schemes exist, but only three of them can come out of a plan you built here. +`UnknownPartitioning` comes from a source, and `RoundRobinBatch` and `Hash` from a +repartition — either one you asked for or one the optimizer inserted. `Range`, which spreads +an ordered key space across partitions at chosen split points, has no request form in this +package: `repartition` asks for round-robin, `repartition_by_hash` asks for hash, and SQL +has no range-repartition syntax. + +It is still worth handling, because a plan does not have to have been built here. Both +`datafusion-proto` and `datafusion-ffi` carry range partitioning faithfully, so +{py:meth}`~datafusion.ExecutionPlan.from_bytes` can return a plan reporting it, as can an +extension library whose query planner builds one — see {ref}`extension_planners`. Such a +plan executes normally; only the split points are invisible, since +{py:attr}`~datafusion.PhysicalPartitioning.hash_expressions` returns `None` for every scheme +but `Hash`. Read {py:func}`repr` of the partitioning to see them. + ### Benchmark Example The repository includes a benchmark script that demonstrates how to maximize CPU usage diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 72fcbb932..d92df4c2a 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -351,8 +351,9 @@ def scheme( ``"UnknownPartitioning"`` means the plan knows how many partitions it has but nothing about how rows are distributed between them, which is the usual case for a file scan. ``"RoundRobinBatch"`` and ``"Hash"`` - come from a repartition the optimizer inserted. ``"Range"`` is defined - upstream but no plan reports it yet. + come from a repartition the optimizer inserted. ``"Range"`` only + appears on a plan this package did not build; see + :ref:`checking_partitioning`. Examples: >>> from datafusion import SessionContext @@ -381,7 +382,8 @@ def hash_expressions(self) -> list[str] | None: """The expressions rows are hashed on, or ``None`` for other schemes. Physical expressions have no Python representation, so these are - returned in their displayed form. + returned in their displayed form. ``None`` covers ``"Range"`` as well, + whose ordering and split points this class does not expose. Examples: >>> import pyarrow as pa From 6205a25aee736a3864adc1f7bdad16dfe6b75c63 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 11:00:45 -0400 Subject: [PATCH 12/14] Describe expr.Partitioning as what it is, not as an argument Both docstrings introducing `PhysicalPartitioning` distinguished it from `datafusion.expr.Partitioning` by calling that one "the partitioning `repartition_by_hash` asks for". No method takes one: `repartition` takes a count and `repartition_by_hash` takes expressions and a count. The type cannot be constructed from Python at all -- `Partitioning()` raises `TypeError` -- has no public members, and is only ever handed back by `Repartition.partitioning_scheme()`. Describe it that way instead: the logical partitioning a `Repartition` node records. The request-versus-result contrast the docstrings were reaching for is real and worth drawing, so keep it, but attach it to the node that holds the request rather than to a parameter that does not exist. The Rust comment also notes that the two enums differ, since the logical one has `DistributeBy` and no `UnknownPartitioning`. Pin the contrast with a test rather than only asserting it in prose. A `repartition_by_hash(num=8)` with nothing above it to consume the redistribution is dropped by the optimizer, so the request records Hash into 8 while the built plan reports `UnknownPartitioning(2)` -- disagreeing on both scheme and count, which is the reason the two types stay separate. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/physical_plan.rs | 11 ++++++++--- python/datafusion/plan.py | 8 +++++--- python/tests/test_plans.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 6 deletions(-) diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index 87386ee85..c7696ce12 100644 --- a/crates/core/src/physical_plan.rs +++ b/crates/core/src/physical_plan.rs @@ -136,9 +136,14 @@ impl PyExecutionPlan { /// How a physical plan's output rows are spread across its partitions. /// -/// Distinct from `datafusion.expr.Partitioning`, which is the *logical* -/// partitioning `DataFrame.repartition` takes as a request. This one reports -/// what a built plan actually does. +/// Distinct from `datafusion.expr.Partitioning`, the *logical* partitioning +/// recorded on a `Repartition` node and read back with +/// `Repartition.partitioning_scheme()`. Neither is an argument to anything: +/// `DataFrame.repartition` takes a count and `repartition_by_hash` takes +/// expressions and a count. The logical one records the request; this one +/// reports what the built plan does with it, and they disagree whenever the +/// optimizer rewrites or drops the repartition. The two Rust enums differ +/// too -- the logical one has `DistributeBy` and no `UnknownPartitioning`. // `skip_from_py_object` because this is a read-only report: nothing accepts a // partitioning as an argument, so there is no inbound direction to support. // Rust callers that need the `Partitioning` read it off the plan instead. diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index d92df4c2a..ace2c4775 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -333,9 +333,11 @@ class PhysicalPartitioning: """How a physical plan's output rows are spread across its partitions. Returned by :py:attr:`ExecutionPlan.output_partitioning`. This is the - partitioning a built plan *has*, which is different from - :py:class:`datafusion.expr.Partitioning` — the partitioning - :py:meth:`~datafusion.DataFrame.repartition_by_hash` *asks* for. + partitioning a built plan *has*. Distinct from + :py:class:`datafusion.expr.Partitioning`, the *logical* partitioning a + ``Repartition`` node records and hands back from + ``partitioning_scheme()`` — a request, which the plan need not honour. See + :ref:`checking_partitioning`. """ def __init__(self, partitioning: df_internal.PhysicalPartitioning) -> None: diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 07072a1f0..690182c80 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -25,11 +25,13 @@ LogicalPlan, Metric, MetricsSet, + PhysicalPartitioning, SessionConfig, SessionContext, col, udf, ) +from datafusion.expr import Partitioning # Note: CSV because a *logical* plan cannot carry a memory table. The physical @@ -146,6 +148,34 @@ def test_output_partitioning_reports_the_scheme_not_just_the_count() -> None: assert repr(partitioning) == "Hash([a@0], 4)" +def test_a_requested_partitioning_and_the_resulting_one_disagree() -> None: + """The logical request and the physical result are different things. + + `datafusion.expr.Partitioning` is what a `Repartition` node records — the + request. `PhysicalPartitioning` is what the built plan does. Here the + optimizer drops the repartition outright, because nothing above it needs + the rows redistributed, so the two do not even agree on the scheme. + """ + ctx = SessionContext(SessionConfig().with_target_partitions(4)) + ctx.register_record_batches( + "t", + [[pa.record_batch({"a": [1, 2, 3]})], [pa.record_batch({"a": [4, 5, 6]})]], + ) + df = ctx.table("t").repartition_by_hash(col("a"), num=8) + + # The request survives on the logical plan, as an opaque object of the + # other Partitioning type. + requested = df.logical_plan().to_variant().partitioning_scheme() + assert isinstance(requested, Partitioning) + assert not isinstance(requested, PhysicalPartitioning) + + # The result honours neither the scheme nor the count that was asked for. + resulting = df.execution_plan().output_partitioning + assert isinstance(resulting, PhysicalPartitioning) + assert resulting.scheme == "UnknownPartitioning" + assert resulting.partition_count == 2 + + def test_output_partitioning_reports_round_robin(tmp_path) -> None: """A round-robin repartition reports `RoundRobinBatch`. From b836628cc80497b0f1be0b762146c208ca78c0fe Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 11:06:53 -0400 Subject: [PATCH 13/14] Give PhysicalPartitioning equality, and tighten four small things Drop the guide sentence claiming every config method mutates in place and returns itself. It is true, but every `with_*` docstring still promises "a new `SessionConfig` object", and a guide that contradicts the API docs it links to is worse than a guide that stays quiet. Correcting fifteen docstrings is its own change. `PhysicalPartitioning` gains `__eq__` and `__hash__`, computed structurally from scheme, partition count and hash expressions. Deliberately not delegated to `Partitioning`'s `PartialEq`, whose match lists no `UnknownPartitioning` arm and so falls through to `false`: two identical `UnknownPartitioning(2)` values are unequal there. That is defensible for deciding whether a partitioning satisfies a distribution requirement, but a non-reflexive `__eq__` would be a trap in Python, and `UnknownPartitioning` is what an ordinary file scan reports. `__hash__` comes along so the class stays usable in a set. `test_output_partitioning_reports_round_robin` asserted set equality over every scheme in the tree, pinning optimizer output that is not the property under test. Membership instead, and 50 rows rather than 2000 -- the round robin appears either way, so the larger file bought nothing. The PyO3 parameter is now `partition` to match the wrapper, so `ctx.execute(plan, partition=0)` works on the internal binding too, and the test covers the keyword form alongside the positional one. `from_bytes` had its false memory-table sentence removed but got nothing back, so it no longer said that a decoding session need share nothing with the encoder -- the property that makes it useful. Restored from the reader's side. Also cover the `OverflowError` a negative index raises, which was documented on `execute` but never exercised. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 8 ++-- docs/source/user-guide/configuration.md | 4 -- python/datafusion/plan.py | 41 ++++++++++++++++ python/tests/test_plans.py | 64 +++++++++++++++++++++++-- 4 files changed, 105 insertions(+), 12 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7b97e2479..fe75668fb 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1436,20 +1436,20 @@ impl PySessionContext { pub fn execute( &self, plan: PyExecutionPlan, - part: usize, + partition: usize, py: Python, ) -> PyDataFusionResult { let plan = plan.plan.clone(); let partition_count = plan.output_partitioning().partition_count(); - if part >= partition_count { + if partition >= partition_count { return Err(PyValueError::new_err(format!( - "Partition index {part} is out of range for a plan with \ + "Partition index {partition} is out of range for a plan with \ {partition_count} partition(s)" )) .into()); } let ctx: TaskContext = TaskContext::from(&self.ctx.state()); - let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?; + let stream = spawn_future(py, async move { plan.execute(partition, Arc::new(ctx)) })?; Ok(PyRecordBatchStream::new(stream)) } diff --git a/docs/source/user-guide/configuration.md b/docs/source/user-guide/configuration.md index 32d13e0bc..a5e926d04 100644 --- a/docs/source/user-guide/configuration.md +++ b/docs/source/user-guide/configuration.md @@ -60,10 +60,6 @@ or an unparsable value raises rather than being silently ignored: config = SessionConfig().set("datafusion.execution.batch_size", "1024") ``` -Every method above modifies the config in place and returns it, which is what makes the -chained style work — the object you started with is the object you end up passing to -`SessionContext`. - A whole dictionary of options can be applied at once by passing it to the {py:class}`~datafusion.SessionConfig` constructor, which is the shape a replayed set of settings usually arrives in: diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index ace2c4775..8b61c0979 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -231,6 +231,9 @@ def from_bytes(ctx: SessionContext, data: bytes) -> ExecutionPlan: Unlike :py:meth:`datafusion.Expr.from_bytes`, ``ctx`` is required and positional, and there is no fallback to a worker or global context. + ``ctx`` need share nothing with the session that encoded the plan: a + scan over a table registered from record batches decodes here, because + the batches travel inside the encoded scan. See Also: :py:meth:`to_bytes`, :py:meth:`LogicalPlan.from_bytes`. @@ -416,6 +419,44 @@ def __repr__(self) -> str: """ return self._raw_partitioning.__repr__() + def _key(self) -> tuple[str, int, tuple[str, ...] | None]: + exprs = self.hash_expressions + return (self.scheme, self.partition_count, tuple(exprs) if exprs else None) + + def __eq__(self, other: object) -> bool: + """Compare two partitionings by scheme, count and hash expressions. + + Equality is structural, and does not mirror DataFusion's own + comparison of the underlying type, under which two + ``UnknownPartitioning`` values of the same width are unequal. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> a = ctx.from_pydict({"a": [1, 2, 3]}).execution_plan() + >>> b = ctx.from_pydict({"b": [4, 5, 6]}).execution_plan() + >>> a.output_partitioning == b.output_partitioning + True + >>> a.output_partitioning == "UnknownPartitioning(1)" + False + """ + if not isinstance(other, PhysicalPartitioning): + return NotImplemented + return self._key() == other._key() + + def __hash__(self) -> int: + """Hash the partitioning, consistently with :py:meth:`__eq__`. + + Examples: + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> a = ctx.from_pydict({"a": [1, 2, 3]}).execution_plan() + >>> b = ctx.from_pydict({"b": [4, 5, 6]}).execution_plan() + >>> len({a.output_partitioning, b.output_partitioning}) + 1 + """ + return hash(self._key()) + class MetricsSet: """A set of metrics for a single execution plan operator. diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 690182c80..76b11739c 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -184,7 +184,7 @@ def test_output_partitioning_reports_round_robin(tmp_path) -> None: survives at the root, so reach it by walking `children`. """ path = tmp_path / "rr.parquet" - pq.write_table(pa.table({"a": list(range(2000)), "b": [1] * 2000}), path) + pq.write_table(pa.table({"a": list(range(50)), "b": [1] * 50}), path) ctx = SessionContext(SessionConfig().with_target_partitions(8)) ctx.register_parquet("t", str(path)) @@ -197,9 +197,11 @@ def test_output_partitioning_reports_round_robin(tmp_path) -> None: schemes.add(node.output_partitioning.scheme) stack.extend(node.children()) - # The single-file scan, the round-robin above it, and the hash repartition - # for the grouping are all in one tree. - assert schemes == {"UnknownPartitioning", "RoundRobinBatch", "Hash"} + # Membership, not equality: which other nodes the optimizer puts in this + # tree is its business, and pinning the whole set here would make an + # unrelated planner change look like a failure of this accessor. The other + # schemes are asserted directly where they are the subject. + assert "RoundRobinBatch" in schemes def test_execute_rejects_an_out_of_range_partition() -> None: @@ -212,6 +214,60 @@ def test_execute_rejects_an_out_of_range_partition() -> None: with pytest.raises(ValueError, match="Partition index 5 is out of range"): ctx.execute(plan, 5) + # The keyword is `partition`, as the upgrade guide says. + with pytest.raises(ValueError, match="Partition index 5 is out of range"): + ctx.execute(plan, partition=5) + + +def test_execute_rejects_a_negative_partition() -> None: + """A negative index cannot reach the bounds check, so it overflows first. + + Documented on `execute` as `OverflowError` because that is what PyO3 + raises converting to `usize`, before any DataFusion code runs. + """ + ctx = SessionContext() + ctx.register_record_batches("t", [[pa.record_batch({"a": [1, 2, 3]})]]) + plan = ctx.sql("select a from t").execution_plan() + + with pytest.raises(OverflowError): + ctx.execute(plan, -1) + + +def test_physical_partitioning_equality_is_structural() -> None: + """Two partitionings are equal when scheme, count and keys agree. + + Not DataFusion's own comparison of the underlying type, which reports two + `UnknownPartitioning` values of the same width as unequal. A reflexive + `__eq__` is the Python expectation, and the count-only alternative would + make `Hash` on different keys compare equal. + """ + ctx = SessionContext(SessionConfig().with_target_partitions(4)) + ctx.register_record_batches( + "t", + [[pa.record_batch({"a": [1, 2, 3]})], [pa.record_batch({"a": [4, 5, 6]})]], + ) + plan = ctx.sql("select a from t").execution_plan() + other_scan = ctx.sql("select a as b from t").execution_plan().output_partitioning + grouped = ( + ctx.sql("select a, count(*) from t group by a") + .execution_plan() + .output_partitioning + ) + + # The property builds a fresh wrapper per access, so these are two objects + # over one partitioning. `UnknownPartitioning` is precisely the scheme + # DataFusion's own comparison reports as unequal to itself. + scan, scan_again = plan.output_partitioning, plan.output_partitioning + assert scan is not scan_again + assert scan == scan_again + + assert scan == other_scan + assert scan != grouped + assert scan != "UnknownPartitioning(2)" + + # Hashing agrees, so these collapse in a set the way equality implies. + assert len({scan, other_scan, grouped}) == 2 + def test_session_config_set_rejects_an_unknown_namespace() -> None: """A bad config key raises rather than aborting through a Rust panic. From 7123e7d1d79252d45544f9704d3a7875318602f0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 11:45:26 -0400 Subject: [PATCH 14/14] File the SessionConfig tests with the other SessionConfig tests The two `SessionConfig.set` tests went into `test_plans.py` because they were committed alongside the `execute` bounds check, not because they have anything to do with plans. `test_context.py` is where `SessionConfig` construction is already covered, and it now also holds the three constructor tests for the other half of the same panic defect. Move them there, ahead of the constructor cases so the method they refer back to is read first, and fold the duplicated note about `PanicException` deriving from `BaseException` into the first of the five. `test_plans.py` keeps its `SessionConfig` import for the `with_target_partitions` calls in the partitioning tests. No assertion changed. Co-Authored-By: Claude Opus 5 (1M context) --- python/tests/test_context.py | 23 ++++++++++++++++++++--- python/tests/test_plans.py | 20 -------------------- 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index d8bf9cc79..9fbc4744f 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -103,6 +103,26 @@ def test_create_context_with_all_valid_args(): ctx.catalog("datafusion") +def test_session_config_set_rejects_an_unknown_namespace(): + """A bad config key raises rather than aborting through a Rust panic. + + `datafusion.runtime.*` appears in `information_schema.df_settings` but has + no `ConfigOptions` namespace, so it is the key a naive "read the settings + back and replay them on the worker" loop hits first. + """ + # `ValueError`, not a bare `Exception`: a panic would arrive as + # `PanicException`, which derives from `BaseException` and so would not be + # caught here at all. Both this and the constructor cases below rely on it. + with pytest.raises(ValueError, match="runtime"): + SessionConfig().set("datafusion.runtime.memory_limit", "unlimited") + + +def test_session_config_set_rejects_an_unparsable_value(): + """A well-known key with a value of the wrong type raises too.""" + with pytest.raises(ValueError, match="batch_size"): + SessionConfig().set("datafusion.execution.batch_size", "not_an_int") + + def test_session_config_constructor_applies_options(): """A dict passed to the constructor reaches the session's options.""" config = SessionConfig( @@ -130,9 +150,6 @@ def test_session_config_constructor_rejects_an_unknown_namespace(): The same defect as `SessionConfig.set` had, reached through the argument that a replayed `information_schema.df_settings` dictionary arrives in. - `ValueError`, not a bare `Exception`: a panic would arrive as - `PanicException`, which derives from `BaseException` and so would not be - caught here at all. """ with pytest.raises(ValueError, match="runtime"): SessionConfig({"datafusion.runtime.memory_limit": "unlimited"}) diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 76b11739c..e0c6e2c0c 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -269,26 +269,6 @@ def test_physical_partitioning_equality_is_structural() -> None: assert len({scan, other_scan, grouped}) == 2 -def test_session_config_set_rejects_an_unknown_namespace() -> None: - """A bad config key raises rather than aborting through a Rust panic. - - `datafusion.runtime.*` appears in `information_schema.df_settings` but has - no `ConfigOptions` namespace, so it is the key a naive "read the settings - back and replay them on the worker" loop hits first. - """ - # `ValueError`, not a bare `Exception`: a panic would arrive as - # `PanicException`, which derives from `BaseException` and so would not be - # caught here at all. - with pytest.raises(ValueError, match="runtime"): - SessionConfig().set("datafusion.runtime.memory_limit", "unlimited") - - -def test_session_config_set_rejects_an_unparsable_value() -> None: - """A well-known key with a value of the wrong type raises too.""" - with pytest.raises(ValueError, match="batch_size"): - SessionConfig().set("datafusion.execution.batch_size", "not_an_int") - - def test_installing_a_physical_codec_preserves_strict_mode() -> None: """Installing a physical extension codec must not re-enable inlining.