diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 087da503654d7..4872db9fd3329 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -54,6 +54,10 @@ use datafusion_physical_expr_adapter::rewrite::{ }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::physical_expr::fmt_sql; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::sort_expr::{ + optional_ordering_try_from_proto, sort_exprs_try_to_proto, +}; use datafusion_physical_plan::DisplayFormatType; use datafusion_physical_plan::SortOrderPushdownResult; use datafusion_physical_plan::filter_pushdown::PushedDown; @@ -1089,11 +1093,25 @@ impl FileSource for ParquetSource { .filter() .map(|pred| ctx.encode_expr(&pred)) .transpose()?; + let sort_order_for_reorder = self + .sort_order_for_reorder + .as_ref() + .map(|ordering| -> datafusion_common::Result<_> { + Ok(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes: sort_exprs_try_to_proto( + ordering.iter(), + &ctx.expr_ctx(), + )?, + }) + }) + .transpose()?; let node = protobuf::ParquetScanExecNode { base_conf: Some(base.try_to_proto(ctx)?), predicate, parquet_options: Some(self.table_parquet_options().try_into()?), + sort_order_for_reorder, + reverse_row_groups: self.reverse_row_groups, }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::ParquetScan(node)), @@ -1163,6 +1181,17 @@ impl ParquetSource { .as_ref() .map(|expr| ctx.decode_expr(expr, predicate_schema.as_ref())) .transpose()?; + let sort_order_for_reorder = scan + .sort_order_for_reorder + .as_ref() + .map(|ordering| { + optional_ordering_try_from_proto( + &ordering.physical_sort_expr_nodes, + &ctx.expr_ctx(predicate_schema.as_ref()), + ) + }) + .transpose()? + .flatten(); let mut options = TableParquetOptions::default(); if let Some(table_options) = scan.parquet_options.as_ref() { @@ -1189,6 +1218,8 @@ impl ParquetSource { let mut source = ParquetSource::new(table_schema) .with_parquet_file_reader_factory(reader_factory) .with_table_parquet_options(options); + source.sort_order_for_reorder = sort_order_for_reorder; + source.reverse_row_groups = scan.reverse_row_groups; if let Some(predicate) = predicate { source = source.with_predicate(predicate); diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index e685b947bbaaf..8486410d07a2b 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1280,6 +1280,10 @@ message ParquetScanExecNode { PhysicalExprNode predicate = 3; datafusion_common.TableParquetOptions parquet_options = 4; + + PhysicalSortExprNodeCollection sort_order_for_reorder = 5; + + bool reverse_row_groups = 6; } message CsvScanExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 4fcca074c8199..f568a6d853d56 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -16225,6 +16225,12 @@ impl serde::Serialize for ParquetScanExecNode { if self.parquet_options.is_some() { len += 1; } + if self.sort_order_for_reorder.is_some() { + len += 1; + } + if self.reverse_row_groups { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ParquetScanExecNode", len)?; if let Some(v) = self.base_conf.as_ref() { struct_ser.serialize_field("baseConf", v)?; @@ -16235,6 +16241,12 @@ impl serde::Serialize for ParquetScanExecNode { if let Some(v) = self.parquet_options.as_ref() { struct_ser.serialize_field("parquetOptions", v)?; } + if let Some(v) = self.sort_order_for_reorder.as_ref() { + struct_ser.serialize_field("sortOrderForReorder", v)?; + } + if self.reverse_row_groups { + struct_ser.serialize_field("reverseRowGroups", &self.reverse_row_groups)?; + } struct_ser.end() } } @@ -16250,6 +16262,10 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { "predicate", "parquet_options", "parquetOptions", + "sort_order_for_reorder", + "sortOrderForReorder", + "reverse_row_groups", + "reverseRowGroups", ]; #[allow(clippy::enum_variant_names)] @@ -16257,6 +16273,8 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { BaseConf, Predicate, ParquetOptions, + SortOrderForReorder, + ReverseRowGroups, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -16281,6 +16299,8 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { "baseConf" | "base_conf" => Ok(GeneratedField::BaseConf), "predicate" => Ok(GeneratedField::Predicate), "parquetOptions" | "parquet_options" => Ok(GeneratedField::ParquetOptions), + "sortOrderForReorder" | "sort_order_for_reorder" => Ok(GeneratedField::SortOrderForReorder), + "reverseRowGroups" | "reverse_row_groups" => Ok(GeneratedField::ReverseRowGroups), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -16303,6 +16323,8 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { let mut base_conf__ = None; let mut predicate__ = None; let mut parquet_options__ = None; + let mut sort_order_for_reorder__ = None; + let mut reverse_row_groups__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::BaseConf => { @@ -16323,12 +16345,26 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { } parquet_options__ = map_.next_value()?; } + GeneratedField::SortOrderForReorder => { + if sort_order_for_reorder__.is_some() { + return Err(serde::de::Error::duplicate_field("sortOrderForReorder")); + } + sort_order_for_reorder__ = map_.next_value()?; + } + GeneratedField::ReverseRowGroups => { + if reverse_row_groups__.is_some() { + return Err(serde::de::Error::duplicate_field("reverseRowGroups")); + } + reverse_row_groups__ = Some(map_.next_value()?); + } } } Ok(ParquetScanExecNode { base_conf: base_conf__, predicate: predicate__, parquet_options: parquet_options__, + sort_order_for_reorder: sort_order_for_reorder__, + reverse_row_groups: reverse_row_groups__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index d14195ecc54a0..f5c7c01e7192f 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1970,6 +1970,10 @@ pub struct ParquetScanExecNode { pub parquet_options: ::core::option::Option< super::datafusion_common::TableParquetOptions, >, + #[prost(message, optional, tag = "5")] + pub sort_order_for_reorder: ::core::option::Option, + #[prost(bool, tag = "6")] + pub reverse_row_groups: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct CsvScanExecNode { diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index c17ff47e0f472..39028ac1ca988 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -112,6 +112,32 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { roundtrip_test(DataSourceExec::from_data_source(scan_config)) } +#[tokio::test] +async fn roundtrip_parquet_exec_with_sort_pushdown() -> Result<()> { + let ctx = all_types_context().await?; + let plan = ctx + .sql("SELECT id FROM alltypes_plain ORDER BY id DESC NULLS LAST LIMIT 5") + .await? + .create_physical_plan() + .await?; + let before = displayable(plan.as_ref()).indent(true).to_string(); + assert!( + before.contains("sort_order_for_reorder=[id@0 DESC NULLS LAST]") + && before.contains("reverse_row_groups=true"), + "expected sort pushdown in plan:\n{before}" + ); + + let roundtripped = roundtrip_test_and_return( + plan, + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let after = displayable(roundtripped.as_ref()).indent(true).to_string(); + pretty_assertions::assert_eq!(before, after); + Ok(()) +} + #[test] fn file_scan_rejects_zero_batch_size() -> Result<()> { let schema = Arc::new(Schema::empty());