diff --git a/datafusion-examples/examples/custom_data_source/adapter_serialization.rs b/datafusion-examples/examples/custom_data_source/adapter_serialization.rs index f18b888f3eb56..d8dba6592966c 100644 --- a/datafusion-examples/examples/custom_data_source/adapter_serialization.rs +++ b/datafusion-examples/examples/custom_data_source/adapter_serialization.rs @@ -21,11 +21,11 @@ //! trait's interception methods (`execution_plan_to_proto` and //! `proto_to_execution_plan`) to implement custom serialization logic. //! -//! The key insight is that `FileScanConfig::expr_adapter_factory` is NOT serialized by -//! default. This example shows how to: +//! The key insight is that a non-default `FileScanConfig::expr_adapter_factory` +//! cannot be serialized by the default protobuf path. This example shows how to: //! 1. Detect plans with custom adapters during serialization //! 2. Wrap them as Extension nodes with JSON-serialized adapter metadata -//! 3. Store the inner DataSourceExec (without adapter) as a child in the extension's inputs field +//! 3. Explicitly remove the captured adapter from the inner `DataSourceExec` //! 4. Unwrap and restore the adapter during deserialization //! //! This demonstrates nested serialization (protobuf outer, JSON inner) and the @@ -333,10 +333,14 @@ impl PhysicalProtoConverterExtension for AdapterPreservingCodec { // 1. Create adapter metadata let adapter_metadata = AdapterMetadata { tag }; - // 2. Serialize the inner plan to protobuf - // Note that this will drop the custom adapter since the default serialization cannot handle it + // 2. The metadata is captured above, so explicitly remove the custom + // adapter before using the guarded default serializer for the child. + let mut inner_config = config.clone(); + inner_config.expr_adapter_factory = None; + let inner_plan: Arc = + DataSourceExec::from_data_source(inner_config); let inner_proto = PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(plan), + inner_plan, extension_codec, self, )?; diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 4071ea471b33f..d94fda65cac17 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -42,13 +42,15 @@ 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::Partitioning; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; -use datafusion_physical_expr::{LexOrdering, Partitioning}; use datafusion_physical_expr_common::sort_expr::{ - sort_exprs_try_from_proto, sort_exprs_try_to_proto, + optional_ordering_try_from_proto, sort_exprs_try_to_proto, }; use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; -use datafusion_proto_models::datafusion_common::CompressionTypeVariant as ProtoCompressionTypeVariant; +use datafusion_proto_models::datafusion_common::{ + CompressionTypeVariant as ProtoCompressionTypeVariant, Schema as ProtoSchema, +}; use datafusion_proto_models::protobuf; use crate::file::FileSource; @@ -68,42 +70,67 @@ impl FileScanConfig { &self, ctx: &ExecutionPlanEncodeCtx<'_>, ) -> Result { - let file_groups = self - .file_groups + // Exhaustive destructure: adding a field to `FileScanConfig` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + object_store_url, + file_groups, + constraints, + limit, + preserve_order, + output_ordering, + file_compression_type, + file_source, + batch_size, + expr_adapter_factory, + // Serialized through `statistics()` so its pushed-filter policy + // remains centralized. + statistics: _, + output_partitioning, + } = self; + + // Non-default factories are executable behavior with no protobuf + // representation; silently replacing one can change scan results. + if expr_adapter_factory + .as_ref() + .is_some_and(|factory| !factory.is_equivalent_to_default()) + { + return datafusion_common::not_impl_err!( + "FileScanConfig with a non-default expr_adapter_factory cannot be serialized" + ); + } + + let proto_file_groups = file_groups .iter() .map(TryInto::try_into) .collect::>>()?; - let mut output_ordering = vec![]; - for order in &self.output_ordering { + let mut proto_output_ordering = vec![]; + for order in output_ordering { let nodes = sort_exprs_try_to_proto(order.iter(), &ctx.expr_ctx())?; - output_ordering.push(protobuf::PhysicalSortExprNodeCollection { + proto_output_ordering.push(protobuf::PhysicalSortExprNodeCollection { physical_sort_expr_nodes: nodes, }); } - let output_partitioning = self - .output_partitioning + let proto_output_partitioning = output_partitioning .as_ref() .map(|partitioning| partitioning.try_to_proto(&ctx.expr_ctx())) .transpose()?; - // Fields must be added to the schema so that they can persist in the - // protobuf, and then removed from the schema in `try_from_proto`. - let mut fields = self - .file_schema() - .fields() - .iter() - .cloned() - .collect::>(); - fields.extend(self.table_partition_cols().iter().cloned()); - let schema = - Schema::new(fields).with_metadata(self.file_schema().metadata.clone()); - - let projection_exprs = self - .file_source() + let table_schema = file_source.table_schema(); + let file_schema = table_schema.file_schema(); + let table_partition_cols = table_schema.table_partition_cols(); + + // Partition fields must be added to the schema so they can persist in + // protobuf and then be removed again in `parse_table_schema_from_proto`. + let mut fields = file_schema.fields().iter().cloned().collect::>(); + fields.extend(table_partition_cols.iter().cloned()); + let schema = Schema::new(fields).with_metadata(file_schema.metadata.clone()); + + let projection_exprs = file_source .projection() - .as_ref() .map(|projection_exprs| { Ok::<_, DataFusionError>(protobuf::ProjectionExprs { projections: projection_exprs @@ -119,35 +146,36 @@ impl FileScanConfig { }) .transpose()?; - let file_compression_type = - self.file_compression_type.is_compressed().then(|| { + let proto_file_compression_type = + file_compression_type.is_compressed().then(|| { let compression: ProtoCompressionTypeVariant = - (*self.file_compression_type.get_variant()).into(); + (*file_compression_type.get_variant()).into(); compression as i32 }); + let statistics = self.statistics(); Ok(protobuf::FileScanExecConf { - file_groups, - statistics: Some((&self.statistics()).into()), - limit: self - .limit + file_groups: proto_file_groups, + statistics: Some((&statistics).into()), + limit: limit .map(|limit| usize_to_wire::(limit, "FileScanConfig", "limit")) .transpose()? .map(|limit| protobuf::ScanLimit { limit }), + // Superseded by `projection_exprs`; kept empty for wire compatibility. projection: vec![], schema: Some((&schema).try_into()?), - table_partition_cols: self - .table_partition_cols() + table_partition_cols: table_partition_cols .iter() .map(|x| x.name().clone()) .collect::>(), - object_store_url: self.object_store_url.to_string(), - output_ordering, - constraints: Some(self.constraints.clone().into()), - batch_size: self.batch_size.map(|s| s as u64), + object_store_url: object_store_url.to_string(), + output_ordering: proto_output_ordering, + constraints: Some(constraints.clone().into()), + batch_size: batch_size.map(|size| size as u64), projection_exprs, - output_partitioning, - file_compression_type, + output_partitioning: proto_output_partitioning, + file_compression_type: proto_file_compression_type, + preserve_order: Some(*preserve_order), }) } @@ -162,10 +190,32 @@ impl FileScanConfig { ctx: &ExecutionPlanDecodeCtx<'_>, file_source: Arc, ) -> Result { - let schema = parse_file_scan_schema(conf)?; + // Destructure exhaustively so a newly added protobuf field must be + // handled here instead of being silently ignored. + let protobuf::FileScanExecConf { + file_groups, + schema: proto_schema, + // Superseded by `projection_exprs`; current encoders leave this + // legacy column-index projection empty. + projection: _, + limit, + statistics, + // Used to construct `file_source` via `parse_table_schema_from_proto` + // before this hook is called. + table_partition_cols: _, + object_store_url, + output_ordering, + constraints, + batch_size, + projection_exprs, + output_partitioning, + file_compression_type, + preserve_order, + } = conf; + + let expression_schema = parse_file_scan_schema(proto_schema)?; - let constraints = conf - .constraints + let decoded_constraints = constraints .as_ref() .ok_or_else(|| { internal_datafusion_err!( @@ -173,8 +223,7 @@ impl FileScanConfig { ) })? .try_into()?; - let statistics = conf - .statistics + let decoded_statistics = statistics .as_ref() .ok_or_else(|| { internal_datafusion_err!( @@ -183,37 +232,41 @@ impl FileScanConfig { })? .try_into()?; - let file_groups = conf - .file_groups + let decoded_file_groups = file_groups .iter() .map(TryInto::try_into) .collect::>>()?; - let object_store_url = match conf.object_store_url.is_empty() { - false => ObjectStoreUrl::parse(&conf.object_store_url)?, + let decoded_object_store_url = match object_store_url.is_empty() { + false => ObjectStoreUrl::parse(object_store_url)?, true => ObjectStoreUrl::local_filesystem(), }; - let mut output_ordering = vec![]; - for node_collection in &conf.output_ordering { - let sort_exprs = sort_exprs_try_from_proto( - &node_collection.physical_sort_expr_nodes, - &ctx.expr_ctx(&schema), - )?; - output_ordering.extend(LexOrdering::new(sort_exprs)); + let mut decoded_output_ordering = vec![]; + for node_collection in output_ordering { + let protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + } = node_collection; + if let Some(ordering) = optional_ordering_try_from_proto( + physical_sort_expr_nodes, + &ctx.expr_ctx(&expression_schema), + )? { + decoded_output_ordering.push(ordering); + } } - let output_partitioning = conf - .output_partitioning + let decoded_output_partitioning = output_partitioning .as_ref() .map(|partitioning| { - Partitioning::try_from_proto(partitioning, &ctx.expr_ctx(&schema)) + Partitioning::try_from_proto( + partitioning, + &ctx.expr_ctx(&expression_schema), + ) }) .transpose()? .flatten(); - let file_compression_type = conf - .file_compression_type + let decoded_file_compression_type = file_compression_type .map(|value| { let compression = ProtoCompressionTypeVariant::try_from(value).map_err(|_| { @@ -226,54 +279,64 @@ impl FileScanConfig { .unwrap_or(FileCompressionType::UNCOMPRESSED); // Parse projection expressions if present and apply to the file source. - let file_source = if let Some(proto_projection_exprs) = &conf.projection_exprs { - let projection_exprs: Vec = proto_projection_exprs - .projections + let decoded_file_source = if let Some(proto_projection_exprs) = projection_exprs { + let protobuf::ProjectionExprs { projections } = proto_projection_exprs; + let decoded_projection_exprs: Vec = projections .iter() .map(|proto_expr| { + let protobuf::ProjectionExpr { alias, expr } = proto_expr; let expr = ctx.decode_expr( - proto_expr.expr.as_ref().ok_or_else(|| { + expr.as_ref().ok_or_else(|| { internal_datafusion_err!("ProjectionExpr missing expr field") })?, - &schema, + &expression_schema, )?; - Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) + Ok(ProjectionExpr::new(expr, alias.clone())) }) .collect::>>()?; - let projection_exprs = ProjectionExprs::new(projection_exprs); + let projection = ProjectionExprs::new(decoded_projection_exprs); file_source - .try_pushdown_projection(&projection_exprs)? + .try_pushdown_projection(&projection)? .unwrap_or(file_source) } else { file_source }; - let limit = conf - .limit + let decoded_limit = limit .as_ref() - .map(|limit| usize_from_wire(limit.limit, "FileScanConfig", "limit")) + .map(|limit| { + let protobuf::ScanLimit { limit } = limit; + usize_from_wire(*limit, "FileScanConfig", "limit") + }) .transpose()?; - let batch_size = conf - .batch_size + let decoded_batch_size = batch_size .map(|size| usize_from_wire(size, "FileScanConfig", "batch_size")) .transpose()?; - if batch_size == Some(0) { + if decoded_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(limit) - .with_output_ordering(output_ordering) - .with_output_partitioning(output_partitioning) - .with_batch_size(batch_size) - .with_file_compression_type(file_compression_type); - Ok(config_builder.build()) + let mut config = + FileScanConfigBuilder::new(decoded_object_store_url, decoded_file_source) + .with_file_groups(decoded_file_groups) + .with_constraints(decoded_constraints) + .with_statistics(decoded_statistics) + .with_limit(decoded_limit) + .with_output_ordering(decoded_output_ordering) + .with_output_partitioning(decoded_output_partitioning) + .with_batch_size(decoded_batch_size) + .with_file_compression_type(decoded_file_compression_type) + .build(); + + // Presence distinguishes a new explicit `false` from a legacy payload, + // which must retain the builder's ordering-derived behavior. + if let Some(preserve_order) = preserve_order { + config.preserve_order = *preserve_order; + } + Ok(config) } /// Parse a [`TableSchema`] (file schema + partition columns) from a @@ -284,7 +347,7 @@ impl FileScanConfig { pub fn parse_table_schema_from_proto( conf: &protobuf::FileScanExecConf, ) -> Result { - let schema = parse_file_scan_schema(conf)?; + let schema = parse_file_scan_schema(&conf.schema)?; // Reacquire the partition column types from the schema before removing // them below. @@ -317,15 +380,9 @@ impl FileScanConfig { } /// Parse the full (file + partition columns) schema off the base conf. -fn parse_file_scan_schema(conf: &protobuf::FileScanExecConf) -> Result> { - let schema: Schema = conf - .schema - .as_ref() - .ok_or_else(|| { - internal_datafusion_err!( - "FileScanExecConf is missing required field 'schema'" - ) - })? - .try_into()?; - Ok(Arc::new(schema)) +fn parse_file_scan_schema(schema: &Option) -> Result> { + let proto_schema = schema.as_ref().ok_or_else(|| { + internal_datafusion_err!("FileScanExecConf is missing required field 'schema'") + })?; + Ok(Arc::new(proto_schema.try_into()?)) } diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 7c10dba981c82..4c79cf4a9851d 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -284,25 +284,38 @@ impl DataSource for MemorySourceConfig { use datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto; use datafusion_proto_models::protobuf; - let partitions = self - .partitions + // Exhaustive destructure: adding a field to `MemorySourceConfig` + // without deciding how it is serialized is a compile error, not a + // silent round-trip gap. + let Self { + partitions: source_partitions, + schema, + // Derived from `schema` and `projection` by `try_new` on decode. + projected_schema: _, + projection: source_projection, + sort_information, + show_sizes, + fetch, + } = self; + + let proto_partitions = source_partitions .iter() .map(|batches| record_batches_to_ipc_bytes(batches)) .collect::>>()?; // Proto3 can't tell `None` from `Some(vec![])`; encode the latter // as the `[u32::MAX]` sentinel, matching the join/filter nodes. - let projection = match self.projection.as_ref() { + let proto_projection = match source_projection.as_ref() { None => Vec::new(), Some(v) if v.is_empty() => vec![u32::MAX], Some(v) => v.iter().map(|x| *x as u32).collect(), }; - let mut sort_information = Vec::with_capacity(self.sort_information.len()); - for ordering in &self.sort_information { + let mut proto_sort_information = Vec::with_capacity(sort_information.len()); + for ordering in sort_information { let physical_sort_expr_nodes = sort_exprs_try_to_proto(ordering.iter(), &ctx.expr_ctx())?; - sort_information.push(protobuf::PhysicalSortExprNodeCollection { + proto_sort_information.push(protobuf::PhysicalSortExprNodeCollection { physical_sort_expr_nodes, }); } @@ -311,13 +324,12 @@ impl DataSource for MemorySourceConfig { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::MemoryScan( protobuf::MemoryScanExecNode { - partitions, - schema: Some(self.schema.as_ref().try_into()?), - projection, - sort_information, - show_sizes: self.show_sizes, - fetch: self - .fetch + partitions: proto_partitions, + schema: Some(schema.as_ref().try_into()?), + projection: proto_projection, + sort_information: proto_sort_information, + show_sizes: *show_sizes, + fetch: fetch .map(|fetch| { usize_to_wire(fetch, "MemoryScanExecNode", "fetch") }) @@ -507,11 +519,28 @@ impl MemorySourceConfig { mut self, mut sort_information: Vec, ) -> Result { - // All sort expressions must refer to the original schema - let fields = self.schema.fields(); + // All sort expressions must refer to the original schema. + Self::validate_sort_information(&sort_information, &self.schema, "original")?; + + // If there is a projection on the source, we also need to project orderings + if self.projection.is_some() { + sort_information = + project_orderings(&sort_information, &self.projected_schema); + } + + self.sort_information = sort_information; + Ok(self) + } + + fn validate_sort_information( + sort_information: &[LexOrdering], + schema: &Schema, + schema_name: &str, + ) -> Result<()> { + let fields = schema.fields(); let ambiguous_column = sort_information .iter() - .flat_map(|ordering| ordering.clone()) + .flat_map(|ordering| ordering.iter()) .flat_map(|expr| collect_columns(&expr.expr)) .find(|col| { fields @@ -521,18 +550,10 @@ impl MemorySourceConfig { }); assert_or_internal_err!( ambiguous_column.is_none(), - "Column {:?} is not found in the original schema of the MemorySourceConfig", + "Column {:?} is not found in the {schema_name} schema of the MemorySourceConfig", ambiguous_column.as_ref().unwrap() ); - - // If there is a projection on the source, we also need to project orderings - if self.projection.is_some() { - sort_information = - project_orderings(&sort_information, &self.projected_schema); - } - - self.sort_information = sort_information; - Ok(self) + Ok(()) } /// Arc clone of ref to original schema @@ -689,7 +710,7 @@ impl MemorySourceConfig { ) -> 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_physical_expr_common::sort_expr::optional_ordering_try_from_proto; use datafusion_proto_models::protobuf; let scan = datafusion_physical_plan::expect_plan_variant!( @@ -698,41 +719,61 @@ impl MemorySourceConfig { "MemorySourceConfig", ); - let partitions = scan - .partitions + // Destructure exhaustively so a newly added protobuf field must be + // handled here instead of being silently ignored. + let protobuf::MemoryScanExecNode { + partitions: proto_partitions, + schema: encoded_schema, + projection: proto_projection, + sort_information: proto_sort_information, + show_sizes, + fetch: proto_fetch, + } = scan; + + let partitions = proto_partitions .iter() .map(|buf| record_batches_from_ipc_bytes(buf)) .collect::>>()?; - let proto_schema = scan.schema.as_ref().ok_or_else(|| { + let proto_schema = encoded_schema.as_ref().ok_or_else(|| { internal_datafusion_err!("schema in MemoryScanExecNode is missing.") })?; - let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); + let source_schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); // Preserve the empty-projection sentinel written by `try_to_proto`. - let projection = match scan.projection.as_slice() { + let projection = match proto_projection.as_slice() { [] => None, [u32::MAX] => Some(Vec::new()), indices => Some(indices.iter().map(|i| *i as usize).collect()), }; - - let mut sort_information = vec![]; - for ordering in &scan.sort_information { - let sort_exprs = sort_exprs_try_from_proto( - &ordering.physical_sort_expr_nodes, - &ctx.expr_ctx(&schema), - )?; - sort_information.extend(LexOrdering::new(sort_exprs)); - } - - let fetch = scan - .fetch + let fetch = proto_fetch .map(|fetch| usize_from_wire(fetch, "MemoryScanExecNode", "fetch")) .transpose()?; - let source = Self::try_new(&partitions, schema, projection)? + let mut source = Self::try_new(&partitions, source_schema, projection)? .with_limit(fetch) - .with_show_sizes(scan.show_sizes) - .try_with_sort_information(sort_information)?; + .with_show_sizes(*show_sizes); + + // Stored sort information has already been projected by + // `try_with_sort_information`; decode it against that same schema and + // do not project it a second time. + let mut decoded_sort_information = vec![]; + for ordering in proto_sort_information { + let protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + } = ordering; + if let Some(ordering) = optional_ordering_try_from_proto( + physical_sort_expr_nodes, + &ctx.expr_ctx(&source.projected_schema), + )? { + decoded_sort_information.push(ordering); + } + } + Self::validate_sort_information( + &decoded_sort_information, + &source.projected_schema, + "projected", + )?; + source.sort_information = decoded_sort_information; Ok(DataSourceExec::from_data_source(source)) } diff --git a/datafusion/physical-expr-adapter/src/schema_rewriter.rs b/datafusion/physical-expr-adapter/src/schema_rewriter.rs index 2548ffc6fb1b7..b6364fcfe93cb 100644 --- a/datafusion/physical-expr-adapter/src/schema_rewriter.rs +++ b/datafusion/physical-expr-adapter/src/schema_rewriter.rs @@ -179,6 +179,13 @@ pub trait PhysicalExprAdapterFactory: Send + Sync + std::fmt::Debug { logical_file_schema: SchemaRef, physical_file_schema: SchemaRef, ) -> Result>; + + /// Whether replacing this factory with [`DefaultPhysicalExprAdapterFactory`] + /// preserves execution behavior. Plan serializers may safely omit factories + /// that return `true` because decoders use the default when none is configured. + fn is_equivalent_to_default(&self) -> bool { + false + } } #[derive(Debug, Clone)] @@ -195,6 +202,10 @@ impl PhysicalExprAdapterFactory for DefaultPhysicalExprAdapterFactory { physical_file_schema, })) } + + fn is_equivalent_to_default(&self) -> bool { + true + } } /// Default implementation of [`PhysicalExprAdapter`] for rewriting physical diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 4b63631613ae1..54c39a9267b5c 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1269,6 +1269,9 @@ message FileScanExecConf { // Compression used by formats such as CSV and JSON. Absent means uncompressed // for compatibility with payloads written before this field existed. optional datafusion_common.CompressionTypeVariant file_compression_type = 16; + // Whether file processing order must be preserved. Absent payloads retain the + // legacy behavior of deriving this from output_ordering. + optional bool preserve_order = 17; } message ParquetScanExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 9811357f1dd5d..2e6f5845f27d7 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -7133,6 +7133,9 @@ impl serde::Serialize for FileScanExecConf { if self.file_compression_type.is_some() { len += 1; } + if self.preserve_order.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.FileScanExecConf", len)?; if !self.file_groups.is_empty() { struct_ser.serialize_field("fileGroups", &self.file_groups)?; @@ -7177,6 +7180,9 @@ impl serde::Serialize for FileScanExecConf { .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", *v)))?; struct_ser.serialize_field("fileCompressionType", &v)?; } + if let Some(v) = self.preserve_order.as_ref() { + struct_ser.serialize_field("preserveOrder", v)?; + } struct_ser.end() } } @@ -7208,6 +7214,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "outputPartitioning", "file_compression_type", "fileCompressionType", + "preserve_order", + "preserveOrder", ]; #[allow(clippy::enum_variant_names)] @@ -7225,6 +7233,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { ProjectionExprs, OutputPartitioning, FileCompressionType, + PreserveOrder, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -7259,6 +7268,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "projectionExprs" | "projection_exprs" => Ok(GeneratedField::ProjectionExprs), "outputPartitioning" | "output_partitioning" => Ok(GeneratedField::OutputPartitioning), "fileCompressionType" | "file_compression_type" => Ok(GeneratedField::FileCompressionType), + "preserveOrder" | "preserve_order" => Ok(GeneratedField::PreserveOrder), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -7291,6 +7301,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { let mut projection_exprs__ = None; let mut output_partitioning__ = None; let mut file_compression_type__ = None; + let mut preserve_order__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::FileGroups => { @@ -7376,6 +7387,12 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { } file_compression_type__ = map_.next_value::<::std::option::Option>()?.map(|x| x as i32); } + GeneratedField::PreserveOrder => { + if preserve_order__.is_some() { + return Err(serde::de::Error::duplicate_field("preserveOrder")); + } + preserve_order__ = map_.next_value()?; + } } } Ok(FileScanExecConf { @@ -7392,6 +7409,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { projection_exprs: projection_exprs__, output_partitioning: output_partitioning__, file_compression_type: file_compression_type__, + preserve_order: preserve_order__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index a20632860a4c7..1a54515a78ba6 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1959,6 +1959,10 @@ pub struct FileScanExecConf { tag = "16" )] pub file_compression_type: ::core::option::Option, + /// Whether file processing order must be preserved. Absent payloads retain the + /// legacy behavior of deriving this from output_ordering. + #[prost(bool, optional, tag = "17")] + pub preserve_order: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetScanExecNode { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 26d1a8ed83d36..887337e291f43 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -105,7 +105,11 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { #[cfg(test)] mod file_scan_config_serde { use super::*; - use arrow::datatypes::{DataType, Field}; + use arrow::datatypes::{DataType, Field, SchemaRef}; + use datafusion::physical_expr_adapter::{ + DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, + PhysicalExprAdapterFactory, + }; use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_compression_type::FileCompressionType; @@ -146,6 +150,28 @@ mod file_scan_config_serde { } } + #[derive(Debug)] + struct SerdeTestExprAdapter; + + impl PhysicalExprAdapter for SerdeTestExprAdapter { + fn rewrite(&self, expr: Arc) -> Result> { + Ok(expr) + } + } + + #[derive(Debug)] + struct SerdeTestExprAdapterFactory; + + impl PhysicalExprAdapterFactory for SerdeTestExprAdapterFactory { + fn create( + &self, + _logical_file_schema: SchemaRef, + _physical_file_schema: SchemaRef, + ) -> Result> { + Ok(Arc::new(SerdeTestExprAdapter)) + } + } + impl FileSource for SerdeTestSource { fn create_file_opener( &self, @@ -374,6 +400,55 @@ mod file_scan_config_serde { Ok(()) } + #[test] + fn new_file_scan_config_serde_preserves_explicit_order_flag() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + let preserve_without_ordering = FileScanConfigBuilder::from(test_config(None)) + .with_output_ordering(vec![]) + .with_preserve_order(true) + .build(); + let encoded = serde.encode(&preserve_without_ordering)?; + assert_eq!(encoded.preserve_order, Some(true)); + let decoded = serde.decode(&encoded)?; + assert!(decoded.preserve_order); + assert!(decoded.output_ordering.is_empty()); + + let mut do_not_preserve_with_ordering = test_config(None); + do_not_preserve_with_ordering.preserve_order = false; + assert!(!do_not_preserve_with_ordering.output_ordering.is_empty()); + let mut encoded = serde.encode(&do_not_preserve_with_ordering)?; + assert_eq!(encoded.preserve_order, Some(false)); + assert!(!serde.decode(&encoded)?.preserve_order); + + // Older payloads had no flag and derived it from the ordering. + encoded.preserve_order = None; + assert!(serde.decode(&encoded)?.preserve_order); + Ok(()) + } + + #[test] + fn new_file_scan_config_encode_handles_expr_adapter_factories() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let default_config = FileScanConfigBuilder::from(test_config(None)) + .with_expr_adapter(Some(Arc::new(DefaultPhysicalExprAdapterFactory))) + .build(); + let encoded = serde.encode(&default_config)?; + assert!(serde.decode(&encoded)?.expr_adapter_factory.is_none()); + + let custom_config = FileScanConfigBuilder::from(test_config(None)) + .with_expr_adapter(Some(Arc::new(SerdeTestExprAdapterFactory))) + .build(); + let err = serde + .encode(&custom_config) + .expect_err("custom expression adapter must not be dropped"); + assert!( + err.to_string().contains("expr_adapter_factory"), + "unexpected error: {err}" + ); + Ok(()) + } + #[test] fn new_file_scan_config_decode_without_compression_uses_legacy_default() -> Result<()> { diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index c17ff47e0f472..967bef60c526e 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -882,7 +882,7 @@ async fn roundtrip_memory_source() -> Result<()> { } #[tokio::test] -async fn roundtrip_memory_source_sort_information_and_fetch() -> Result<()> { +async fn roundtrip_memory_source_projected_sort_information_and_fetch() -> Result<()> { use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSource as _; @@ -905,10 +905,11 @@ async fn roundtrip_memory_source_sort_information_and_fetch() -> Result<()> { }, )]) .unwrap(); - let source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None)? - .with_limit(Some(1)) - .with_show_sizes(false) - .try_with_sort_information(vec![ordering])?; + let source = + MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), Some(vec![1]))? + .with_limit(Some(1)) + .with_show_sizes(false) + .try_with_sort_information(vec![ordering])?; let exec_plan = DataSourceExec::from_data_source(source.clone()); let ctx = SessionContext::new();