Merge remote-tracking branch 'apache/master' into file-partitioning

Author: alamb

.github/workflows/python_build.yml

@@ -106,6 +106,7 @@ jobs:
       - run: cat LICENSE.txt
       - name: Build wheels
         run: |
+          export RUSTFLAGS='-C target-cpu=skylake'
           docker run --rm -v $(pwd)/..:/io \
             --workdir /io/python \
             konstin2/maturin:v0.11.2 \

Cargo.toml

@@ -29,3 +29,7 @@ members = [
 ]
 
 exclude = ["python"]
+
+[profile.release]
+lto = true
+codegen-units = 1

datafusion/src/execution/context.rs

@@ -893,15 +893,6 @@ impl ExecutionConfig {
         Default::default()
     }
 
-    /// Deprecated. Use with_target_partitions instead.
-    #[deprecated(
-        since = "5.1.0",
-        note = "This method is deprecated in favor of `with_target_partitions`."
-    )]
-    pub fn with_concurrency(self, n: usize) -> Self {
-        self.with_target_partitions(n)
-    }
-
     /// Customize target_partitions
     pub fn with_target_partitions(mut self, n: usize) -> Self {
         // partition count must be greater than zero

datafusion/src/field_util.rs

@@ -0,0 +1,50 @@
+// 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.
+
+//! Utility functions for complex field access
+
+use arrow::datatypes::{DataType, Field};
+
+use crate::error::{DataFusionError, Result};
+use crate::scalar::ScalarValue;
+
+/// Returns the field access indexed by `key` from a [`DataType::List`]
+/// # Error
+/// Errors if
+/// * the `data_type` is not a Struct or,
+/// * there is no field key is not of the required index type
+pub fn get_indexed_field(data_type: &DataType, key: &ScalarValue) -> Result<Field> {
+    match (data_type, key) {
+        (DataType::List(lt), ScalarValue::Int64(Some(i))) => {
+            if *i < 0 {
+                Err(DataFusionError::Plan(format!(
+                    "List based indexed access requires a positive int, was {0}",
+                    i
+                )))
+            } else {
+                Ok(Field::new(&i.to_string(), lt.data_type().clone(), false))
+            }
+        }
+        (DataType::List(_), _) => Err(DataFusionError::Plan(
+            "Only ints are valid as an indexed field in a list".to_string(),
+        )),
+        _ => Err(DataFusionError::Plan(
+            "The expression to get an indexed field is only valid for `List` types"
+                .to_string(),
+        )),
+    }
+}

datafusion/src/lib.rs

@@ -231,6 +231,7 @@ pub mod variable;
 pub use arrow;
 pub use parquet;
 
+pub(crate) mod field_util;
 #[cfg(test)]
 pub mod test;
 pub mod test_util;

datafusion/src/logical_plan/expr.rs

@@ -20,6 +20,7 @@
 
 pub use super::Operator;
 use crate::error::{DataFusionError, Result};
+use crate::field_util::get_indexed_field;
 use crate::logical_plan::{window_frames, DFField, DFSchema, LogicalPlan};
 use crate::physical_plan::functions::Volatility;
 use crate::physical_plan::{
@@ -245,6 +246,13 @@ pub enum Expr {
     IsNull(Box<Expr>),
     /// arithmetic negation of an expression, the operand must be of a signed numeric data type
     Negative(Box<Expr>),
+    /// Returns the field of a [`ListArray`] by key
+    GetIndexedField {
+        /// the expression to take the field from
+        expr: Box<Expr>,
+        /// The name of the field to take
+        key: ScalarValue,
+    },
     /// Whether an expression is between a given range.
     Between {
         /// The value to compare
@@ -433,6 +441,11 @@ impl Expr {
             Expr::Wildcard => Err(DataFusionError::Internal(
                 "Wildcard expressions are not valid in a logical query plan".to_owned(),
             )),
+            Expr::GetIndexedField { ref expr, key } => {
+                let data_type = expr.get_type(schema)?;
+
+                get_indexed_field(&data_type, key).map(|x| x.data_type().clone())
+            }
         }
     }
 
@@ -488,6 +501,10 @@ impl Expr {
             Expr::Wildcard => Err(DataFusionError::Internal(
                 "Wildcard expressions are not valid in a logical query plan".to_owned(),
             )),
+            Expr::GetIndexedField { ref expr, key } => {
+                let data_type = expr.get_type(input_schema)?;
+                get_indexed_field(&data_type, key).map(|x| x.is_nullable())
+            }
         }
     }
 
@@ -763,6 +780,7 @@ impl Expr {
                     .try_fold(visitor, |visitor, arg| arg.accept(visitor))
             }
             Expr::Wildcard => Ok(visitor),
+            Expr::GetIndexedField { ref expr, .. } => expr.accept(visitor),
         }?;
 
         visitor.post_visit(self)
@@ -923,6 +941,10 @@ impl Expr {
                 negated,
             },
             Expr::Wildcard => Expr::Wildcard,
+            Expr::GetIndexedField { expr, key } => Expr::GetIndexedField {
+                expr: rewrite_boxed(expr, rewriter)?,
+                key,
+            },
         };
 
         // now rewrite this expression itself
