You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/27 13:22:37 UTC
(datafusion) branch main updated: fix(proto): check plan integer conversions across usize boundaries (#24483)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 2a7a1e3060 fix(proto): check plan integer conversions across usize boundaries (#24483)
2a7a1e3060 is described below
commit 2a7a1e3060c42e5ade6ad88b9be667e1b72639bb
Author: Burak Şen <bu...@gmail.com>
AuthorDate: Thu Aug 27 12:53:53 2026 +0000
fix(proto): check plan integer conversions across usize boundaries (#24483)
## Which issue does this PR close?
- Part of #24170.
## Rationale for this change
Several logical and physical plan protobuf paths convert wire integers
to `usize` with unchecked casts. On 32-bit targets, values above
`u32::MAX` silently truncate. For row limits, a fetch of `1 << 32`
becomes `0`, causing a plan that should return rows to return an empty
result.
The reverse conversion can also truncate when a 64-bit `usize` does not
fit the protobuf field.
## What changes are included in this PR?
- Add shared helpers for checked conversions between wire integers and
`usize`, with errors identifying the plan node and field.
- Apply checked conversions to limits, partition counts, batch sizes,
capacities, and indexes across logical and physical plan serialization.
- Preserve the documented `-1` sentinel for limits without a fetch.
- Reject invalid zero batch sizes while preserving proto3 defaults for
older `FilterExec` payloads.
- Treat out-of-range `usize` statistics as unknown rather than
truncating them.
The protobuf wire format is unchanged.
Follow-up work will make common protobuf option and constraint
conversions fallible.
## Are these changes tested?
Yes. Added unit and protobuf roundtrip coverage for checked conversions,
architecture-dependent integer boundaries, oversized encode/decode
values, limits without a fetch, malformed zero-column `Values` nodes,
and zero batch sizes.
## Are there any user-facing changes?
Out-of-range plan integers now return a clear planning error instead of
silently truncating. Invalid zero batch sizes are rejected, and
out-of-range statistics become unknown. There are no breaking public API
or protobuf wire-format changes.
---
datafusion/common/src/utils/mod.rs | 54 +++++++++++++++-
.../datasource/src/file_scan_config/proto.rs | 25 +++++++-
datafusion/datasource/src/memory.rs | 15 ++++-
datafusion/expr/src/logical_plan/plan.rs | 43 ++++++++++++-
datafusion/functions-table/src/generate_series.rs | 24 ++++++-
datafusion/physical-expr/src/partitioning.rs | 73 ++++-----------------
datafusion/physical-expr/src/scalar_subquery.rs | 12 ++--
datafusion/physical-plan/src/aggregates/mod.rs | 8 +--
datafusion/physical-plan/src/buffer.rs | 4 +-
datafusion/physical-plan/src/coalesce_batches.rs | 27 ++++++--
.../physical-plan/src/coalesce_partitions.rs | 12 +++-
datafusion/physical-plan/src/empty.rs | 4 +-
datafusion/physical-plan/src/filter.rs | 48 ++++++++++++--
.../physical-plan/src/joins/hash_join/exec.rs | 17 +----
datafusion/physical-plan/src/limit.rs | 24 +++----
datafusion/physical-plan/src/placeholder_row.rs | 4 +-
datafusion/physical-plan/src/sorts/sort.rs | 36 ++++++-----
.../src/sorts/sort_preserving_merge.rs | 37 +++++++----
datafusion/physical-plan/src/unnest.rs | 5 +-
.../physical-plan/src/windows/window_agg_exec.rs | 8 ++-
datafusion/proto-common/src/from_proto/mod.rs | 6 +-
datafusion/proto/src/logical_plan/mod.rs | 39 ++++++++----
datafusion/proto/src/physical_plan/mod.rs | 20 +++++-
datafusion/proto/tests/cases/plans/leaves.rs | 29 ++++++++-
datafusion/proto/tests/cases/plans/limits.rs | 74 +++++++++++++++++++++-
datafusion/proto/tests/cases/plans/misc.rs | 9 +--
datafusion/proto/tests/cases/plans/sources.rs | 35 ++++++++++
.../proto/tests/cases/roundtrip_logical_plan.rs | 44 +++++++++++++
28 files changed, 556 insertions(+), 180 deletions(-)
diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs
index 142d3cd8e2..e047db39a5 100644
--- a/datafusion/common/src/utils/mod.rs
+++ b/datafusion/common/src/utils/mod.rs
@@ -25,7 +25,9 @@ pub mod proxy;
pub mod string_utils;
use crate::assert_or_internal_err;
-use crate::error::{_exec_datafusion_err, _exec_err, _internal_datafusion_err};
+use crate::error::{
+ _exec_datafusion_err, _exec_err, _internal_datafusion_err, _plan_datafusion_err,
+};
use crate::{Result, ScalarValue};
use arrow::array::{
Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,
@@ -1142,6 +1144,32 @@ pub fn combine_limit(
(combined_skip, combined_fetch)
}
+/// Converts a wire integer to `usize`, rejecting out-of-range values.
+/// `context` and `field` identify the value in the error message.
+pub fn usize_from_wire<T>(value: T, context: &str, field: &str) -> Result<usize>
+where
+ T: TryInto<usize> + std::fmt::Display + Copy,
+{
+ value.try_into().map_err(|_| {
+ _plan_datafusion_err!(
+ "{context}: {field} wire value {value} is out of range for usize"
+ )
+ })
+}
+
+/// Converts a `usize` to a wire integer, rejecting out-of-range values.
+pub fn usize_to_wire<T: TryFrom<usize>>(
+ value: usize,
+ context: &str,
+ field: &str,
+) -> Result<T> {
+ T::try_from(value).map_err(|_| {
+ _plan_datafusion_err!(
+ "{context}: {field} value {value} is out of range for the plan wire format"
+ )
+ })
+}
+
/// Returns the estimated number of threads available for parallel execution.
///
/// This is a wrapper around `std::thread::available_parallelism`, providing a default value
@@ -1482,6 +1510,30 @@ mod tests {
#[cfg(feature = "sql")]
use sqlparser::ast::Ident;
+ #[test]
+ fn test_usize_wire_conversions() {
+ assert_eq!(usize_from_wire(42_u64, "SomeExec", "fetch").unwrap(), 42);
+ let err = usize_from_wire(-1_i64, "SomeExec", "skip").unwrap_err();
+ assert_eq!(
+ err.strip_backtrace(),
+ "Error during planning: SomeExec: skip wire value -1 is out of range for usize"
+ );
+
+ assert_eq!(usize_to_wire::<u32>(42, "SomeExec", "fetch").unwrap(), 42);
+ let err = usize_to_wire::<u8>(256, "SomeExec", "fetch").unwrap_err();
+ assert_eq!(
+ err.strip_backtrace(),
+ "Error during planning: SomeExec: fetch value 256 is out of range for the plan wire format"
+ );
+
+ let max = usize_from_wire(u64::MAX, "SomeExec", "fetch");
+ if usize::BITS >= u64::BITS {
+ assert_eq!(max.unwrap(), usize::MAX);
+ } else {
+ assert!(max.is_err());
+ }
+ }
+
#[test]
fn test_bisect_linear_left_and_right() -> Result<()> {
let arrays: Vec<ArrayRef> = vec![
diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs
index efd17e4f57..4071ea471b 100644
--- a/datafusion/datasource/src/file_scan_config/proto.rs
+++ b/datafusion/datasource/src/file_scan_config/proto.rs
@@ -39,6 +39,7 @@ use std::sync::Arc;
use arrow::datatypes::Schema;
use datafusion_common::parsers::CompressionTypeVariant;
+use datafusion_common::utils::{usize_from_wire, usize_to_wire};
use datafusion_common::{DataFusionError, Result, internal_datafusion_err};
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs};
@@ -128,7 +129,11 @@ impl FileScanConfig {
Ok(protobuf::FileScanExecConf {
file_groups,
statistics: Some((&self.statistics()).into()),
- limit: self.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }),
+ limit: self
+ .limit
+ .map(|limit| usize_to_wire::<u32>(limit, "FileScanConfig", "limit"))
+ .transpose()?
+ .map(|limit| protobuf::ScanLimit { limit }),
projection: vec![],
schema: Some((&schema).try_into()?),
table_partition_cols: self
@@ -245,14 +250,28 @@ impl FileScanConfig {
file_source
};
+ let limit = conf
+ .limit
+ .as_ref()
+ .map(|limit| usize_from_wire(limit.limit, "FileScanConfig", "limit"))
+ .transpose()?;
+ let batch_size = conf
+ .batch_size
+ .map(|size| usize_from_wire(size, "FileScanConfig", "batch_size"))
+ .transpose()?;
+ if batch_size == Some(0) {
+ return datafusion_common::plan_err!(
+ "FileScanConfig: batch_size must be greater than 0"
+ );
+ }
let config_builder = FileScanConfigBuilder::new(object_store_url, file_source)
.with_file_groups(file_groups)
.with_constraints(constraints)
.with_statistics(statistics)
- .with_limit(conf.limit.as_ref().map(|sl| sl.limit as usize))
+ .with_limit(limit)
.with_output_ordering(output_ordering)
.with_output_partitioning(output_partitioning)
- .with_batch_size(conf.batch_size.map(|s| s as usize))
+ .with_batch_size(batch_size)
.with_file_compression_type(file_compression_type);
Ok(config_builder.build())
}
diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs
index a4f51c28f8..7c10dba981 100644
--- a/datafusion/datasource/src/memory.rs
+++ b/datafusion/datasource/src/memory.rs
@@ -280,6 +280,7 @@ impl DataSource for MemorySourceConfig {
&self,
ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto;
use datafusion_proto_models::protobuf;
@@ -315,7 +316,12 @@ impl DataSource for MemorySourceConfig {
projection,
sort_information,
show_sizes: self.show_sizes,
- fetch: self.fetch.map(|f| f as u32),
+ fetch: self
+ .fetch
+ .map(|fetch| {
+ usize_to_wire(fetch, "MemoryScanExecNode", "fetch")
+ })
+ .transpose()?,
},
),
),
@@ -682,6 +688,7 @@ impl MemorySourceConfig {
ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
use datafusion_common::internal_datafusion_err;
+ use datafusion_common::utils::usize_from_wire;
use datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto;
use datafusion_proto_models::protobuf;
@@ -718,8 +725,12 @@ impl MemorySourceConfig {
sort_information.extend(LexOrdering::new(sort_exprs));
}
+ let fetch = scan
+ .fetch
+ .map(|fetch| usize_from_wire(fetch, "MemoryScanExecNode", "fetch"))
+ .transpose()?;
let source = Self::try_new(&partitions, schema, projection)?
- .with_limit(scan.fetch.map(|f| f as usize))
+ .with_limit(fetch)
.with_show_sizes(scan.show_sizes)
.try_with_sort_information(sort_information)?;
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index ab3c26795e..1a8cd81aa7 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -65,7 +65,7 @@ use datafusion_common::{
FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result,
ScalarValue, Spans, SplitPoint, TableReference, UnnestOptions,
aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err,
- internal_err, plan_err, validate_range_split_points,
+ internal_err, plan_datafusion_err, plan_err, validate_range_split_points,
};
use indexmap::IndexSet;
use itertools::Itertools as _;
@@ -3782,7 +3782,12 @@ impl Limit {
// `skip = NULL` is equivalent to `skip = 0`
let s = s.unwrap_or(0);
if s >= 0 {
- Ok(SkipType::Literal(s as usize))
+ let s = usize::try_from(s).map_err(|_| {
+ plan_datafusion_err!(
+ "OFFSET value {s} cannot be represented as usize"
+ )
+ })?;
+ Ok(SkipType::Literal(s))
} else {
plan_err!("OFFSET must be >=0, '{}' was provided", s)
}
@@ -3800,7 +3805,12 @@ impl Limit {
Some(expr) => match *expr {
Expr::Literal(ScalarValue::Int64(Some(s)), _) => {
if s >= 0 {
- Ok(FetchType::Literal(Some(s as usize)))
+ let s = usize::try_from(s).map_err(|_| {
+ plan_datafusion_err!(
+ "LIMIT value {s} cannot be represented as usize"
+ )
+ })?;
+ Ok(FetchType::Literal(Some(s)))
} else {
plan_err!("LIMIT must be >= 0, '{}' was provided", s)
}
@@ -5031,6 +5041,33 @@ mod tests {
);
}
+ #[test]
+ fn limit_literals_use_checked_usize_conversion() -> Result<()> {
+ let value = i64::from(u32::MAX) + 1;
+ let input = Arc::new(LogicalPlanBuilder::empty(false).build()?);
+ let limit = Limit {
+ skip: Some(Box::new(lit(value))),
+ fetch: Some(Box::new(lit(value))),
+ input,
+ };
+
+ if usize::BITS < 64 {
+ assert!(limit.get_skip_type().is_err());
+ assert!(limit.get_fetch_type().is_err());
+ } else {
+ let expected = usize::try_from(value).unwrap();
+ let SkipType::Literal(skip) = limit.get_skip_type()? else {
+ panic!("expected literal skip")
+ };
+ let FetchType::Literal(Some(fetch)) = limit.get_fetch_type()? else {
+ panic!("expected literal fetch")
+ };
+ assert_eq!(skip, expected);
+ assert_eq!(fetch, expected);
+ }
+ Ok(())
+ }
+
fn employee_schema() -> Schema {
Schema::new(vec![
Field::new("id", DataType::Int32, false),
diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs
index 8d9327c5ac..4b4fc01af9 100644
--- a/datafusion/functions-table/src/generate_series.rs
+++ b/datafusion/functions-table/src/generate_series.rs
@@ -291,6 +291,9 @@ impl GenerateSeriesTable {
&self,
batch_size: usize,
) -> Result<Arc<RwLock<dyn LazyBatchGenerator>>> {
+ if batch_size == 0 {
+ return plan_err!("GenerateSeriesTable: batch_size must be greater than 0");
+ }
let generator: Arc<RwLock<dyn LazyBatchGenerator>> = match &self.args {
GenSeriesArgs::ContainsNull { name } => Arc::new(RwLock::new(Empty { name })),
GenSeriesArgs::Int64Args {
@@ -877,7 +880,26 @@ mod generate_series_tests {
use datafusion_common::Result;
use datafusion_physical_plan::memory::LazyBatchGenerator;
- use crate::generate_series::GenericSeriesState;
+ use crate::generate_series::{
+ GenSeriesArgs, GenerateSeriesTable, GenericSeriesState,
+ };
+
+ #[test]
+ fn generate_series_rejects_zero_batch_size() {
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
+ let table = GenerateSeriesTable::new(
+ schema,
+ GenSeriesArgs::Int64Args {
+ start: 1,
+ end: 2,
+ step: 1,
+ include_end: true,
+ name: "generate_series",
+ },
+ );
+
+ assert!(table.as_generator(0).is_err());
+ }
#[test]
fn test_generic_series_state_reset() -> Result<()> {
diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs
index 98f082f725..70ba92a0a9 100644
--- a/datafusion/physical-expr/src/partitioning.rs
+++ b/datafusion/physical-expr/src/partitioning.rs
@@ -536,19 +536,20 @@ impl Partitioning {
&self,
ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
) -> Result<datafusion_proto_models::protobuf::Partitioning> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
+ let partition_count =
+ |count: usize| usize_to_wire::<u64>(count, "Partitioning", "partition_count");
let partition_method = match self {
Partitioning::RoundRobinBatch(n) => {
- protobuf::partitioning::PartitionMethod::RoundRobin(wire_partition_count(
- *n,
- )?)
+ protobuf::partitioning::PartitionMethod::RoundRobin(partition_count(*n)?)
}
Partitioning::Hash(exprs, n) => {
protobuf::partitioning::PartitionMethod::Hash(
protobuf::PhysicalHashRepartition {
hash_expr: ctx.encode_children_expressions(exprs)?,
- partition_count: wire_partition_count(*n)?,
+ partition_count: partition_count(*n)?,
},
)
}
@@ -574,9 +575,7 @@ impl Partitioning {
)
}
Partitioning::UnknownPartitioning(n) => {
- protobuf::partitioning::PartitionMethod::Unknown(wire_partition_count(
- *n,
- )?)
+ protobuf::partitioning::PartitionMethod::Unknown(partition_count(*n)?)
}
};
Ok(protobuf::Partitioning {
@@ -593,9 +592,12 @@ impl Partitioning {
node: &datafusion_proto_models::protobuf::Partitioning,
ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
) -> Result<Option<Self>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_common::{ScalarValue, internal_datafusion_err, internal_err};
use datafusion_proto_models::protobuf;
+ let partition_count =
+ |count: u64| usize_from_wire(count, "Partitioning", "partition_count");
let Some(partition_method) = node.partition_method.as_ref() else {
return Ok(None);
};
@@ -646,29 +648,6 @@ impl Partitioning {
}
}
-/// Narrow a wire partition count to `usize`.
-#[cfg(feature = "proto")]
-fn partition_count(count: u64) -> Result<usize> {
- usize::try_from(count).map_err(|_| {
- datafusion_common::internal_datafusion_err!(
- "Partition count {count} exceeds usize::MAX"
- )
- })
-}
-
-/// Widen a partition count to its `u64` wire representation.
-///
-/// The mirror of [`partition_count`]: an out-of-range count is an error on both
-/// sides rather than a silent truncation on the way out.
-#[cfg(feature = "proto")]
-fn wire_partition_count(count: usize) -> Result<u64> {
- u64::try_from(count).map_err(|_| {
- datafusion_common::internal_datafusion_err!(
- "Partition count {count} exceeds u64::MAX"
- )
- })
-}
-
impl PartialEq for Partitioning {
fn eq(&self, other: &Partitioning) -> bool {
match (self, other) {
@@ -1416,7 +1395,7 @@ mod partition_count_proto_tests {
use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
use datafusion_proto_models::protobuf;
- use super::{Partitioning, partition_count, wire_partition_count};
+ use super::Partitioning;
use crate::expressions::Column;
use crate::proto_test_util::{StubDecoder, StubEncoder, column_node};
@@ -1444,36 +1423,6 @@ mod partition_count_proto_tests {
]
}
- #[test]
- fn partition_count_round_trips_at_the_usize_ceiling() {
- // `usize::MAX` is the largest count that can exist in memory, so it has
- // to widen onto the wire and narrow back unchanged.
- let wire = wire_partition_count(usize::MAX).unwrap();
- assert_eq!(wire, u64::try_from(usize::MAX).unwrap());
- assert_eq!(partition_count(wire).unwrap(), usize::MAX);
- }
-
- #[test]
- fn out_of_range_partition_count_is_reported_not_wrapped() {
- // A count wider than the target's `usize` can only be reached by
- // decoding on a narrower host than the one that encoded. That used to
- // wrap (`as usize`) or panic (`unwrap`); it is an error now. On a
- // 64-bit target every `u64` fits, so the same input has to decode
- // losslessly instead of being rejected.
- let narrowed = partition_count(u64::MAX);
-
- #[cfg(target_pointer_width = "64")]
- assert_eq!(narrowed.unwrap(), usize::MAX);
-
- #[cfg(not(target_pointer_width = "64"))]
- assert!(
- narrowed
- .unwrap_err()
- .to_string()
- .contains("Partition count 18446744073709551615 exceeds usize::MAX")
- );
- }
-
#[test]
fn try_from_proto_narrows_every_counted_variant() {
let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
@@ -1492,7 +1441,7 @@ mod partition_count_proto_tests {
decoded
.unwrap_err()
.to_string()
- .contains("exceeds usize::MAX")
+ .contains("is out of range for usize")
);
}
}
diff --git a/datafusion/physical-expr/src/scalar_subquery.rs b/datafusion/physical-expr/src/scalar_subquery.rs
index 473b52a5cb..ac473ea9dd 100644
--- a/datafusion/physical-expr/src/scalar_subquery.rs
+++ b/datafusion/physical-expr/src/scalar_subquery.rs
@@ -157,6 +157,7 @@ impl PhysicalExpr for ScalarSubqueryExpr {
&self,
_ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
Ok(Some(protobuf::PhysicalExprNode {
expr_id: None,
@@ -164,12 +165,11 @@ impl PhysicalExpr for ScalarSubqueryExpr {
protobuf::PhysicalScalarSubqueryExprNode {
data_type: Some((&self.data_type).try_into()?),
nullable: self.nullable,
- index: u32::try_from(self.index.as_usize()).map_err(|_| {
- internal_datafusion_err!(
- "scalar subquery index {} does not fit in u32",
- self.index.as_usize()
- )
- })?,
+ index: usize_to_wire(
+ self.index.as_usize(),
+ "ScalarSubqueryExpr",
+ "index",
+ )?,
},
)),
}))
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 2dc41adfd6..69a24ce17e 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -2527,6 +2527,7 @@ impl AggregateExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_physical_expr::aggregate::AggregateExprBuilder;
use datafusion_proto_models::protobuf;
use protobuf::physical_aggregate_expr_node::AggregateFunction;
@@ -2701,11 +2702,10 @@ impl AggregateExec {
)
}?;
let aggregate = if let Some(limit) = limit {
+ let fetch = usize_from_wire(limit.limit, "AggregateExec", "limit")?;
let options = match limit.descending {
- Some(descending) => {
- LimitOptions::new_with_order(limit.limit as usize, descending)
- }
- None => LimitOptions::new(limit.limit as usize),
+ Some(descending) => LimitOptions::new_with_order(fetch, descending),
+ None => LimitOptions::new(fetch),
};
aggregate.with_limit_options(Some(options))
} else {
diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs
index 72df0f5345..e1452b75bc 100644
--- a/datafusion/physical-plan/src/buffer.rs
+++ b/datafusion/physical-plan/src/buffer.rs
@@ -368,6 +368,7 @@ impl BufferExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
let buffer = crate::expect_plan_variant!(
node,
@@ -378,7 +379,8 @@ impl BufferExec {
// compile error here rather than a silently dropped field.
let protobuf::BufferExecNode { input, capacity } = &**buffer;
let input = ctx.decode_required_child(input.as_deref(), "BufferExec", "input")?;
- Ok(Arc::new(BufferExec::new(input, *capacity as usize)))
+ let capacity = usize_from_wire(*capacity, "BufferExec", "capacity")?;
+ Ok(Arc::new(BufferExec::new(input, capacity)))
}
}
diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs
index 87957ced7b..0b3359bc7c 100644
--- a/datafusion/physical-plan/src/coalesce_batches.rs
+++ b/datafusion/physical-plan/src/coalesce_batches.rs
@@ -322,6 +322,7 @@ impl ExecutionPlan for CoalesceBatchesExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `CoalesceBatchesExec` is a compile error here until it is either
@@ -336,13 +337,22 @@ impl ExecutionPlan for CoalesceBatchesExec {
cache: _,
} = self;
let input = ctx.encode_child(input)?;
+ let to_wire =
+ |value, field| usize_to_wire::<u32>(value, "CoalesceBatchesExec", field);
+ if *target_batch_size == 0 {
+ return datafusion_common::plan_err!(
+ "CoalesceBatchesExec: target_batch_size must be greater than 0"
+ );
+ }
+ let target_batch_size = to_wire(*target_batch_size, "target_batch_size")?;
+ let fetch = fetch.map(|fetch| to_wire(fetch, "fetch")).transpose()?;
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(
protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches(
Box::new(protobuf::CoalesceBatchesExecNode {
input: Some(Box::new(input)),
- target_batch_size: *target_batch_size as u32,
- fetch: fetch.map(|n| n as u32),
+ target_batch_size,
+ fetch,
}),
),
),
@@ -367,6 +377,7 @@ impl CoalesceBatchesExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
let coalesce_batches = crate::expect_plan_variant!(
node,
@@ -383,9 +394,17 @@ impl CoalesceBatchesExec {
} = &**coalesce_batches;
let input =
ctx.decode_required_child(input.as_deref(), "CoalesceBatchesExec", "input")?;
+ let from_wire =
+ |value, field| usize_from_wire(value, "CoalesceBatchesExec", field);
+ let target_batch_size = from_wire(*target_batch_size, "target_batch_size")?;
+ if target_batch_size == 0 {
+ return datafusion_common::plan_err!(
+ "CoalesceBatchesExec: target_batch_size must be greater than 0"
+ );
+ }
+ let fetch = fetch.map(|fetch| from_wire(fetch, "fetch")).transpose()?;
Ok(Arc::new(
- CoalesceBatchesExec::new(input, *target_batch_size as usize)
- .with_fetch(fetch.map(|f| f as usize)),
+ CoalesceBatchesExec::new(input, target_batch_size).with_fetch(fetch),
))
}
}
diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs
index 7c8bca772e..7438f1c0ca 100644
--- a/datafusion/physical-plan/src/coalesce_partitions.rs
+++ b/datafusion/physical-plan/src/coalesce_partitions.rs
@@ -383,6 +383,7 @@ impl ExecutionPlan for CoalescePartitionsExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `CoalescePartitionsExec` is a compile error here until it is either
@@ -396,12 +397,15 @@ impl ExecutionPlan for CoalescePartitionsExec {
fetch,
} = self;
let input = ctx.encode_child(input)?;
+ let fetch = fetch
+ .map(|fetch| usize_to_wire(fetch, "CoalescePartitionsExec", "fetch"))
+ .transpose()?;
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(
protobuf::physical_plan_node::PhysicalPlanType::Merge(Box::new(
protobuf::CoalescePartitionsExecNode {
input: Some(Box::new(input)),
- fetch: fetch.map(|f| f as u32),
+ fetch,
},
)),
),
@@ -422,6 +426,7 @@ impl CoalescePartitionsExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
let merge = crate::expect_plan_variant!(
node,
@@ -437,8 +442,11 @@ impl CoalescePartitionsExec {
"CoalescePartitionsExec",
"input",
)?;
+ let fetch = fetch
+ .map(|f| usize_from_wire(f, "CoalescePartitionsExec", "fetch"))
+ .transpose()?;
Ok(Arc::new(
- CoalescePartitionsExec::new(input).with_fetch(fetch.map(|f| f as usize)),
+ CoalescePartitionsExec::new(input).with_fetch(fetch),
))
}
}
diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs
index a1d79890b0..f96fb274a8 100644
--- a/datafusion/physical-plan/src/empty.rs
+++ b/datafusion/physical-plan/src/empty.rs
@@ -213,6 +213,7 @@ impl ExecutionPlan for EmptyExec {
&self,
_ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `EmptyExec` is a compile error here until it is either serialized or
@@ -224,12 +225,13 @@ impl ExecutionPlan for EmptyExec {
cache: _,
} = self;
let schema = schema.as_ref().try_into()?;
+ let partitions = usize_to_wire(*partitions, "EmptyExec", "partitions")?;
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(
protobuf::physical_plan_node::PhysicalPlanType::Empty(
protobuf::EmptyExecNode {
schema: Some(schema),
- partitions: *partitions as u32,
+ partitions,
},
),
),
diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs
index 7cf14ac57d..12771eec78 100644
--- a/datafusion/physical-plan/src/filter.rs
+++ b/datafusion/physical-plan/src/filter.rs
@@ -181,6 +181,10 @@ impl FilterExecBuilder {
/// Build the FilterExec, computing properties once with all configured parameters
pub fn build(self) -> Result<FilterExec> {
+ if self.batch_size == 0 {
+ return plan_err!("FilterExec: batch_size must be greater than 0");
+ }
+
// Validate predicate type
match self.predicate.data_type(self.input.schema().as_ref())? {
DataType::Boolean => {}
@@ -282,6 +286,9 @@ impl FilterExec {
/// Set the batch size
pub fn with_batch_size(&self, batch_size: usize) -> Result<Self> {
+ if batch_size == 0 {
+ return plan_err!("FilterExec: batch_size must be greater than 0");
+ }
Ok(Self {
predicate: Arc::clone(&self.predicate),
input: Arc::clone(&self.input),
@@ -858,6 +865,7 @@ impl ExecutionPlan for FilterExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `FilterExec` is a compile error here until it is either serialized or
@@ -876,6 +884,13 @@ impl ExecutionPlan for FilterExec {
} = self;
let input_node = ctx.encode_child(input)?;
let expr = ctx.encode_expr(predicate)?;
+ if *batch_size == 0 {
+ return plan_err!("FilterExec: batch_size must be greater than 0");
+ }
+ let batch_size = usize_to_wire(*batch_size, "FilterExec", "batch_size")?;
+ let fetch = fetch
+ .map(|fetch| usize_to_wire(fetch, "FilterExec", "fetch"))
+ .transpose()?;
// The identity projection `[0, 1, ..., num_fields - 1]` is the
// canonical wire representation of a full projection, so `None` is
// encoded that way (and decodes back to `None`).
@@ -894,8 +909,8 @@ impl ExecutionPlan for FilterExec {
expr: Some(expr),
default_filter_selectivity: *default_selectivity as u32,
projection,
- batch_size: *batch_size as u32,
- fetch: fetch.map(|f| f as u32),
+ batch_size,
+ fetch,
},
)),
),
@@ -916,6 +931,7 @@ impl FilterExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
let filter_node = crate::expect_plan_variant!(
node,
@@ -957,10 +973,18 @@ impl FilterExec {
} else {
Some(projection_vec)
};
+ // Proto3's zero default means "use the builder default."
+ let batch_size = match *batch_size {
+ 0 => FILTER_EXEC_DEFAULT_BATCH_SIZE,
+ batch_size => usize_from_wire(batch_size, "FilterExec", "batch_size")?,
+ };
+ let fetch = fetch
+ .map(|f| usize_from_wire(f, "FilterExec", "fetch"))
+ .transpose()?;
let filter = FilterExecBuilder::new(predicate, input)
.apply_projection(projection)?
- .with_batch_size(*batch_size as usize)
- .with_fetch(fetch.map(|f| f as usize))
+ .with_batch_size(batch_size)
+ .with_fetch(fetch)
.build()?;
match filter_selectivity {
Ok(filter_selectivity) => Ok(Arc::new(
@@ -1520,6 +1544,22 @@ mod tests {
use crate::test::exec::StatisticsExec;
use arrow::datatypes::{Field, Schema, UnionFields, UnionMode};
+ #[test]
+ fn filter_rejects_zero_batch_size() -> Result<()> {
+ let input: Arc<dyn ExecutionPlan> =
+ Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
+ assert!(
+ FilterExecBuilder::new(lit(true), Arc::clone(&input))
+ .with_batch_size(0)
+ .build()
+ .is_err()
+ );
+
+ let filter = FilterExec::try_new(lit(true), input)?;
+ assert!(filter.with_batch_size(0).is_err());
+ Ok(())
+ }
+
#[tokio::test]
async fn collect_columns_predicates() -> Result<()> {
let schema = test::aggr_test_schema();
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index df10a8a5fc..9a4afe3cdb 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -1996,7 +1996,7 @@ impl HashJoinExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
- use datafusion_common::{internal_datafusion_err, plan_datafusion_err};
+ use datafusion_common::{internal_datafusion_err, utils::usize_from_wire};
use datafusion_proto_models::protobuf;
use std::any::Any;
@@ -2081,21 +2081,8 @@ impl HashJoinExec {
// Restore the row limit that `limit_pushdown` may have pushed into the
// join. The field is presence-tracked, so a message written before it
// existed decodes to `None` (no limit) rather than to `Some(0)`.
- //
- // The conversion is checked, not `as usize`: `fetch` is a `u64` on the
- // wire but a `usize` in the plan, and on a 32-bit target `as usize`
- // truncates. A fetch of `1 << 32` would become `0` -- not merely a
- // wrong limit but the worst one, silently turning the query into an
- // empty result. Report the out-of-range value instead. Please do not
- // "simplify" this back to `as usize`.
let fetch = fetch
- .map(|f| {
- usize::try_from(f).map_err(|_| {
- plan_datafusion_err!(
- "HashJoinExec: fetch value {f} cannot be represented as usize on this target"
- )
- })
- })
+ .map(|fetch| usize_from_wire(fetch, "HashJoinExec", "fetch"))
.transpose()?;
let mut hash_join = HashJoinExecBuilder::new(left, right, on, join_type)
diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs
index 4e5d9200d9..73d98d8105 100644
--- a/datafusion/physical-plan/src/limit.rs
+++ b/datafusion/physical-plan/src/limit.rs
@@ -36,6 +36,8 @@ use crate::{
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::tree_node::TreeNodeRecursion;
+#[cfg(feature = "proto")]
+use datafusion_common::utils::{usize_from_wire, usize_to_wire};
use datafusion_common::{Result, assert_eq_or_internal_err};
use datafusion_execution::TaskContext;
@@ -299,10 +301,10 @@ impl ExecutionPlan for GlobalLimitExec {
protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new(
protobuf::GlobalLimitExecNode {
input: Some(Box::new(input)),
- skip: *skip as u32,
+ skip: usize_to_wire(*skip, "GlobalLimitExec", "skip")?,
fetch: match fetch {
- Some(n) => *n as i64,
- _ => -1, // no limit
+ Some(n) => usize_to_wire(*n, "GlobalLimitExec", "fetch")?,
+ None => -1, // no limit
},
required_ordering,
},
@@ -336,16 +338,15 @@ impl GlobalLimitExec {
} = &**limit;
let input =
ctx.decode_required_child(input.as_deref(), "GlobalLimitExec", "input")?;
- let fetch = if *fetch >= 0 {
- Some(*fetch as usize)
- } else {
- None
- };
+ let fetch = (*fetch >= 0)
+ .then(|| usize_from_wire(*fetch, "GlobalLimitExec", "fetch"))
+ .transpose()?;
let required_ordering = optional_ordering_try_from_proto(
required_ordering,
&ctx.expr_ctx(input.schema().as_ref()),
)?;
- let mut exec = GlobalLimitExec::new(input, *skip as usize, fetch);
+ let skip = usize_from_wire(*skip, "GlobalLimitExec", "skip")?;
+ let mut exec = GlobalLimitExec::new(input, skip, fetch);
exec.set_required_ordering(required_ordering);
Ok(Arc::new(exec))
}
@@ -575,7 +576,7 @@ impl ExecutionPlan for LocalLimitExec {
protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new(
protobuf::LocalLimitExecNode {
input: Some(Box::new(input)),
- fetch: *fetch as u32,
+ fetch: usize_to_wire(*fetch, "LocalLimitExec", "fetch")?,
required_ordering,
},
)),
@@ -610,7 +611,8 @@ impl LocalLimitExec {
required_ordering,
&ctx.expr_ctx(input.schema().as_ref()),
)?;
- let mut exec = LocalLimitExec::new(input, *fetch as usize);
+ let fetch = usize_from_wire(*fetch, "LocalLimitExec", "fetch")?;
+ let mut exec = LocalLimitExec::new(input, fetch);
exec.set_required_ordering(required_ordering);
Ok(Arc::new(exec))
}
diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs
index 33833cd1e7..eb6fe8f976 100644
--- a/datafusion/physical-plan/src/placeholder_row.rs
+++ b/datafusion/physical-plan/src/placeholder_row.rs
@@ -213,6 +213,7 @@ impl ExecutionPlan for PlaceholderRowExec {
&self,
_ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `PlaceholderRowExec` is a compile error here until it is either
@@ -224,12 +225,13 @@ impl ExecutionPlan for PlaceholderRowExec {
cache: _,
} = self;
let schema = schema.as_ref().try_into()?;
+ let partitions = usize_to_wire(*partitions, "PlaceholderRowExec", "partitions")?;
Ok(Some(protobuf::PhysicalPlanNode {
physical_plan_type: Some(
protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow(
protobuf::PlaceholderRowExecNode {
schema: Some(schema),
- partitions: *partitions as u32,
+ partitions,
},
),
),
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index 9880c8b5e5..09655648ab 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -1619,6 +1619,7 @@ impl ExecutionPlan for SortExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `SortExec` is a compile error here until it is either serialized or
@@ -1667,8 +1668,8 @@ impl ExecutionPlan for SortExec {
input: Some(Box::new(input)),
expr,
fetch: match fetch {
- Some(n) => *n as i64,
- None => -1,
+ Some(n) => usize_to_wire(*n, "SortExec", "fetch")?,
+ None => -1, // no limit
},
preserve_partitioning: *preserve_partitioning,
dynamic_filter,
@@ -1685,6 +1686,7 @@ impl SortExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
use protobuf::physical_expr_node::ExprType;
let sort = crate::expect_plan_variant!(
@@ -1728,7 +1730,9 @@ impl SortExec {
let Some(ordering) = LexOrdering::new(exprs) else {
return datafusion_common::internal_err!("SortExec requires an ordering");
};
- let fetch = (*fetch >= 0).then_some(*fetch as usize);
+ let fetch = (*fetch >= 0)
+ .then(|| usize_from_wire(*fetch, "SortExec", "fetch"))
+ .transpose()?;
let new_sort = SortExec::new(ordering, input)
.with_fetch(fetch)
.with_preserve_partitioning(*preserve_partitioning);
@@ -1857,25 +1861,23 @@ mod proto_tests {
assert_eq!(node.fetch, 0);
}
- /// The other end of the range does *not* survive: `fetch` goes onto the
- /// wire as `usize as i64`, so on a 64-bit target `usize::MAX` wraps to
- /// `-1` — the very value that means "absent" — and reads back as an
- /// unlimited sort. That is pre-existing behavior of the `i64` wire field
- /// rather than something this tier introduces; pinning it keeps any change
- /// to the encoding a deliberate one instead of a silent fix.
+ /// `usize::MAX` does not fit the signed wire field on a 64-bit target and
+ /// must be rejected instead of wrapping to the `-1` "absent" encoding.
#[test]
#[cfg(target_pointer_width = "64")]
- fn try_to_proto_wraps_usize_max_fetch_into_the_absent_encoding() {
+ fn try_to_proto_rejects_usize_max_fetch() {
let encoder = StubPlanEncoder::ok();
- let node = encode(&sort_fixture(Some(usize::MAX)), &encoder);
-
- assert_eq!(node.fetch, -1);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+ let err = sort_fixture(Some(usize::MAX))
+ .try_to_proto(&ctx)
+ .unwrap_err();
- // ... and so the limit is gone by the time the node is read back.
- let decoder = StubPlanDecoder::ok();
assert_eq!(
- decode(decodable_node(node.fetch, false), &decoder).fetch(),
- None
+ err.strip_backtrace(),
+ format!(
+ "Error during planning: SortExec: fetch value {} is out of range for the plan wire format",
+ usize::MAX
+ )
);
}
diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
index 9fadeb972d..c988105108 100644
--- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
+++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
@@ -472,6 +472,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
&self,
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ use datafusion_common::utils::usize_to_wire;
use datafusion_proto_models::protobuf;
// Destructure exhaustively (no `..`) so that adding a field to
// `SortPreservingMergeExec` is a compile error here until it is either
@@ -511,7 +512,12 @@ impl ExecutionPlan for SortPreservingMergeExec {
Box::new(protobuf::SortPreservingMergeExecNode {
input: Some(Box::new(input)),
expr,
- fetch: fetch.map(|f| f as i64).unwrap_or(-1),
+ fetch: match fetch {
+ Some(n) => {
+ usize_to_wire(*n, "SortPreservingMergeExec", "fetch")?
+ }
+ None => -1, // no limit
+ },
}),
),
),
@@ -526,6 +532,7 @@ impl SortPreservingMergeExec {
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
use arrow::compute::SortOptions;
+ use datafusion_common::utils::usize_from_wire;
use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
use datafusion_proto_models::protobuf;
let spm = crate::expect_plan_variant!(
@@ -571,7 +578,9 @@ impl SortPreservingMergeExec {
let Some(ordering) = LexOrdering::new(exprs) else {
return internal_err!("SortPreservingMergeExec requires an ordering");
};
- let fetch = (*fetch >= 0).then_some(*fetch as usize);
+ let fetch = (*fetch >= 0)
+ .then(|| usize_from_wire(*fetch, "SortPreservingMergeExec", "fetch"))
+ .transpose()?;
Ok(Arc::new(
SortPreservingMergeExec::new(ordering, input).with_fetch(fetch),
))
@@ -688,20 +697,24 @@ mod proto_tests {
assert_eq!(encode(&spm_fixture(Some(0)), &encoder).fetch, 0);
}
- /// ... and `usize::MAX` does not survive it at all: `fetch` goes onto the
- /// wire as `usize as i64`, so on a 64-bit target it wraps to `-1`, the
- /// "absent" value, and reads back as an unlimited merge. Same pre-existing
- /// `i64` wire behavior as `SortExec`, pinned here for the same reason.
+ /// `usize::MAX` does not fit the signed wire field on a 64-bit target and
+ /// must be rejected instead of wrapping to the `-1` "absent" encoding.
#[test]
#[cfg(target_pointer_width = "64")]
- fn try_to_proto_wraps_usize_max_fetch_into_the_absent_encoding() {
+ fn try_to_proto_rejects_usize_max_fetch() {
let encoder = StubPlanEncoder::ok();
- let node = encode(&spm_fixture(Some(usize::MAX)), &encoder);
-
- assert_eq!(node.fetch, -1);
+ let ctx = ExecutionPlanEncodeCtx::new(&encoder);
+ let err = spm_fixture(Some(usize::MAX))
+ .try_to_proto(&ctx)
+ .unwrap_err();
- let decoder = StubPlanDecoder::ok();
- assert_eq!(decode(decodable_node(node.fetch), &decoder).fetch(), None);
+ assert_eq!(
+ err.strip_backtrace(),
+ format!(
+ "Error during planning: SortPreservingMergeExec: fetch value {} is out of range for the plan wire format",
+ usize::MAX
+ )
+ );
}
#[test]
diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs
index 964225d71c..d93e028051 100644
--- a/datafusion/physical-plan/src/unnest.rs
+++ b/datafusion/physical-plan/src/unnest.rs
@@ -414,6 +414,7 @@ impl UnnestExec {
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>> {
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
let unnest = crate::expect_plan_variant!(
@@ -449,8 +450,8 @@ impl UnnestExec {
.collect();
let struct_column_indices = struct_type_columns
.iter()
- .map(|index| *index as _)
- .collect();
+ .map(|index| usize_from_wire(*index, "UnnestExec", "struct_type_columns"))
+ .collect::<Result<Vec<_>>>()?;
let options = options.as_ref().ok_or_else(|| {
datafusion_common::internal_datafusion_err!(
"UnnestExec is missing required field 'options'"
diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs
index d794e7df9d..dba4c063d1 100644
--- a/datafusion/physical-plan/src/windows/window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs
@@ -423,6 +423,7 @@ impl WindowAggExec {
) -> Result<Arc<dyn ExecutionPlan>> {
use super::BoundedWindowAggExec;
use crate::InputOrderMode;
+ use datafusion_common::utils::usize_from_wire;
use datafusion_proto_models::protobuf;
use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode;
@@ -458,7 +459,12 @@ impl WindowAggExec {
ProtoInputOrderMode::PartiallySorted(
protobuf::PartiallySortedInputOrderMode { columns },
) => InputOrderMode::PartiallySorted(
- columns.iter().map(|column| *column as usize).collect(),
+ columns
+ .iter()
+ .map(|column| {
+ usize_from_wire(*column, "WindowAggExec", "columns")
+ })
+ .collect::<Result<Vec<_>>>()?,
),
ProtoInputOrderMode::Sorted(_) => InputOrderMode::Sorted,
};
diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs
index 169ff7f3d9..cf3a3cb75a 100644
--- a/datafusion/proto-common/src/from_proto/mod.rs
+++ b/datafusion/proto-common/src/from_proto/mod.rs
@@ -813,7 +813,9 @@ impl From<protobuf::Precision> for Precision<usize> {
if let Ok(ScalarValue::UInt64(Some(val))) =
ScalarValue::try_from(&val)
{
- Precision::Exact(val as usize)
+ // A value that does not fit in `usize` decodes as
+ // unknown rather than a silently truncated statistic.
+ usize::try_from(val).map_or(Precision::Absent, Precision::Exact)
} else {
Precision::Absent
}
@@ -826,7 +828,7 @@ impl From<protobuf::Precision> for Precision<usize> {
if let Ok(ScalarValue::UInt64(Some(val))) =
ScalarValue::try_from(&val)
{
- Precision::Inexact(val as usize)
+ usize::try_from(val).map_or(Precision::Absent, Precision::Inexact)
} else {
Precision::Absent
}
diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs
index 08c1a59e46..73cfc6e710 100644
--- a/datafusion/proto/src/logical_plan/mod.rs
+++ b/datafusion/proto/src/logical_plan/mod.rs
@@ -39,6 +39,7 @@ use datafusion_common::file_options::file_type::FileType;
use datafusion_common::format::{
ExplainAnalyzeCategories, ExplainFormat, MetricCategory, MetricType,
};
+use datafusion_common::utils::{usize_from_wire, usize_to_wire};
use datafusion_common::{
NullEquality, Result, TableReference, assert_or_internal_err, context,
internal_datafusion_err, internal_err, not_impl_err, plan_err,
@@ -507,9 +508,11 @@ impl AsLogicalPlan for LogicalPlanNode {
})?;
match plan {
LogicalPlanType::Values(values) => {
- let n_cols = values.n_cols as usize;
+ let n_cols = usize_from_wire(values.n_cols, "Values", "n_cols")?;
let values: Vec<Vec<Expr>> = if values.values_list.is_empty() {
Ok(Vec::new())
+ } else if n_cols == 0 {
+ internal_err!("ValuesNode n_cols must be greater than 0")
} else if values.values_list.len() % n_cols != 0 {
internal_err!(
"Invalid values list length, expect {} to be divisible by {}",
@@ -740,7 +743,9 @@ impl AsLogicalPlan for LogicalPlanNode {
into_logical_plan!(sort.input, ctx, extension_codec)?;
let sort_expr: Vec<SortExpr> =
from_proto::parse_sorts(&sort.expr, ctx, extension_codec)?;
- let fetch: Option<usize> = sort.fetch.try_into().ok();
+ let fetch = (sort.fetch >= 0)
+ .then(|| usize_from_wire(sort.fetch, "Sort", "fetch"))
+ .transpose()?;
LogicalPlanBuilder::from(input)
.sort_with_limit(sort_expr, fetch)?
.build()
@@ -756,16 +761,20 @@ impl AsLogicalPlan for LogicalPlanNode {
)
})?;
+ let decode_partition_count =
+ |count: u64| usize_from_wire(count, "Repartition", "partition_count");
let partitioning_scheme = match pb_partition_method {
PartitionMethod::Hash(protobuf::HashRepartition {
hash_expr: pb_hash_expr,
partition_count,
}) => Partitioning::Hash(
from_proto::parse_exprs(pb_hash_expr, ctx, extension_codec)?,
- *partition_count as usize,
+ decode_partition_count(*partition_count)?,
),
PartitionMethod::RoundRobin(partition_count) => {
- Partitioning::RoundRobinBatch(*partition_count as usize)
+ Partitioning::RoundRobinBatch(decode_partition_count(
+ *partition_count,
+ )?)
}
PartitionMethod::Range(protobuf::RangeRepartition {
sort_expr: pb_sort_expr,
@@ -982,13 +991,11 @@ impl AsLogicalPlan for LogicalPlanNode {
LogicalPlanType::Limit(limit) => {
let input: LogicalPlan =
into_logical_plan!(limit.input, ctx, extension_codec)?;
- let skip = limit.skip.max(0) as usize;
+ let skip = usize_from_wire(limit.skip.max(0), "Limit", "skip")?;
- let fetch = if limit.fetch < 0 {
- None
- } else {
- Some(limit.fetch as usize)
- };
+ let fetch = (limit.fetch >= 0)
+ .then(|| usize_from_wire(limit.fetch, "Limit", "fetch"))
+ .transpose()?;
LogicalPlanBuilder::from(input).limit(skip, fetch)?.build()
}
@@ -1753,8 +1760,11 @@ impl AsLogicalPlan for LogicalPlanNode {
logical_plan_type: Some(LogicalPlanType::Limit(Box::new(
protobuf::LimitNode {
input: Some(Box::new(input)),
- skip: skip as i64,
- fetch: fetch.unwrap_or(i64::MAX as usize) as i64,
+ skip: usize_to_wire(skip, "Limit", "skip")?,
+ fetch: match fetch {
+ Some(f) => usize_to_wire(f, "Limit", "fetch")?,
+ None => -1, // no limit
+ },
},
))),
})
@@ -1771,7 +1781,10 @@ impl AsLogicalPlan for LogicalPlanNode {
protobuf::SortNode {
input: Some(Box::new(input)),
expr: sort_expr,
- fetch: fetch.map(|f| f as i64).unwrap_or(-1i64),
+ fetch: match fetch {
+ Some(f) => usize_to_wire(*f, "Sort", "fetch")?,
+ None => -1, // no limit
+ },
},
))),
})
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
index 5f534ef633..6d43477c1e 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -23,8 +23,10 @@ use std::sync::Arc;
use arrow::datatypes::{IntervalMonthDayNanoType, Schema, SchemaRef};
use datafusion_catalog::memory::MemorySourceConfig;
+use datafusion_common::utils::{usize_from_wire, usize_to_wire};
use datafusion_common::{
DataFusionError, Result, internal_datafusion_err, internal_err, not_impl_err,
+ plan_err,
};
use datafusion_datasource_arrow::source::ArrowSource;
#[cfg(feature = "avro")]
@@ -1401,7 +1403,17 @@ pub trait PhysicalPlanNodeExt: Sized {
};
let table = GenerateSeriesTable::new(Arc::clone(&schema), args);
- let generator = table.as_generator(generate_series.target_batch_size as usize)?;
+ let target_batch_size = usize_from_wire(
+ generate_series.target_batch_size,
+ "GenerateSeriesNode",
+ "target_batch_size",
+ )?;
+ if target_batch_size == 0 {
+ return plan_err!(
+ "GenerateSeriesNode: target_batch_size must be greater than 0"
+ );
+ }
+ let generator = table.as_generator(target_batch_size)?;
Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?))
}
@@ -1446,6 +1458,8 @@ pub trait PhysicalPlanNodeExt: Sized {
}));
}
+ let encode_target_batch_size =
+ |size| usize_to_wire::<u32>(size, "GenerateSeriesNode", "target_batch_size");
if let Some(int_64) = generator_guard
.as_any()
.downcast_ref::<GenericSeriesState<i64>>()
@@ -1453,7 +1467,7 @@ pub trait PhysicalPlanNodeExt: Sized {
let schema = exec.schema();
let node = protobuf::GenerateSeriesNode {
schema: Some(schema.as_ref().try_into()?),
- target_batch_size: int_64.batch_size() as u32,
+ target_batch_size: encode_target_batch_size(int_64.batch_size())?,
args: Some(protobuf::generate_series_node::Args::Int64Args(
protobuf::GenerateSeriesArgsInt64 {
start: *int_64.start(),
@@ -1517,7 +1531,7 @@ pub trait PhysicalPlanNodeExt: Sized {
let node = protobuf::GenerateSeriesNode {
schema: Some(schema.as_ref().try_into()?),
- target_batch_size: timestamp_args.batch_size() as u32,
+ target_batch_size: encode_target_batch_size(timestamp_args.batch_size())?,
args: Some(args),
};
diff --git a/datafusion/proto/tests/cases/plans/leaves.rs b/datafusion/proto/tests/cases/plans/leaves.rs
index afcab2dda2..b9341c4ce0 100644
--- a/datafusion/proto/tests/cases/plans/leaves.rs
+++ b/datafusion/proto/tests/cases/plans/leaves.rs
@@ -19,9 +19,9 @@
use super::{roundtrip_test, roundtrip_test_and_return};
use datafusion::arrow::datatypes::Schema;
-use datafusion::physical_plan::ExecutionPlanProperties;
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::placeholder_row::PlaceholderRowExec;
+use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
use datafusion::prelude::SessionContext;
use datafusion_common::Result;
use datafusion_proto::physical_plan::{
@@ -59,6 +59,33 @@ fn roundtrip_placeholder_row_with_partitions() -> Result<()> {
Ok(())
}
+#[cfg(target_pointer_width = "64")]
+#[test]
+fn leaf_plans_reject_partition_counts_above_wire_range() {
+ let partitions = u32::MAX as usize + 1;
+ let schema = Arc::new(Schema::empty());
+ let plans: [(Arc<dyn ExecutionPlan>, &str); 2] = [
+ (
+ Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(partitions)),
+ "EmptyExec",
+ ),
+ (
+ Arc::new(PlaceholderRowExec::new(schema).with_partitions(partitions)),
+ "PlaceholderRowExec",
+ ),
+ ];
+
+ for (plan, name) in plans {
+ let err = PhysicalPlanNode::try_from_physical_plan(
+ plan,
+ &DefaultPhysicalExtensionCodec {},
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains(name));
+ assert!(err.to_string().contains("partitions"));
+ }
+}
+
/// Plans encoded before `partitions` was added carry no value for it, which
/// decodes as zero and must be treated as the previous default of one.
#[test]
diff --git a/datafusion/proto/tests/cases/plans/limits.rs b/datafusion/proto/tests/cases/plans/limits.rs
index a832d46d53..abcdd35fd8 100644
--- a/datafusion/proto/tests/cases/plans/limits.rs
+++ b/datafusion/proto/tests/cases/plans/limits.rs
@@ -38,13 +38,18 @@ use datafusion::physical_plan::coop::CooperativeExec;
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::expressions::{PhysicalSortExpr, col};
use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
-use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
+use datafusion::physical_plan::{
+ ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
+};
use datafusion::prelude::SessionContext;
use datafusion_common::Result;
use datafusion_common::config::ConfigOptions;
use datafusion_proto::physical_plan::{
- DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter,
+ AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter,
};
+#[cfg(target_pointer_width = "32")]
+use datafusion_proto::protobuf;
+use datafusion_proto::protobuf::PhysicalPlanNode;
use std::sync::Arc;
use std::vec;
@@ -74,6 +79,54 @@ fn roundtrip_global_skip_no_limit() -> Result<()> {
)))
}
+#[cfg(target_pointer_width = "64")]
+#[test]
+fn local_limit_rejects_fetch_above_wire_range() {
+ let fetch = u32::MAX as usize + 1;
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(LocalLimitExec::new(
+ Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
+ fetch,
+ ));
+
+ let err =
+ PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {})
+ .unwrap_err();
+ assert_eq!(
+ err.strip_backtrace(),
+ format!(
+ "Error during planning: LocalLimitExec: fetch value {fetch} is out of range for the plan wire format"
+ )
+ );
+}
+
+#[cfg(target_pointer_width = "32")]
+#[test]
+fn global_limit_rejects_fetch_above_usize() -> Result<()> {
+ let codec = DefaultPhysicalExtensionCodec {};
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(
+ Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
+ 0,
+ Some(1),
+ ));
+ let mut node = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?;
+ let Some(protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(limit)) =
+ &mut node.physical_plan_type
+ else {
+ panic!("expected GlobalLimitExecNode");
+ };
+ limit.fetch = i64::from(u32::MAX) + 1;
+
+ let ctx = SessionContext::new();
+ let err = node
+ .try_into_physical_plan(&ctx.task_ctx(), &codec)
+ .unwrap_err();
+ assert_eq!(
+ err.strip_backtrace(),
+ "Error during planning: GlobalLimitExec: fetch wire value 4294967296 is out of range for usize"
+ );
+ Ok(())
+}
+
/// Sort key at index 1, so a decoder that misbinds column name vs index
/// cannot pass.
fn limit_test_schema() -> Arc<Schema> {
@@ -198,6 +251,23 @@ fn roundtrip_coalesce_batches_with_fetch() -> Result<()> {
))
}
+#[test]
+#[expect(deprecated)]
+fn coalesce_batches_rejects_zero_batch_size_on_encode() {
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(CoalesceBatchesExec::new(
+ Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
+ 0,
+ ));
+
+ let err =
+ PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {})
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("CoalesceBatchesExec: target_batch_size must be greater than 0")
+ );
+}
+
#[test]
fn roundtrip_coalesce_partitions_with_fetch() -> Result<()> {
let field_a = Field::new("a", DataType::Boolean, false);
diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs
index 41bb051c28..fb2fd7f1c1 100644
--- a/datafusion/proto/tests/cases/plans/misc.rs
+++ b/datafusion/proto/tests/cases/plans/misc.rs
@@ -343,12 +343,9 @@ fn parse_hash_partitioning_delegates_to_shared_decoder() -> Result<()> {
}
#[cfg(not(target_pointer_width = "64"))]
- assert!(
- decoded
- .unwrap_err()
- .to_string()
- .contains("Partition count 18446744073709551615 exceeds usize::MAX")
- );
+ assert!(decoded.unwrap_err().to_string().contains(
+ "Partitioning: partition_count wire value 18446744073709551615 is out of range for usize"
+ ));
Ok(())
}
diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs
index 4392ef5238..c17ff47e0f 100644
--- a/datafusion/proto/tests/cases/plans/sources.rs
+++ b/datafusion/proto/tests/cases/plans/sources.rs
@@ -112,6 +112,41 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> {
roundtrip_test(DataSourceExec::from_data_source(scan_config))
}
+#[test]
+fn file_scan_rejects_zero_batch_size() -> Result<()> {
+ let schema = Arc::new(Schema::empty());
+ let scan_config = FileScanConfigBuilder::new(
+ ObjectStoreUrl::local_filesystem(),
+ Arc::new(ParquetSource::new(schema)),
+ )
+ .build();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let mut node = PhysicalPlanNode::try_from_physical_plan(
+ DataSourceExec::from_data_source(scan_config),
+ &codec,
+ )?;
+ let Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetScan(scan)) =
+ node.physical_plan_type.as_mut()
+ else {
+ return internal_err!("Expected ParquetScan node");
+ };
+ scan.base_conf
+ .as_mut()
+ .expect("Parquet scan has a base config")
+ .batch_size = Some(0);
+
+ let ctx = SessionContext::new();
+ let err = node
+ .try_into_physical_plan(ctx.task_ctx().as_ref(), &codec)
+ .expect_err("zero file scan batch size must fail");
+ assert!(
+ err.to_string()
+ .contains("FileScanConfig: batch_size must be greater than 0"),
+ "unexpected error: {err}"
+ );
+ Ok(())
+}
+
#[test]
fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Result<()> {
let file_schema =
diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
index a450f7a7e8..7c916db9c3 100644
--- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -3559,6 +3559,50 @@ async fn roundtrip_custom_listing_tables_schema() -> Result<()> {
Ok(())
}
+#[test]
+fn logical_values_reject_nonempty_rows_with_zero_columns() {
+ let node = protobuf::LogicalPlanNode {
+ logical_plan_type: Some(protobuf::logical_plan_node::LogicalPlanType::Values(
+ protobuf::ValuesNode {
+ n_cols: 0,
+ values_list: vec![protobuf::LogicalExprNode::default()],
+ },
+ )),
+ };
+ let ctx = SessionContext::new();
+ let err =
+ logical_plan_from_bytes(&node.encode_to_vec(), &ctx.task_ctx()).unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("ValuesNode n_cols must be greater than 0")
+ );
+}
+
+#[test]
+fn roundtrip_logical_limit_without_fetch() -> Result<()> {
+ let plan = LogicalPlanBuilder::empty(false).limit(7, None)?.build()?;
+ let bytes = logical_plan_to_bytes(&plan)?;
+ let ctx = SessionContext::new();
+ let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?;
+
+ assert_eq!(plan, round_trip);
+ Ok(())
+}
+
+#[cfg(target_pointer_width = "64")]
+#[test]
+fn roundtrip_logical_limit_at_i64_max() -> Result<()> {
+ let plan = LogicalPlanBuilder::empty(false)
+ .limit(0, Some(i64::MAX as usize))?
+ .build()?;
+ let bytes = logical_plan_to_bytes(&plan)?;
+ let ctx = SessionContext::new();
+ let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?;
+
+ assert_eq!(plan, round_trip);
+ Ok(())
+}
+
#[tokio::test]
async fn roundtrip_custom_listing_tables_schema_table_scan_projection() -> Result<()> {
let ctx = SessionContext::new();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org