This is an automated email from the ASF dual-hosted git repository. timsaucer pushed a commit to branch feat/plan-partitioning-and-errors in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit a26b90a08583bbe2100c28ee0d7f46a992ee7d09 Author: Tim Saucer <[email protected]> AuthorDate: Wed Sep 9 11:44:08 2026 -0400 Report physical partitioning, and stop two panics escaping as panics 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) <[email protected]> --- 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 c711a62d..8aa0504f 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<Self> { + let mut config = self.config.clone(); + config.options_mut().set(key, value)?; + Ok(Self::from(config)) } pub fn with_extension(&self, extension: Bound<PyAny>) -> PyResult<Self> { @@ -1412,6 +1422,18 @@ impl PySessionContext { ) -> PyDataFusionResult<PyRecordBatchStream> { 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 7f0f9cb3..492d5764 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::<metrics::PyMetricsSet>()?; m.add_class::<metrics::PyMetric>()?; m.add_class::<physical_plan::PyExecutionPlan>()?; + m.add_class::<physical_plan::PyPhysicalPartitioning>()?; m.add_class::<record_batch::PyRecordBatch>()?; m.add_class::<record_batch::PyRecordBatchStream>()?; diff --git a/crates/core/src/physical_plan.rs b/crates/core/src/physical_plan.rs index 594655a6..ddb344a9 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<Vec<String>> { + match &self.partitioning { + Partitioning::Hash(exprs, _) => { + Some(exprs.iter().map(|expr| format!("{expr}")).collect()) + } + _ => None, + } + } + + fn __repr__(&self) -> String { + format!("{}", self.partitioning) + } +} + +impl From<Partitioning> for PyPhysicalPartitioning { + fn from(partitioning: Partitioning) -> Self { + Self { partitioning } + } +} + +impl From<PyPhysicalPartitioning> for Partitioning { + fn from(partitioning: PyPhysicalPartitioning) -> Self { + partitioning.partitioning + } } impl From<PyExecutionPlan> for Arc<dyn ExecutionPlan> { diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index f98590b0..d71f9370 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 f9e96382..e8fd1969 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<u8>, 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::<DataSourceExec>() || node.is::<ForeignExecutionPlan>() { self.counters .encode_execution_plan diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 3696d92a..1b44f8a7 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -100,7 +100,13 @@ from .extensions import ( ) 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 @@ __all__ = [ "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", + "PhysicalPartitioning", "QueryPlannerExportable", "RecordBatch", "RecordBatchStream", diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 8d03bae2..49b2035d 100644 --- a/python/datafusion/plan.py +++ b/python/datafusion/plan.py @@ -34,6 +34,7 @@ __all__ = [ "LogicalPlan", "Metric", "MetricsSet", + "PhysicalPartitioning", ] @@ -178,17 +179,54 @@ class ExecutionPlan: @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 @@ class ExecutionPlan: 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 @@ class ExecutionPlan: 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 0145d123..95b5a77f 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -24,14 +24,15 @@ from datafusion import ( 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. --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
