diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 142d3cd8e2cee..e047db39a5740 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 efd17e4f57091..4071ea471b33f 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::(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 a4f51c28f83ac..7c10dba981c82 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> { + 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> { 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 ab3c26795e74a..1a8cd81aa74bd 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 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 2dc41adfd6c62..69a24ce17e852 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> { + 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 72df0f5345041..e1452b75bc4c3 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> { + 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 87957ced7b11c..0b3359bc7c583 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> { + 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::(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> { + 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 7c8bca772e21d..7438f1c0ca4e3 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> { + 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> { + 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 a1d79890b0815..f96fb274a80b3 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> { + 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 7cf14ac57d56e..12771eec78470 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), @@ -858,6 +865,7 @@ impl ExecutionPlan for FilterExec { &self, ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { + 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> { + 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 = + 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 df10a8a5fcad5..9a4afe3cdbce3 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> { - 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 4e5d9200d9df9..73d98d8105215 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 33833cd1e7811..eb6fe8f97649e 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> { + 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 9880c8b5e5eb1..09655648ab62b 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> { + 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> { + 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 9fadeb972d4ed..c988105108679 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; // 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> { 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 964225d71ccd6..d93e0280515c6 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> { + 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::>>()?; 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 5f534ef63336d..6d43477c1e542 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::(size, "GenerateSeriesNode", "target_batch_size"); if let Some(int_64) = generator_guard .as_any() .downcast_ref::>() @@ -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 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/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index 4392ef52381ce..c17ff47e0f472 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 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();