@@ -1799,6 +1821,9 @@ impl fmt::Debug for Expr {
                 }
             }
             Expr::Wildcard => write!(f, "*"),
+            Expr::GetIndexedField { ref expr, key } => {
+                write!(f, "({:?})[{}]", expr, key)
+            }
         }
     }
 }
@@ -1879,6 +1904,10 @@ fn create_name(e: &Expr, input_schema: &DFSchema) -> Result<String> {
             let expr = create_name(expr, input_schema)?;
             Ok(format!("{} IS NOT NULL", expr))
         }
+        Expr::GetIndexedField { expr, key } => {
+            let expr = create_name(expr, input_schema)?;
+            Ok(format!("{}[{}]", expr, key))
+        }
         Expr::ScalarFunction { fun, args, .. } => {
             create_function_name(&fun.to_string(), false, args, input_schema)
         }

datafusion/src/optimizer/common_subexpr_eliminate.rs

@@ -442,6 +442,10 @@ impl ExprIdentifierVisitor<'_> {
             Expr::Wildcard => {
                 desc.push_str("Wildcard-");
             }
+            Expr::GetIndexedField { key, .. } => {
+                desc.push_str("GetIndexedField-");
+                desc.push_str(&key.to_string());
+            }
         }
 
         desc

datafusion/src/optimizer/constant_folding.rs

@@ -17,16 +17,13 @@
 
 //! Constant folding and algebraic simplification
 
-use std::sync::Arc;
-
 use arrow::datatypes::DataType;
 
 use crate::error::Result;
 use crate::execution::context::ExecutionProps;
 use crate::logical_plan::{DFSchemaRef, Expr, ExprRewriter, LogicalPlan, Operator};
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::utils;
-use crate::physical_plan::functions::BuiltinScalarFunction;
 use crate::scalar::ScalarValue;
 
 /// Simplifies plans by rewriting [`Expr`]`s evaluating constants
