Merge branch 'main' into case_reduce_filtering

Author: pepijnve

.github/workflows/audit.yml

@@ -42,7 +42,7 @@ jobs:
     steps:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
       - name: Install cargo-audit
-        uses: taiki-e/install-action@e43a5023a747770bfcb71ae048541a681714b951  # v2.62.33
+        uses: taiki-e/install-action@80466ef8efa80486cdfbddf929453a4f3565c791  # v2.62.34
         with:
           tool: cargo-audit
       - name: Run audit check

.github/workflows/rust.yml

@@ -425,7 +425,7 @@ jobs:
           sudo apt-get update -qq
           sudo apt-get install -y -qq clang
       - name: Setup wasm-pack
-        uses: taiki-e/install-action@e43a5023a747770bfcb71ae048541a681714b951  # v2.62.33
+        uses: taiki-e/install-action@80466ef8efa80486cdfbddf929453a4f3565c791  # v2.62.34
         with:
           tool: wasm-pack
       - name: Run tests with headless mode
@@ -752,7 +752,7 @@ jobs:
       - name: Setup Rust toolchain
         uses: ./.github/actions/setup-builder
       - name: Install cargo-msrv
-        uses: taiki-e/install-action@e43a5023a747770bfcb71ae048541a681714b951  # v2.62.33
+        uses: taiki-e/install-action@80466ef8efa80486cdfbddf929453a4f3565c791  # v2.62.34
         with:
           tool: cargo-msrv
 

datafusion/common/src/config.rs

