From e80c835f00543fad180058f872e648b8d5801230 Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Tue, 18 Aug 2026 22:33:20 +0300 Subject: [PATCH 1/2] fix(proto): check integer conversions in plan serde --- datafusion/common/src/utils/mod.rs | 54 +++++++++++++- .../datasource/src/file_scan_config/proto.rs | 20 ++++- datafusion/datasource/src/memory.rs | 15 +++- datafusion/expr/src/logical_plan/plan.rs | 43 ++++++++++- .../functions-table/src/generate_series.rs | 24 +++++- datafusion/physical-expr/src/partitioning.rs | 73 +++--------------- .../physical-expr/src/scalar_subquery.rs | 12 +-- .../physical-plan/src/aggregates/mod.rs | 8 +- datafusion/physical-plan/src/buffer.rs | 4 +- .../physical-plan/src/coalesce_batches.rs | 34 ++++++++- .../physical-plan/src/coalesce_partitions.rs | 15 +++- datafusion/physical-plan/src/empty.rs | 11 ++- datafusion/physical-plan/src/filter.rs | 51 ++++++++++++- .../physical-plan/src/joins/hash_join/exec.rs | 17 +---- datafusion/physical-plan/src/limit.rs | 24 +++--- .../physical-plan/src/placeholder_row.rs | 11 ++- datafusion/physical-plan/src/sorts/sort.rs | 10 ++- .../src/sorts/sort_preserving_merge.rs | 13 +++- datafusion/physical-plan/src/unnest.rs | 5 +- .../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 +-- .../tests/cases/roundtrip_logical_plan.rs | 44 +++++++++++ 27 files changed, 510 insertions(+), 163 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 8df160187ac7d..127504d62779f 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(value: T, context: &str, field: &str) -> Result +where + T: TryInto + 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>( + value: usize, + context: &str, + field: &str, +) -> Result { + 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::(42, "SomeExec", "fetch").unwrap(), 42); + let err = usize_to_wire::(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 = vec![ diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index d7135173c8934..1c08d817ef401 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -38,6 +38,7 @@ use std::sync::Arc; use arrow::datatypes::Schema; +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}; @@ -117,7 +118,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::(limit, "FileScanConfig", "limit")) + .transpose()? + .map(|limit| protobuf::ScanLimit { limit }), projection: vec![], schema: Some((&schema).try_into()?), table_partition_cols: self @@ -219,14 +224,23 @@ 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()?; 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); Ok(config_builder.build()) } diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 2370ed87a2954..c99d63005db65 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -275,6 +275,7 @@ impl DataSource for MemorySourceConfig { &self, ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto; use datafusion_proto_models::protobuf; @@ -310,7 +311,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()?, }, ), ), @@ -677,6 +683,7 @@ impl MemorySourceConfig { ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { 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; @@ -713,8 +720,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 1a141ea52a13a..b6b8b9d7d8937 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 _; @@ -3704,7 +3704,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) } @@ -3722,7 +3727,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) } @@ -4954,6 +4964,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 8d9327c5acfa2..4b4fc01af9a93 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>> { + if batch_size == 0 { + return plan_err!("GenerateSeriesTable: batch_size must be greater than 0"); + } let generator: Arc> = 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 98f082f7256db..70ba92a0a9ecc 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 { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; + let partition_count = + |count: usize| usize_to_wire::(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> { + 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::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::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 473b52a5cb45c..ac473ea9dd78d 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> { + 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 1671735ec30a3..d589214e8da6b 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -2471,6 +2471,7 @@ impl AggregateExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + 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; @@ -2645,11 +2646,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 24cca6b0b17f4..bd52f1bab0513 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -357,6 +357,7 @@ impl BufferExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_from_wire; use datafusion_proto_models::protobuf; let buffer = crate::expect_plan_variant!( node, @@ -365,7 +366,8 @@ impl BufferExec { ); let input = ctx.decode_required_child(buffer.input.as_deref(), "BufferExec", "input")?; - Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) + let capacity = usize_from_wire(buffer.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 cb0f9b2ce4b36..a8c8b176e074f 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -322,15 +322,28 @@ impl ExecutionPlan for CoalesceBatchesExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; + let to_wire = + |value, field| usize_to_wire::(value, "CoalesceBatchesExec", field); + if self.target_batch_size() == 0 { + return datafusion_common::plan_err!( + "CoalesceBatchesExec: target_batch_size must be greater than 0" + ); + } + let target_batch_size = to_wire(self.target_batch_size(), "target_batch_size")?; + let fetch = self + .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: self.target_batch_size() as u32, - fetch: self.fetch().map(|n| n as u32), + target_batch_size, + fetch, }), ), ), @@ -355,6 +368,7 @@ impl CoalesceBatchesExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_from_wire; use datafusion_proto_models::protobuf; let coalesce_batches = crate::expect_plan_variant!( node, @@ -366,9 +380,21 @@ impl CoalesceBatchesExec { "CoalesceBatchesExec", "input", )?; + let from_wire = + |value, field| usize_from_wire(value, "CoalesceBatchesExec", field); + let target_batch_size = + from_wire(coalesce_batches.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 = coalesce_batches + .fetch + .map(|fetch| from_wire(fetch, "fetch")) + .transpose()?; Ok(Arc::new( - CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) - .with_fetch(coalesce_batches.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 6f58eb2f1e6be..246c34ff2701f 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -383,14 +383,19 @@ impl ExecutionPlan for CoalescePartitionsExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; + let fetch = self + .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: self.fetch().map(|f| f as u32), + fetch, }, )), ), @@ -411,6 +416,7 @@ impl CoalescePartitionsExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_from_wire; use datafusion_proto_models::protobuf; let merge = crate::expect_plan_variant!( node, @@ -422,9 +428,12 @@ impl CoalescePartitionsExec { "CoalescePartitionsExec", "input", )?; + let fetch = merge + .fetch + .map(|f| usize_from_wire(f, "CoalescePartitionsExec", "fetch")) + .transpose()?; Ok(Arc::new( - CoalescePartitionsExec::new(input) - .with_fetch(merge.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 dd08ff36a9d88..79e55e3f4a157 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -213,17 +213,20 @@ impl ExecutionPlan for EmptyExec { &self, _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let schema = self.schema().as_ref().try_into()?; + let partitions = usize_to_wire( + self.properties().output_partitioning().partition_count(), + "EmptyExec", + "partitions", + )?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::Empty( protobuf::EmptyExecNode { schema: Some(schema), - partitions: self - .properties() - .output_partitioning() - .partition_count() as u32, + partitions, }, ), ), diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index c62db109ead1c..0c23e3a78a544 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 { + 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 { + 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), @@ -849,9 +856,18 @@ impl ExecutionPlan for FilterExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; let expr = ctx.encode_expr(self.predicate())?; + if self.batch_size() == 0 { + return plan_err!("FilterExec: batch_size must be greater than 0"); + } + let batch_size = usize_to_wire(self.batch_size(), "FilterExec", "batch_size")?; + let fetch = self + .fetch() + .map(|fetch| usize_to_wire(fetch, "FilterExec", "fetch")) + .transpose()?; // Preserve the exact wire format: `None` (full projection) is serialized // as the identity projection `[0, 1, ..., num_fields - 1]` so that it is // distinguishable from an explicit projection on decode. @@ -870,8 +886,8 @@ impl ExecutionPlan for FilterExec { expr: Some(expr), default_filter_selectivity: self.default_selectivity() as u32, projection, - batch_size: self.batch_size() as u32, - fetch: self.fetch().map(|f| f as u32), + batch_size, + fetch, }, )), ), @@ -892,6 +908,7 @@ impl FilterExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_from_wire; use datafusion_proto_models::protobuf; let filter = crate::expect_plan_variant!( node, @@ -924,10 +941,20 @@ impl FilterExec { } else { Some(projection_vec) }; + // Proto3's zero default means "use the builder default." + let batch_size = match filter.batch_size { + 0 => FILTER_EXEC_DEFAULT_BATCH_SIZE, + batch_size => usize_from_wire(batch_size, "FilterExec", "batch_size")?, + }; let filter = FilterExecBuilder::new(predicate, input) .apply_projection(projection)? - .with_batch_size(filter.batch_size as usize) - .with_fetch(filter.fetch.map(|f| f as usize)) + .with_batch_size(batch_size) + .with_fetch( + filter + .fetch + .map(|f| usize_from_wire(f, "FilterExec", "fetch")) + .transpose()?, + ) .build()?; match filter_selectivity { Ok(filter_selectivity) => Ok(Arc::new( @@ -1417,6 +1444,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 = + 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 08d209003ad91..78171339fd7bf 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1921,7 +1921,7 @@ impl HashJoinExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { - 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; @@ -2001,22 +2001,9 @@ 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 = hashjoin .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 dd62c93d1cfe0..372b227990e97 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; @@ -288,10 +290,10 @@ impl ExecutionPlan for GlobalLimitExec { protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new( protobuf::GlobalLimitExecNode { input: Some(Box::new(input)), - skip: self.skip() as u32, + skip: usize_to_wire(self.skip(), "GlobalLimitExec", "skip")?, fetch: match self.fetch() { - Some(n) => n as i64, - _ => -1, // no limit + Some(n) => usize_to_wire(n, "GlobalLimitExec", "fetch")?, + None => -1, // no limit }, required_ordering, }, @@ -319,16 +321,15 @@ impl GlobalLimitExec { "GlobalLimitExec", "input", )?; - let fetch = if limit.fetch >= 0 { - Some(limit.fetch as usize) - } else { - None - }; + let fetch = (limit.fetch >= 0) + .then(|| usize_from_wire(limit.fetch, "GlobalLimitExec", "fetch")) + .transpose()?; let required_ordering = optional_ordering_try_from_proto( &limit.required_ordering, &ctx.expr_ctx(input.schema().as_ref()), )?; - let mut exec = GlobalLimitExec::new(input, limit.skip as usize, fetch); + let skip = usize_from_wire(limit.skip, "GlobalLimitExec", "skip")?; + let mut exec = GlobalLimitExec::new(input, skip, fetch); exec.set_required_ordering(required_ordering); Ok(Arc::new(exec)) } @@ -548,7 +549,7 @@ impl ExecutionPlan for LocalLimitExec { protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new( protobuf::LocalLimitExecNode { input: Some(Box::new(input)), - fetch: self.fetch() as u32, + fetch: usize_to_wire(self.fetch(), "LocalLimitExec", "fetch")?, required_ordering, }, )), @@ -576,7 +577,8 @@ impl LocalLimitExec { &limit.required_ordering, &ctx.expr_ctx(input.schema().as_ref()), )?; - let mut exec = LocalLimitExec::new(input, limit.fetch as usize); + let fetch = usize_from_wire(limit.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 67c063b65cbc6..1df406ec5db51 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -213,17 +213,20 @@ impl ExecutionPlan for PlaceholderRowExec { &self, _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let schema = self.schema().as_ref().try_into()?; + let partitions = usize_to_wire( + self.properties().output_partitioning().partition_count(), + "PlaceholderRowExec", + "partitions", + )?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( protobuf::PlaceholderRowExecNode { schema: Some(schema), - partitions: self - .properties() - .output_partitioning() - .partition_count() as u32, + partitions, }, ), ), diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 6c782f5134484..c8d5698e43811 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1603,6 +1603,7 @@ impl ExecutionPlan for SortExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; let expr = self @@ -1635,8 +1636,8 @@ impl ExecutionPlan for SortExec { input: Some(Box::new(input)), expr, fetch: match self.fetch() { - Some(n) => n as i64, - None => -1, + Some(n) => usize_to_wire(n, "SortExec", "fetch")?, + None => -1, // no limit }, preserve_partitioning: self.preserve_partitioning(), dynamic_filter, @@ -1653,6 +1654,7 @@ impl SortExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_from_wire; use datafusion_proto_models::protobuf; use protobuf::physical_expr_node::ExprType; let sort = crate::expect_plan_variant!( @@ -1689,7 +1691,9 @@ impl SortExec { let Some(ordering) = LexOrdering::new(exprs) else { return datafusion_common::internal_err!("SortExec requires an ordering"); }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as usize); + let fetch = (sort.fetch >= 0) + .then(|| usize_from_wire(sort.fetch, "SortExec", "fetch")) + .transpose()?; let new_sort = SortExec::new(ordering, input) .with_fetch(fetch) .with_preserve_partitioning(sort.preserve_partitioning); diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index ad17f2c2136af..81b7b31472f05 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> { + use datafusion_common::utils::usize_to_wire; use datafusion_proto_models::protobuf; let input = ctx.encode_child(self.input())?; let expr = self @@ -496,7 +497,12 @@ impl ExecutionPlan for SortPreservingMergeExec { Box::new(protobuf::SortPreservingMergeExecNode { input: Some(Box::new(input)), expr, - fetch: self.fetch().map(|f| f as i64).unwrap_or(-1), + fetch: match self.fetch() { + Some(n) => { + usize_to_wire(n, "SortPreservingMergeExec", "fetch")? + } + None => -1, // no limit + }, }), ), ), @@ -511,6 +517,7 @@ impl SortPreservingMergeExec { ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { 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!( @@ -554,7 +561,9 @@ impl SortPreservingMergeExec { let Some(ordering) = LexOrdering::new(exprs) else { return internal_err!("SortPreservingMergeExec requires an ordering"); }; - let fetch = (spm.fetch >= 0).then_some(spm.fetch as usize); + let fetch = (spm.fetch >= 0) + .then(|| usize_from_wire(spm.fetch, "SortPreservingMergeExec", "fetch")) + .transpose()?; Ok(Arc::new( SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), )) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 3fa274b27a7bd..95f1a08a368b2 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -396,6 +396,7 @@ impl UnnestExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::utils::usize_from_wire; use datafusion_proto_models::protobuf; let unnest = crate::expect_plan_variant!( @@ -431,8 +432,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::>>()?; 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 d794e7df9d0a9..dba4c063d1911 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> { 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::>>()?, ), 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 169ff7f3d9ff2..cf3a3cb75a0f0 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 for Precision { 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 for Precision { 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 08c1a59e46c6c..73cfc6e710d15 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> = 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 = from_proto::parse_sorts(&sort.expr, ctx, extension_codec)?; - let fetch: Option = 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 222901aff5211..e9bdd6a756ba0 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")] @@ -1372,7 +1374,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])?)) } @@ -1417,6 +1429,8 @@ pub trait PhysicalPlanNodeExt: Sized { })); } + let encode_target_batch_size = + |size| usize_to_wire::(size, "GenerateSeriesNode", "target_batch_size"); if let Some(int_64) = generator_guard .as_any() .downcast_ref::>() @@ -1424,7 +1438,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(), @@ -1488,7 +1502,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 afcab2dda24bc..b9341c4ce0df3 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, &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 a832d46d53152..abcdd35fd840f 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 = 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 = 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 { @@ -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 = 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 41bb051c28730..fb2fd7f1c1b5e 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/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index a450f7a7e888f..7c916db9c3cd1 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(); From 2fa2eea0ec2683cbe883046676b295ed2b040a2f Mon Sep 17 00:00:00 2001 From: Burak Sen Date: Thu, 20 Aug 2026 11:27:36 +0300 Subject: [PATCH 2/2] test(proto): expect sort fetch overflow errors --- datafusion/physical-plan/src/sorts/sort.rs | 26 +++++++++---------- .../src/sorts/sort_preserving_merge.rs | 24 ++++++++++------- 2 files changed, 26 insertions(+), 24 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 2364c8cf586f1..3d31566a890b3 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1821,25 +1821,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 4e14fb3d6bd37..1fbe6f4ea97ab 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -680,20 +680,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]