zhuqi-lucas commented on code in PR #18817: URL: https://github.com/apache/datafusion/pull/18817#discussion_r2554026346
########## datafusion/physical-optimizer/src/reverse_order.rs: ########## @@ -0,0 +1,1446 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::PhysicalOptimizerRule; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, +}; +use datafusion_common::Result; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::source::{DataSource, DataSourceExec}; +use datafusion_datasource_parquet::source::ParquetSource; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_plan::joins::SortMergeJoinExec; +use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use std::sync::Arc; + +/// A PhysicalOptimizerRule that attempts to reverse a SortExec's input if doing so would make the +/// input ordering compatible with the SortExec's required output ordering. +/// It also removes unnecessary sorts when the input already satisfies the required ordering. +#[derive(Debug, Clone, Default)] +pub struct ReverseOrder; + +impl ReverseOrder { + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for ReverseOrder { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + // Check if reverse scan optimization is enabled + let enable_reverse_scan = config.execution.parquet.enable_reverse_scan; + + // Return early if not enabled + if !enable_reverse_scan { + return Ok(plan); + } + + // Search for any SortExec nodes and try to optimize them + plan.transform_up(&|plan: Arc<dyn ExecutionPlan>| { + // First check if this is a GlobalLimitExec -> SortExec pattern + if let Some(limit_exec) = plan.as_any().downcast_ref::<GlobalLimitExec>() { + if let Some(sort_exec) = + limit_exec.input().as_any().downcast_ref::<SortExec>() + { + return optimize_limit_sort(limit_exec, sort_exec); + } + } + + // Otherwise, check if this is just a SortExec + let sort_exec = match plan.as_any().downcast_ref::<SortExec>() { + Some(sort_exec) => sort_exec, + None => return Ok(Transformed::no(plan)), + }; + + let sort_input: Arc<dyn ExecutionPlan> = Arc::clone(sort_exec.input()); + + // First, check if the sort is already satisfied by input ordering + if let Some(_input_ordering) = sort_input.output_ordering() { + let input_eq_properties = sort_input.equivalence_properties(); + + // Check if input already satisfies the sort requirement + if input_eq_properties.ordering_satisfy(sort_exec.expr().clone())? { + return remove_unnecessary_sort(sort_exec, sort_input); + } + } + + // If not satisfied, try to reverse the input + let reversed_eq_properties = { + let mut new = sort_input.properties().equivalence_properties().clone(); + new.clear_orderings(); + + // Build reversed sort exprs for each ordering class + let reversed_orderings = sort_input + .equivalence_properties() + .oeq_class() + .iter() + .map(|ordering| { + ordering + .iter() + .map(|expr| expr.reverse()) + .collect::<Vec<_>>() + }) + .collect::<Vec<_>>(); + + new.add_orderings(reversed_orderings); + new + }; + + match reversed_eq_properties.ordering_satisfy(sort_exec.expr().clone())? { + true => { + // Reverse the input and then remove the sort + let reversed_input = + sort_input.rewrite(&mut ReverseRewriter).unwrap().data; + + // After reversing, check if we can remove the sort + if reversed_input + .equivalence_properties() + .ordering_satisfy(sort_exec.expr().clone())? + { + remove_unnecessary_sort(sort_exec, reversed_input) + } else { + // Keep the sort but with reversed input + Ok(Transformed::yes(Arc::new( + SortExec::new(sort_exec.expr().clone(), reversed_input) + .with_fetch(sort_exec.fetch()) + .with_preserve_partitioning( + sort_exec.preserve_partitioning(), + ), + ))) + } + } + false => Ok(Transformed::no(plan)), + } + }) + .data() + } + + fn name(&self) -> &str { + "ReverseOrder" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Handle the GlobalLimitExec -> SortExec pattern +fn optimize_limit_sort( + limit_exec: &GlobalLimitExec, + sort_exec: &SortExec, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let sort_input = Arc::clone(sort_exec.input()); + + // Check if input is already sorted + if let Some(_input_ordering) = sort_input.output_ordering() { + let input_eq_properties = sort_input.equivalence_properties(); + if input_eq_properties.ordering_satisfy(sort_exec.expr().clone())? { + // Input is already sorted correctly, remove sort and keep limit + return Ok(Transformed::yes(Arc::new(GlobalLimitExec::new( + sort_input, + limit_exec.skip(), + limit_exec.fetch(), + )))); + } + } + + // Check if we can reverse the input to satisfy the sort + let reversed_eq_properties = { + let mut new = sort_input.properties().equivalence_properties().clone(); + new.clear_orderings(); + + let reversed_orderings = sort_input + .equivalence_properties() + .oeq_class() + .iter() + .map(|ordering| { + ordering + .iter() + .map(|expr| expr.reverse()) + .collect::<Vec<_>>() + }) + .collect::<Vec<_>>(); + + new.add_orderings(reversed_orderings); + new + }; + + if reversed_eq_properties.ordering_satisfy(sort_exec.expr().clone())? { + // Can reverse! Apply reversal + let reversed_input = sort_input.rewrite(&mut ReverseRewriter).unwrap().data; + + // Check if reversed input satisfies the sort requirement + if reversed_input + .equivalence_properties() + .ordering_satisfy(sort_exec.expr().clone())? + { + // Check if this is a single-partition DataSourceExec with reverse_scan enabled + // In that case, the limit is already handled internally by ReversedParquetStreamWithLimit + if is_single_partition_reverse_scan_datasource(&reversed_input) { + let total_fetch = limit_exec.skip() + limit_exec.fetch().unwrap_or(0); + + if let Some(with_fetch) = reversed_input.with_fetch(Some(total_fetch)) { + if limit_exec.skip() > 0 { + return Ok(Transformed::yes(Arc::new(GlobalLimitExec::new( + with_fetch, + limit_exec.skip(), + limit_exec.fetch(), + )))); + } else { + return Ok(Transformed::yes(with_fetch)); + } + } + + return Ok(Transformed::yes(Arc::new(GlobalLimitExec::new( + reversed_input, + limit_exec.skip(), + limit_exec.fetch(), + )))); + } + + // Otherwise, remove sort but keep limit with reversed input + return Ok(Transformed::yes(Arc::new(GlobalLimitExec::new( + reversed_input, + limit_exec.skip(), + limit_exec.fetch(), + )))); + } + } + + // Can't optimize, return original pattern + Ok(Transformed::no(Arc::new(GlobalLimitExec::new( + Arc::new(sort_exec.clone()), + limit_exec.skip(), + limit_exec.fetch(), + )))) +} + +/// Check if the plan is a single-partition DataSourceExec with reverse_scan enabled +fn is_single_partition_reverse_scan_datasource(plan: &Arc<dyn ExecutionPlan>) -> bool { + // Only optimize for single partition + if plan.output_partitioning().partition_count() != 1 { + return false; + } + + if let Some(data_source_exec) = plan.as_any().downcast_ref::<DataSourceExec>() { + if let Some(scan_config) = data_source_exec + .data_source() + .as_any() + .downcast_ref::<FileScanConfig>() + { + if let Some(parquet_source) = scan_config + .file_source + .as_any() + .downcast_ref::<ParquetSource>() + { + return parquet_source.reverse_scan(); + } + } + } + false +} + +/// Remove unnecessary sort based on the logic from EnforceSorting::analyze_immediate_sort_removal +fn remove_unnecessary_sort( Review Comment: Note, i add this to reverse order because after reverse order, we can optimize more to remove the sort, so we don't need to execute the enforce sort optimization again after this optimization. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