@@ -938,6 +938,11 @@ config_namespace! {
         /// HashJoin can work more efficiently than SortMergeJoin but consumes more memory
         pub prefer_hash_join: bool, default = true
 
+        /// When set to true, piecewise merge join is enabled. PiecewiseMergeJoin is currently
+        /// experimental. Physical planner will opt for PiecewiseMergeJoin when there is only
+        /// one range filter.
+        pub enable_piecewise_merge_join: bool, default = false
+
         /// The maximum estimated size in bytes for one input side of a HashJoin
         /// will be collected into a single partition
         pub hash_join_single_partition_threshold: usize, default = 1024 * 1024

datafusion/core/src/physical_planner.rs

@@ -78,10 +78,11 @@ use datafusion_expr::expr::{
 };
 use datafusion_expr::expr_rewriter::unnormalize_cols;
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
+use datafusion_expr::utils::split_conjunction;
 use datafusion_expr::{
-    Analyze, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, FetchType,
-    Filter, JoinType, RecursiveQuery, SkipType, StringifiedPlan, WindowFrame,
-    WindowFrameBound, WriteOp,
+    Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension,
+    FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan,
+    WindowFrame, WindowFrameBound, WriteOp,
 };
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
 use datafusion_physical_expr::expressions::Literal;
@@ -91,6 +92,7 @@ use datafusion_physical_expr::{
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::execution_plan::InvariantLevel;
+use datafusion_physical_plan::joins::PiecewiseMergeJoinExec;
 use datafusion_physical_plan::metrics::MetricType;
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion_physical_plan::recursive_query::RecursiveQueryExec;
@@ -1133,8 +1135,42 @@ impl DefaultPhysicalPlanner {
                     })
                     .collect::<Result<join_utils::JoinOn>>()?;
 
+                // TODO: `num_range_filters` can be used later on for ASOF joins (`num_range_filters > 1`)
+                let mut num_range_filters = 0;
+                let mut range_filters: Vec<Expr> = Vec::new();
+                let mut total_filters = 0;
+
                 let join_filter = match filter {
                     Some(expr) => {
+                        let split_expr = split_conjunction(expr);
+                        for expr in split_expr.iter() {
+                            match *expr {
+                                Expr::BinaryExpr(BinaryExpr {
+                                    left: _,
+                                    right: _,
+                                    op,
+                                }) => {
+                                    if matches!(
+                                        op,
+                                        Operator::Lt
+                                            | Operator::LtEq
+                                            | Operator::Gt
+                                            | Operator::GtEq
+                                    ) {
+                                        range_filters.push((**expr).clone());
+                                        num_range_filters += 1;
+                                    }
+                                    total_filters += 1;
+                                }
+                                // TODO: Want to deal with `Expr::Between` for IEJoins, it counts as two range predicates
+                                // which is why it is not dealt with in PWMJ
+                                // Expr::Between(_) => {},
+                                _ => {
+                                    total_filters += 1;
+                                }
+                            }
+                        }
+
                         // Extract columns from filter expression and saved in a HashSet
                         let cols = expr.column_refs();
 
@@ -1190,6 +1226,7 @@ impl DefaultPhysicalPlanner {
                         )?;
                         let filter_schema =
                             Schema::new_with_metadata(filter_fields, metadata);
+
                         let filter_expr = create_physical_expr(
                             expr,
                             &filter_df_schema,
@@ -1212,10 +1249,125 @@ impl DefaultPhysicalPlanner {
                 let prefer_hash_join =
                     session_state.config_options().optimizer.prefer_hash_join;
 
+                // TODO: Allow PWMJ to deal with residual equijoin conditions
                 let join: Arc<dyn ExecutionPlan> = if join_on.is_empty() {
                     if join_filter.is_none() && matches!(join_type, JoinType::Inner) {
                         // cross join if there is no join conditions and no join filter set
                         Arc::new(CrossJoinExec::new(physical_left, physical_right))
+                    } else if num_range_filters == 1
+                        && total_filters == 1
+                        && !matches!(
+                            join_type,
+                            JoinType::LeftSemi
+                                | JoinType::RightSemi
+                                | JoinType::LeftAnti
+                                | JoinType::RightAnti
+                                | JoinType::LeftMark
+                                | JoinType::RightMark
+                        )
+                        && session_state
+                            .config_options()
+                            .optimizer
+                            .enable_piecewise_merge_join
+                    {
+                        let Expr::BinaryExpr(be) = &range_filters[0] else {
+                            return plan_err!(
+                                "Unsupported expression for PWMJ: Expected `Expr::BinaryExpr`"
+                            );
+                        };
+
+                        let mut op = be.op;
+                        if !matches!(
+                            op,
+                            Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq
+                        ) {
+                            return plan_err!(
+                                "Unsupported operator for PWMJ: {:?}. Expected one of <, <=, >, >=",
+                                op
+                            );
+                        }
+
+                        fn reverse_ineq(op: Operator) -> Operator {
+                            match op {
+                                Operator::Lt => Operator::Gt,
+                                Operator::LtEq => Operator::GtEq,
+                                Operator::Gt => Operator::Lt,
+                                Operator::GtEq => Operator::LtEq,
+                                _ => op,
+                            }
+                        }
+
+                        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+                        enum Side {
+                            Left,
+                            Right,
+                            Both,
+                        }
+
+                        let side_of = |e: &Expr| -> Result<Side> {
+                            let cols = e.column_refs();
+                            let any_left = cols
+                                .iter()
+                                .any(|c| left_df_schema.index_of_column(c).is_ok());
+                            let any_right = cols
+                                .iter()
+                                .any(|c| right_df_schema.index_of_column(c).is_ok());
+
+                            Ok(match (any_left, any_right) {
+                                (true, false) => Side::Left,
+                                (false, true) => Side::Right,
+                                (true, true) => Side::Both,
+                                _ => unreachable!(),
+                            })
+                        };
+
+                        let mut lhs_logical = &be.left;
+                        let mut rhs_logical = &be.right;
+
+                        let left_side = side_of(lhs_logical)?;
+                        let right_side = side_of(rhs_logical)?;
+                        if matches!(left_side, Side::Both)
+                            || matches!(right_side, Side::Both)
+                        {
+                            return Ok(Arc::new(NestedLoopJoinExec::try_new(
+                                physical_left,
+                                physical_right,
+                                join_filter,
+                                join_type,
+                                None,
+                            )?));
+                        }
+
+                        if left_side == Side::Right && right_side == Side::Left {
+                            std::mem::swap(&mut lhs_logical, &mut rhs_logical);
+                            op = reverse_ineq(op);
+                        } else if !(left_side == Side::Left && right_side == Side::Right)
+                        {
+                            return plan_err!(
+                                "Unsupported operator for PWMJ: {:?}. Expected one of <, <=, >, >=",
+                                op
+                            );
+                        }
+
+                        let on_left = create_physical_expr(
+                            lhs_logical,
+                            left_df_schema,
+                            session_state.execution_props(),
+                        )?;
+                        let on_right = create_physical_expr(
+                            rhs_logical,
+                            right_df_schema,
+                            session_state.execution_props(),
+                        )?;
+
+                        Arc::new(PiecewiseMergeJoinExec::try_new(
+                            physical_left,
+                            physical_right,
+                            (on_left, on_right),
+                            op,
+                            *join_type,
+                            session_state.config().target_partitions(),
+                        )?)
                     } else {
                         // there is no equal join condition, use the nested loop join
                         Arc::new(NestedLoopJoinExec::try_new(

datafusion/functions/src/datetime/current_time.rs

@@ -15,24 +15,28 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::array::timezone::Tz;
 use arrow::datatypes::DataType;
 use arrow::datatypes::DataType::Time64;
 use arrow::datatypes::TimeUnit::Nanosecond;
-use std::any::Any;
-
+use chrono::TimeZone;
+use chrono::Timelike;
 use datafusion_common::{internal_err, Result, ScalarValue};
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 use datafusion_expr::{
     ColumnarValue, Documentation, Expr, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
+use std::any::Any;
 
 #[user_doc(
     doc_section(label = "Time and Date Functions"),
     description = r#"
-Returns the current UTC time.
+Returns the current time in the session time zone.
 
 The `current_time()` return value is determined at query time and will return the same time, no matter when in the query plan the function executes.
+
+The session time zone can be set using the statement 'SET datafusion.execution.time_zone = desired time zone'. The time zone can be a value like +00:00, 'Europe/London' etc.
 "#,
     syntax_example = "current_time()"
 )]
@@ -93,7 +97,20 @@ impl ScalarUDFImpl for CurrentTimeFunc {
         info: &dyn SimplifyInfo,
     ) -> Result<ExprSimplifyResult> {
         let now_ts = info.execution_props().query_execution_start_time;
-        let nano = now_ts.timestamp_nanos_opt().map(|ts| ts % 86400000000000);
+
+        // Try to get timezone from config and convert to local time
+        let nano = info
+            .execution_props()
+            .config_options()
+            .and_then(|config| config.execution.time_zone.parse::<Tz>().ok())
+            .map_or_else(
+                || datetime_to_time_nanos(&now_ts),
+                |tz| {
+                    let local_now = tz.from_utc_datetime(&now_ts.naive_utc());
+                    datetime_to_time_nanos(&local_now)
+                },
+            );
+
         Ok(ExprSimplifyResult::Simplified(Expr::Literal(
             ScalarValue::Time64Nanosecond(nano),
             None,
@@ -104,3 +121,97 @@ impl ScalarUDFImpl for CurrentTimeFunc {
         self.doc()
     }
 }
+
+// Helper function for conversion of datetime to a timestamp.
+fn datetime_to_time_nanos<Tz: TimeZone>(dt: &chrono::DateTime<Tz>) -> Option<i64> {
+    let hour = dt.hour() as i64;
+    let minute = dt.minute() as i64;
+    let second = dt.second() as i64;
+    let nanosecond = dt.nanosecond() as i64;
+    Some((hour * 3600 + minute * 60 + second) * 1_000_000_000 + nanosecond)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::datatypes::{DataType, TimeUnit::Nanosecond};
+    use chrono::{DateTime, Utc};
+    use datafusion_common::{Result, ScalarValue};
+    use datafusion_expr::execution_props::ExecutionProps;
+    use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
+    use std::sync::Arc;
+
+    struct MockSimplifyInfo {
+        execution_props: ExecutionProps,
+    }
+
+    impl SimplifyInfo for MockSimplifyInfo {
+        fn is_boolean_type(&self, _expr: &Expr) -> Result<bool> {
+            Ok(false)
+        }
+
+        fn nullable(&self, _expr: &Expr) -> Result<bool> {
+            Ok(true)
+        }
+
+        fn execution_props(&self) -> &ExecutionProps {
+            &self.execution_props
+        }
+
+        fn get_data_type(&self, _expr: &Expr) -> Result<DataType> {
+            Ok(Time64(Nanosecond))
+        }
+    }
+
+    fn set_session_timezone_env(tz: &str, start_time: DateTime<Utc>) -> MockSimplifyInfo {
+        let mut config = datafusion_common::config::ConfigOptions::default();
+        config.execution.time_zone = tz.to_string();
+        let mut execution_props =
+            ExecutionProps::new().with_query_execution_start_time(start_time);
+        execution_props.config_options = Some(Arc::new(config));
+        MockSimplifyInfo { execution_props }
+    }
+
+    #[test]
+    fn test_current_time_timezone_offset() {
+        // Use a fixed start time for consistent testing
+        let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 12, 0, 0).unwrap();
+
+        // Test with UTC+05:00
+        let info_plus_5 = set_session_timezone_env("+05:00", start_time);
+        let result_plus_5 = CurrentTimeFunc::new()
+            .simplify(vec![], &info_plus_5)
+            .unwrap();
+
+        // Test with UTC-05:00
+        let info_minus_5 = set_session_timezone_env("-05:00", start_time);
+        let result_minus_5 = CurrentTimeFunc::new()
+            .simplify(vec![], &info_minus_5)
+            .unwrap();
+
+        // Extract nanoseconds from results
+        let nanos_plus_5 = match result_plus_5 {
+            ExprSimplifyResult::Simplified(Expr::Literal(
+                ScalarValue::Time64Nanosecond(Some(n)),
+                _,
+            )) => n,
+            _ => panic!("Expected Time64Nanosecond literal"),
+        };
+
+        let nanos_minus_5 = match result_minus_5 {
+            ExprSimplifyResult::Simplified(Expr::Literal(
+                ScalarValue::Time64Nanosecond(Some(n)),
+                _,
+            )) => n,
+            _ => panic!("Expected Time64Nanosecond literal"),
+        };
+
+        // Calculate the difference: UTC+05:00 should be 10 hours ahead of UTC-05:00
+        let difference = nanos_plus_5 - nanos_minus_5;
+
+        // 10 hours in nanoseconds
+        let expected_offset = 10i64 * 3600 * 1_000_000_000;
+
+        assert_eq!(difference, expected_offset, "Expected 10-hour offset difference in nanoseconds between UTC+05:00 and UTC-05:00");
+    }
+}

datafusion/physical-plan/src/joins/hash_join/stream.rs

@@ -637,6 +637,7 @@ impl HashJoinStream {
         let (left_side, right_side) = get_final_indices_from_shared_bitmap(
             build_side.left_data.visited_indices_bitmap(),
             self.join_type,
+            true,
         );
         let empty_right_batch = RecordBatch::new_empty(self.right.schema());
         // use the left and right indices to produce the batch result