Indexed field access for List

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/utils.rs

@@ -85,6 +85,7 @@ impl ExpressionVisitor for ColumnNameVisitor<'_> {
             Expr::AggregateUDF { .. } => {}
             Expr::InList { .. } => {}
             Expr::Wildcard => {}
+            Expr::GetIndexedField { .. } => {}
         }
         Ok(Recursion::Continue(self))
     }
@@ -337,6 +338,7 @@ pub fn expr_sub_expressions(expr: &Expr) -> Result<Vec<Expr>> {
         Expr::Wildcard { .. } => Err(DataFusionError::Internal(
             "Wildcard expressions are not valid in a logical query plan".to_owned(),
         )),
+        Expr::GetIndexedField { expr, .. } => Ok(vec![expr.as_ref().to_owned()]),
     }
 }
 
@@ -496,6 +498,10 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> {
         Expr::Wildcard { .. } => Err(DataFusionError::Internal(
             "Wildcard expressions are not valid in a logical query plan".to_owned(),
         )),
+        Expr::GetIndexedField { expr: _, key } => Ok(Expr::GetIndexedField {
+            expr: Box::new(expressions[0].clone()),
+            key: key.clone(),
+        }),
     }
 }
 
@@ -650,6 +656,7 @@ impl ConstEvaluator {
             Expr::Cast { .. } => true,
             Expr::TryCast { .. } => true,
             Expr::InList { .. } => true,
+            Expr::GetIndexedField { .. } => true,
         }
     }