Report physical partitioning and resolve two panics escaping as PanicException

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,
 };
@@ -120,17 +121,28 @@ impl From<SessionConfig> 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<HashMap<String, String>>) -> Self {
+    fn new(config_options: Option<HashMap<String, String>>) -> PyResult<Self> {
         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 {
@@ -193,8 +205,25 @@ 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.
+    ///
+    /// 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<Self> {
+        let mut config = self.config.clone();
+        config
+            .options_mut()
+            .set(key, value)
+            .map_err(from_datafusion_error)?;
+        Ok(Self::from(config))
     }
 
     pub fn with_extension(&self, extension: Bound<PyAny>) -> PyResult<Self> {
@@ -1407,12 +1436,20 @@ impl PySessionContext {
     pub fn execute(
         &self,
         plan: PyExecutionPlan,
-        part: usize,
+        partition: usize,
         py: Python,
     ) -> PyDataFusionResult<PyRecordBatchStream> {
-        let ctx: TaskContext = TaskContext::from(&self.ctx.state());
         let plan = plan.plan.clone();
-        let stream = spawn_future(py, async move { plan.execute(part, Arc::new(ctx)) })?;
+        let partition_count = plan.output_partitioning().partition_count();
+        if partition >= partition_count {
+            return Err(PyValueError::new_err(format!(
+                "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(partition, Arc::new(ctx)) })?;
         Ok(PyRecordBatchStream::new(stream))
     }
 

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>()?;
 

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,95 @@ 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`, 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.
+#[pyclass(
+    skip_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.
+    ///
+    /// `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. `RoundRobinBatch` and
+    /// `Hash` come from a `RepartitionExec`.
+    ///
+    /// `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 {
+            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.
+    ///
+    /// `None` for `Range` too, whose ordering and split points this class does
+    /// not expose yet.
+    #[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<PyExecutionPlan> for Arc<dyn ExecutionPlan> {

docs/source/user-guide/configuration.md

@@ -49,6 +49,38 @@ 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")
+```
+
+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 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
 
 DataFusion uses partitions to parallelize work. For small queries the
@@ -96,6 +128,58 @@ 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.
+
+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

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,6 +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.
 
+### `SessionContext.execute` renamed its second parameter
+
+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
+ctx.execute(plan, partitions=0)  # before
+ctx.execute(plan, partition=0)  # after
+```
+
 ### Changes to the `datafusion-python-util` crate
 
 Extension libraries written in Rust usually depend on the

examples/datafusion-ffi-example/src/physical_extension_codec.rs

@@ -122,9 +122,31 @@ 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.
+        //
+        // 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::<DataSourceExec>() || node.is::<ForeignExecutionPlan>() {
             self.counters
                 .encode_execution_plan

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",