Support field access for literal Null (#10655)

apache · May 25, 2024 · 0f994af · 0f994af
1 parent d10b1a4
commit 0f994af
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 3 deletions.
diff --git a/datafusion/core/tests/expr_api/mod.rs b/datafusion/core/tests/expr_api/mod.rs
@@ -20,7 +20,7 @@ use arrow_array::builder::{ListBuilder, StringBuilder};
 use arrow_array::{ArrayRef, RecordBatch, StringArray, StructArray};
 use arrow_schema::{DataType, Field};
 use datafusion::prelude::*;
-use datafusion_common::DFSchema;
+use datafusion_common::{DFSchema, ScalarValue};
 use datafusion_functions::core::expr_ext::FieldAccessor;
 use datafusion_functions_array::expr_ext::{IndexAccessor, SliceAccessor};
 /// Tests of using and evaluating `Expr`s outside the context of a LogicalPlan
@@ -78,6 +78,21 @@ fn test_get_field() {
     );
 }
 
+#[test]
+fn test_get_field_null() {
+    #[rustfmt::skip]
+    evaluate_expr_test(
+        lit(ScalarValue::Null).field("a"),
+        vec![
+            "+------+",
+            "| expr |",
+            "+------+",
+            "|      |",
+            "+------+",
+        ],
+    );
+}
+
 #[test]
 fn test_nested_get_field() {
     evaluate_expr_test(
@@ -98,11 +113,17 @@ fn test_nested_get_field() {
 }
 
 #[test]
-fn test_list() {
+fn test_list_index() {
+    #[rustfmt::skip]
     evaluate_expr_test(
         col("list").index(lit(1i64)),
         vec![
-            "+------+", "| expr |", "+------+", "| one  |", "| two  |", "| five |",
+            "+------+",
+            "| expr |",
+            "+------+",
+            "| one  |",
+            "| two  |",
+            "| five |",
             "+------+",
         ],
     );

diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs
@@ -106,6 +106,9 @@ impl ScalarUDFImpl for GetFieldFunc {
         };
         let access_schema = GetFieldAccessSchema::NamedStructField { name: name.clone() };
         let arg_dt = args[0].get_type(schema)?;
+        if arg_dt.is_null() {
+            return Ok(DataType::Null);
+        }
         access_schema
             .get_accessed_field(&arg_dt)
             .map(|f| f.data_type().clone())
@@ -119,6 +122,10 @@ impl ScalarUDFImpl for GetFieldFunc {
             );
         }
 
+        if args[0].data_type().is_null() {
+            return Ok(ColumnarValue::Scalar(ScalarValue::Null));
+        }
+
         let arrays = ColumnarValue::values_to_arrays(args)?;
         let array = arrays[0].clone();