compile datafusion_ext

crates/datafusion_ext/src/lib.rs

@@ -1,4 +1,3 @@
-pub mod cast;
 pub mod errors;
 pub mod metrics;
 pub mod planner;

crates/datafusion_ext/src/planner/expr/function.rs

@@ -204,7 +204,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     }
 
     pub(super) async fn sql_named_function_to_expr(
-        &self,
+        &mut self,
         expr: SQLExpr,
         fun: BuiltinScalarFunction,
         schema: &DFSchema,
@@ -239,7 +239,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     }
 
     async fn sql_fn_arg_to_logical_expr(
-        &self,
+        &mut self,
         sql: FunctionArg,
         schema: &DFSchema,
         planner_context: &mut PlannerContext,
@@ -268,7 +268,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     }
 
     pub(super) async fn function_args_to_expr(
-        &self,
+        &mut self,
         args: Vec<FunctionArg>,
         schema: &DFSchema,
         planner_context: &mut PlannerContext,

crates/datafusion_ext/src/planner/expr/identifier.rs

@@ -186,7 +186,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     }
 
     pub(super) async fn sql_case_identifier_to_expr(
-        &self,
+        &mut self,
         operand: Option<Box<SQLExpr>>,
         conditions: Vec<SQLExpr>,
         results: Vec<SQLExpr>,

crates/datafusion_ext/src/planner/expr/mod.rs

@@ -47,7 +47,7 @@ use datafusion::sql::sqlparser::parser::ParserError::ParserError;
 impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     #[async_recursion]
     pub(crate) async fn sql_expr_to_logical_expr(
-        &self,
+        &mut self,
         sql: SQLExpr,
         schema: &DFSchema,
         planner_context: &mut PlannerContext,
@@ -714,7 +714,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     }
 
     async fn sql_agg_with_filter_to_expr(
-        &self,
+        &mut self,
         expr: SQLExpr,
         filter: SQLExpr,
         schema: &DFSchema,

crates/datafusion_ext/src/planner/expr/order_by.rs

@@ -29,7 +29,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     /// into the SELECT list (e.g. `SELECT a, b FROM table ORDER BY 2`).
     /// If false, interpret numeric literals as constant values.
     pub(crate) async fn order_by_to_sort_expr(
-        &self,
+        &mut self,
         exprs: &[OrderByExpr],
         schema: &DFSchema,
         planner_context: &mut PlannerContext,

crates/datafusion_ext/src/planner/expr/value.rs

@@ -109,7 +109,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     }
 
     pub(super) async fn sql_array_literal(
-        &self,
+        &mut self,
         elements: Vec<SQLExpr>,
         schema: &DFSchema,
     ) -> Result<Expr> {

crates/datafusion_ext/src/planner/select.rs

@@ -352,7 +352,7 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
     /// Generate a relational expression from a select SQL expression
     #[async_recursion]
     async fn sql_select_to_rex(
-        &self,
+        &mut self,
         sql: SelectItem,
         plan: &LogicalPlan,
         empty_from: bool,

crates/datafusion_ext/src/planner/set_expr.rs

@@ -17,7 +17,7 @@
 
 use crate::planner::{AsyncContextProvider, SqlQueryPlanner};
 use async_recursion::async_recursion;
-use datafusion::common::{DataFusionError, Result};
+use datafusion::common::{not_impl_err, DataFusionError, Result};
 use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder};
 use datafusion::sql::planner::PlannerContext;
 use datafusion::sql::sqlparser::ast::{SetExpr, SetOperator, SetQuantifier};
@@ -42,14 +42,13 @@ impl<'a, S: AsyncContextProvider> SqlQueryPlanner<'a, S> {
                     SetQuantifier::All => true,
                     SetQuantifier::Distinct | SetQuantifier::None => false,
                     SetQuantifier::ByName => {
-                        return Err(DataFusionError::NotImplemented(
-                            "UNION BY NAME not implemented".to_string(),
-                        ));
+                        return not_impl_err!("UNION BY NAME not implemented");
                     }
                     SetQuantifier::AllByName => {
-                        return Err(DataFusionError::NotImplemented(
-                            "UNION ALL BY NAME not implemented".to_string(),
-                        ))
+                        return not_impl_err!("UNION ALL BY NAME not implemented");
+                    }
+                    SetQuantifier::DistinctByName => {
+                        return not_impl_err!("UNION DISTINCT BY NAME not implemented");
                     }
                 };
 

crates/datafusion_ext/src/transform.rs

@@ -10,7 +10,7 @@ pub trait TreeNodeExt: TreeNode {
     where
         F: FnMut(Self) -> Result<Transformed<Self>>,
     {
-        let after_op_children = self.map_children(|node| node.transform_up_mut(op))?;
+        let after_op_children = self.map_children(|node| TreeNode::transform_up_mut(node, op))?;
 
         let new_node = op(after_op_children)?.into();
         Ok(new_node)
@@ -23,7 +23,7 @@ pub trait TreeNodeExt: TreeNode {
         F: FnMut(Self) -> Result<Transformed<Self>>,
     {
         let after_op = op(self)?.into();
-        after_op.map_children(|node| node.transform_down_mut(op))
+        after_op.map_children(|node| TreeNode::transform_down_mut(node, op))
     }
 }
 

crates/datafusion_ext/src/vars.rs

@@ -5,6 +5,7 @@ mod inner;
 mod utils;
 mod value;
 use constants::*;
+use datafusion::arrow::array::{ListBuilder, StringBuilder};
 use datafusion::arrow::datatypes::{DataType, Field};
 use datafusion::config::{ConfigExtension, ExtensionOptions};
 use datafusion::scalar::ScalarValue;
@@ -247,25 +248,27 @@ impl VarProvider for SessionVars {
             "current_schema" => ScalarValue::Utf8(self.search_path().first().cloned()),
             "connection_id" => ScalarValue::Utf8(Some(self.connection_id().to_string())),
             "current_schemas" => {
-                let schemas = self
-                    .search_path()
-                    .into_iter()
-                    .map(|path| ScalarValue::Utf8(Some(path)))
-                    .collect::<Vec<_>>();
-                ScalarValue::List(
-                    Some(schemas),
-                    Field::new("item", DataType::Utf8, true).into(),
-                )
+                let search_path = self.search_path();
+                let mut list = ListBuilder::with_capacity(
+                    StringBuilder::with_capacity(search_path.len(), search_path.len() * 10),
+                    /* list capacity = */ 1,
+                );
+                list.append_value(search_path.into_iter().map(Some));
+                let list = list.finish();
+                ScalarValue::List(Arc::new(list))
             }
             "current_schemas_include_implicit" => {
-                let schemas = self
-                    .implicit_search_path_iter()
-                    .map(|path| ScalarValue::Utf8(Some(path)))
-                    .collect::<Vec<_>>();
-                ScalarValue::List(
-                    Some(schemas),
-                    Field::new("item", DataType::Utf8, true).into(),
-                )
+                let implicit_search_path = self.implicit_search_path();
+                let mut list = ListBuilder::with_capacity(
+                    StringBuilder::with_capacity(
+                        implicit_search_path.len(),
+                        implicit_search_path.len() * 10,
+                    ),
+                    /* list capacity = */ 1,
+                );
+                list.append_value(implicit_search_path.into_iter().map(Some));
+                let list = list.finish();
+                ScalarValue::List(Arc::new(list))
             }
             s => Err(datafusion::error::DataFusionError::External(
                 VarError::UnknownVariable(s.to_string()).into(),

crates/sqlbuiltins/src/functions/scalars/mod.rs

@@ -65,16 +65,60 @@ fn get_nth_scalar_value(
     match input.get(n) {
         Some(input) => match input {
             ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar(op(scalar.clone())?)),
-            ColumnarValue::Array(arr) => Ok(ColumnarValue::Array(scalar_iter_to_array(
-                (0..arr.len()).map(|idx| -> Result<ScalarValue, ExtensionError> {
-                    Ok(op(ScalarValue::try_from_array(arr, idx)?)?)
-                }),
-            )?)),
+            ColumnarValue::Array(arr) => Ok(ColumnarValue::Array(apply_op_to_col_array(arr, op)?)),
         },
         None => Err(BuiltinError::MissingValueAtIndex(n)),
     }
 }
 
+fn apply_op_to_col_array(
+    arr: &dyn Array,
+    op: &dyn Fn(ScalarValue) -> Result<ScalarValue, BuiltinError>,
+) -> Result<Arc<dyn Array>> {
+    let mut check_err: Result<()> = Ok(());
+
+    let filter_fn = |res: Result<ScalarValue>| {
+        match (check_err, res) {
+            // If check_err is already an error, return None (so we
+            // don't iterate further).
+            (Err(_), _) => None,
+            (_, Err(e)) => {
+                // Set check_err to the corresponding error.
+                check_err = Err(e);
+                None
+            }
+            (_, Ok(scalar)) => Some(scalar),
+        }
+    };
+
+    let iter = (0..arr.len())
+        .filter_map(|idx| {
+            let res = ScalarValue::try_from_array(arr, idx);
+            filter_fn(res)
+        })
+        .filter_map(|s| {
+            let res = op(s);
+            filter_fn(res)
+        });
+
+    // NB: ScalarValue::iter_to_array accepts an iterator over
+    // Item = ScalarValue but we have an iterator over Item =
+    // Result<ScalarValue>. To convert, we filter-map the errors but
+    // that doesn't let us catch the error. To catch the error, we
+    // have `check_err` (a result) that is initially set to Ok and
+    // set to the error when the iterator is actually iterated over
+    // the values in iter_to_array. We filter out all values once
+    // the check_err is set to an error value so we don't iterate
+    // any further.
+    //
+    // A simpler solution would have been to collect all the values
+    // in another container (such as a Vec) but that would result in
+    // extra space being used.
+    let arr = ScalarValue::iter_to_array(iter)?;
+    check_err?;
+    Ok(arr)
+}
+
 fn try_from_u64_scalar(scalar: ScalarValue) -> Result<u64, BuiltinError> {
     match scalar {
         ScalarValue::Int8(Some(value)) => safe_up_cast_integer_scalar(value as i64),