Support LargeList for ListIndex (#9424)

Signed-off-by: Chojan Shang <[email protected]>
apache · Mar 3, 2024 · 89aea0a · 89aea0a
1 parent 487c53b
commit 89aea0a
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/datafusion/expr/src/field_util.rs b/datafusion/expr/src/field_util.rs
@@ -78,18 +78,19 @@ impl GetFieldAccessSchema {
             Self::ListIndex{ key_dt } => {
                 match (data_type, key_dt) {
                     (DataType::List(lt), DataType::Int64) => Ok(Field::new("list", lt.data_type().clone(), true)),
-                    (DataType::List(_), _) => plan_err!(
-                        "Only ints are valid as an indexed field in a list"
+                    (DataType::LargeList(lt), DataType::Int64) => Ok(Field::new("large_list", lt.data_type().clone(), true)),
+                    (DataType::List(_), _) | (DataType::LargeList(_), _) => plan_err!(
+                        "Only ints are valid as an indexed field in a List/LargeList"
                     ),
-                    (other, _) => plan_err!("The expression to get an indexed field is only valid for `List` or `Struct` types, got {other}"),
+                    (other, _) => plan_err!("The expression to get an indexed field is only valid for `List`, `LargeList` or `Struct` types, got {other}"),
                 }
             }
             Self::ListRange { start_dt, stop_dt, stride_dt } => {
                 match (data_type, start_dt, stop_dt, stride_dt) {
                     (DataType::List(_), DataType::Int64, DataType::Int64, DataType::Int64) => Ok(Field::new("list", data_type.clone(), true)),
                     (DataType::LargeList(_), DataType::Int64, DataType::Int64, DataType::Int64) => Ok(Field::new("large_list", data_type.clone(), true)),
                     (DataType::List(_), _, _, _) | (DataType::LargeList(_), _, _, _)=> plan_err!(
-                        "Only ints are valid as an indexed field in a list"
+                        "Only ints are valid as an indexed field in a List/LargeList"
                     ),
                     (other, _, _, _) => plan_err!("The expression to get an indexed field is only valid for `List`, `LargeList` or `Struct` types, got {other}"),
                 }

diff --git a/datafusion/physical-expr/src/expressions/get_indexed_field.rs b/datafusion/physical-expr/src/expressions/get_indexed_field.rs
@@ -252,14 +252,14 @@ impl PhysicalExpr for GetIndexedFieldExpr {
             GetFieldAccessExpr::ListIndex{key} => {
             let key = key.evaluate(batch)?.into_array(batch.num_rows())?;
             match (array.data_type(), key.data_type()) {
-                (DataType::List(_), DataType::Int64) => Ok(ColumnarValue::Array(array_element(&[
+                (DataType::List(_), DataType::Int64) | (DataType::LargeList(_), DataType::Int64) => Ok(ColumnarValue::Array(array_element(&[
                     array, key
                 ])?)),
-                (DataType::List(_), key) => exec_err!(
-                                "get indexed field is only possible on lists with int64 indexes. \
+                (DataType::List(_), key) | (DataType::LargeList(_), key) => exec_err!(
+                                "get indexed field is only possible on List/LargeList with int64 indexes. \
                                     Tried with {key:?} index"),
                             (dt, key) => exec_err!(
-                                        "get indexed field is only possible on lists with int64 indexes or struct \
+                                        "get indexed field is only possible on List/LargeList with int64 indexes or struct \
                                                  with utf8 indexes. Tried {dt:?} with {key:?} index"),
                         }
                 },

diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt
@@ -883,6 +883,11 @@ select arrow_cast([1, 2, 3], 'LargeList(Int64)')[0:0],
 ----
 [] [1, 2] [h, e, l, l, o]
 
+query I
+select arrow_cast([1, 2, 3], 'LargeList(Int64)')[1];
+----
+1
+
 # TODO: support multiple negative index
 # multiple index with columns #3 (negative index)
 # query II