diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index c711a62dc..fe75668fb 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, }; @@ -120,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 { @@ -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 { + 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) -> PyResult { @@ -1407,12 +1436,20 @@ impl PySessionContext { pub fn execute( &self, plan: PyExecutionPlan, - part: usize, + partition: usize, py: Python, ) -> PyDataFusionResult { - 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)) } 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..c7696ce12 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,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> { + 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 Arc { diff --git a/docs/source/user-guide/configuration.md b/docs/source/user-guide/configuration.md index d1c5c9b44..a5e926d04 100644 --- a/docs/source/user-guide/configuration.md +++ b/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 diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index f98590b0a..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,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 diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index f9e96382e..810e090de 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -122,9 +122,31 @@ 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. + // + // 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/__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/context.py b/python/datafusion/context.py index bbf08e84e..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) @@ -341,14 +363,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: + ValueError: 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 @@ -2351,9 +2396,41 @@ 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``. + OverflowError: If ``partition`` is negative, or too large to fit a + platform-sized unsigned integer. + + 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 + """ + return RecordBatchStream(self.ctx.execute(plan._raw_plan, partition)) @staticmethod def _convert_file_sort_order( diff --git a/python/datafusion/plan.py b/python/datafusion/plan.py index 8d03bae2c..8b61c0979 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 @@ -34,6 +34,7 @@ "LogicalPlan", "Metric", "MetricsSet", + "PhysicalPartitioning", ] @@ -178,20 +179,61 @@ 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. A plan does not necessarily partition the way it was + asked to; see :ref:`checking_partitioning`. + + 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. + ``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`. @@ -204,8 +246,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 +332,132 @@ 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*. 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: + """This constructor should not be called by the end user.""" + self._raw_partitioning = partitioning + + @property + def scheme( + self, + ) -> Literal["RoundRobinBatch", "Hash", "Range", "UnknownPartitioning"]: + """Which partitioning scheme this is. + + ``"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"`` only + appears on a plan this package did not build; see + :ref:`checking_partitioning`. + + 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. ``None`` covers ``"Range"`` as well, + whose ordering and split points this class does not expose. + + 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__() + + 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_context.py b/python/tests/test_context.py index 1a15e4a54..9fbc4744f 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -103,6 +103,64 @@ 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( + { + "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. + """ + 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( diff --git a/python/tests/test_plans.py b/python/tests/test_plans.py index 0145d123e..e0c6e2c0c 100644 --- a/python/tests/test_plans.py +++ b/python/tests/test_plans.py @@ -18,20 +18,24 @@ import datetime import pyarrow as pa +import pyarrow.parquet as pq import pytest from datafusion import ( ExecutionPlan, LogicalPlan, Metric, MetricsSet, + PhysicalPartitioning, + SessionConfig, SessionContext, col, udf, ) +from datafusion.expr import Partitioning -# 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 +99,176 @@ 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() + scanned = scan.output_partitioning + assert scanned.scheme == "UnknownPartitioning" + assert scanned.hash_expressions is None + # Agrees with the count-only accessor it supplements. + assert scanned.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_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`. + + 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(50)), "b": [1] * 50}), 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()) + + # 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: + """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() + assert plan.partition_count == 1 + + 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_installing_a_physical_codec_preserves_strict_mode() -> None: """Installing a physical extension codec must not re-enable inlining.