fix: preserve Parquet sort pushdown across proto roundtrips

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);

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 {

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,13 +16262,19 @@ 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)]
         enum GeneratedField {
             BaseConf,
             Predicate,
             ParquetOptions,
+            SortOrderForReorder,
+            ReverseRowGroups,
         }
         impl<'de> serde::Deserialize<'de> for GeneratedField {
             fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -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(),
                 })
             }
         }

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<PhysicalSortExprNodeCollection>,
+    #[prost(bool, tag = "6")]
+    pub reverse_row_groups: bool,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct CsvScanExecNode {

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());