@@ -61,18 +58,14 @@ impl OptimizerRule for ConstantFolding {
         // children plans.
         let mut simplifier = Simplifier {
             schemas: plan.all_schemas(),
-            execution_props,
         };
 
-        let mut const_evaluator = utils::ConstEvaluator::new();
+        let mut const_evaluator = utils::ConstEvaluator::new(execution_props);
 
         match plan {
-            LogicalPlan::Filter { predicate, input } => Ok(LogicalPlan::Filter {
-                predicate: predicate.clone().rewrite(&mut simplifier)?,
-                input: Arc::new(self.optimize(input, execution_props)?),
-            }),
-            // Rest: recurse into plan, apply optimization where possible
-            LogicalPlan::Projection { .. }
+            // Recurse into plan, apply optimization where possible
+            LogicalPlan::Filter { .. }
+            | LogicalPlan::Projection { .. }
             | LogicalPlan::Window { .. }
             | LogicalPlan::Aggregate { .. }
             | LogicalPlan::Repartition { .. }
@@ -130,7 +123,6 @@ impl OptimizerRule for ConstantFolding {
 struct Simplifier<'a> {
     /// input schemas
     schemas: Vec<&'a DFSchemaRef>,
-    execution_props: &'a ExecutionProps,
 }
 
 impl<'a> Simplifier<'a> {
@@ -228,15 +220,6 @@ impl<'a> ExprRewriter for Simplifier<'a> {
                     Expr::Not(inner)
                 }
             }
-            // convert now() --> the time in `ExecutionProps`
-            Expr::ScalarFunction {
-                fun: BuiltinScalarFunction::Now,
-                ..
-            } => Expr::Literal(ScalarValue::TimestampNanosecond(Some(
-                self.execution_props
-                    .query_execution_start_time
-                    .timestamp_nanos(),
-            ))),
             expr => {
                 // no additional rewrites possible
                 expr
@@ -248,10 +231,13 @@ impl<'a> ExprRewriter for Simplifier<'a> {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use super::*;
     use crate::{
         assert_contains,
         logical_plan::{col, lit, max, min, DFField, DFSchema, LogicalPlanBuilder},
+        physical_plan::functions::BuiltinScalarFunction,
     };
 
     use arrow::datatypes::*;
@@ -282,7 +268,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         assert_eq!(
@@ -298,7 +283,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         // x = null is always null
@@ -334,7 +318,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         assert_eq!(col("c2").get_type(&schema)?, DataType::Boolean);
@@ -365,7 +348,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         // When one of the operand is not of boolean type, folding the other boolean constant will
@@ -405,7 +387,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         assert_eq!(col("c2").get_type(&schema)?, DataType::Boolean);
@@ -441,7 +422,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         // when one of the operand is not of boolean type, folding the other boolean constant will
@@ -477,7 +457,6 @@ mod tests {
         let schema = expr_test_schema();
         let mut rewriter = Simplifier {
             schemas: vec![&schema],
-            execution_props: &ExecutionProps::new(),
         };
 
         assert_eq!(
@@ -753,27 +732,6 @@ mod tests {
         }
     }
 
-    #[test]
-    fn single_now_expr() {
-        let table_scan = test_table_scan().unwrap();
-        let proj = vec![now_expr()];
-        let time = Utc::now();
-        let plan = LogicalPlanBuilder::from(table_scan)
-            .project(proj)
-            .unwrap()
-            .build()
-            .unwrap();
-
-        let expected = format!(
-            "Projection: TimestampNanosecond({})\
-            \n  TableScan: test projection=None",
-            time.timestamp_nanos()
-        );
-        let actual = get_optimized_plan_formatted(&plan, &time);
-
-        assert_eq!(expected, actual);
-    }
-
     #[test]
     fn multiple_now_expr() {
         let table_scan = test_table_scan().unwrap();
@@ -838,17 +796,16 @@ mod tests {
         //  now() < cast(to_timestamp(...) as int) + 5000000000
         let plan = LogicalPlanBuilder::from(table_scan)
             .filter(
-                now_expr()
+                cast_to_int64_expr(now_expr())
                     .lt(cast_to_int64_expr(to_timestamp_expr(ts_string)) + lit(50000)),
             )
             .unwrap()
             .build()
             .unwrap();
 
-        // Note that constant folder should be able to run again and fold
-        // this whole expression down to a single constant;
-        // https://github.com/apache/arrow-datafusion/issues/1160
-        let expected = "Filter: TimestampNanosecond(1599566400000000000) < CAST(totimestamp(Utf8(\"2020-09-08T12:05:00+00:00\")) AS Int64) + Int32(50000)\
+        // Note that constant folder runs and folds the entire
+        // expression down to a single constant (true)
+        let expected = "Filter: Boolean(true)\
                         \n  TableScan: test projection=None";
         let actual = get_optimized_plan_formatted(&plan, &time);