diff --git a/datafusion/expr/src/field_util.rs b/datafusion/expr/src/field_util.rs index 3195ce6f2dfc..8039a211c9e4 100644 --- a/datafusion/expr/src/field_util.rs +++ b/datafusion/expr/src/field_util.rs @@ -87,10 +87,11 @@ impl GetFieldAccessSchema { 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::List(_), _, _, _) => plan_err!( + (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" ), - (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}"), } } } diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 9175ccb46859..3f65c68bc45b 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -525,9 +525,13 @@ impl<'a> ConstEvaluator<'a> { DataFusionError::Execution(format!("Could not evaluate the expression, found a result of length {}", a.len())), expr, ) - } else if as_list_array(&a).is_ok() || as_large_list_array(&a).is_ok() { + } else if as_list_array(&a).is_ok() { ConstSimplifyResult::Simplified(ScalarValue::List( - a.as_list().to_owned().into(), + a.as_list::().to_owned().into(), + )) + } else if as_large_list_array(&a).is_ok() { + ConstSimplifyResult::Simplified(ScalarValue::LargeList( + a.as_list::().to_owned().into(), )) } else { // Non-ListArray diff --git a/datafusion/physical-expr/src/expressions/get_indexed_field.rs b/datafusion/physical-expr/src/expressions/get_indexed_field.rs index 39eef61f963a..773387bf7421 100644 --- a/datafusion/physical-expr/src/expressions/get_indexed_field.rs +++ b/datafusion/physical-expr/src/expressions/get_indexed_field.rs @@ -268,16 +268,18 @@ impl PhysicalExpr for GetIndexedFieldExpr { let stop = stop.evaluate(batch)?.into_array(batch.num_rows())?; let stride = stride.evaluate(batch)?.into_array(batch.num_rows())?; match (array.data_type(), start.data_type(), stop.data_type(), stride.data_type()) { - (DataType::List(_), DataType::Int64, DataType::Int64, DataType::Int64) => { + (DataType::List(_), DataType::Int64, DataType::Int64, DataType::Int64) | + (DataType::LargeList(_), DataType::Int64, DataType::Int64, DataType::Int64)=> { Ok(ColumnarValue::Array((array_slice(&[ array, start, stop, stride ]))?)) }, - (DataType::List(_), start, stop, stride) => exec_err!( - "get indexed field is only possible on lists with int64 indexes. \ + (DataType::List(_), start, stop, stride) | + (DataType::LargeList(_), start, stop, stride)=> exec_err!( + "get indexed field is only possible on List/LargeList with int64 indexes. \ Tried with {start:?}, {stop:?} and {stride:?} indices"), (dt, start, stop, stride) => 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 {start:?}, {stop:?} and {stride:?} indices"), } } diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index 640bf82b5520..4e6cb4d59d14 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/datafusion/sqllogictest/test_files/array.slt @@ -861,12 +861,28 @@ select make_array(1, 2, 3)[1:2], make_array(1.0, 2.0, 3.0)[2:3], make_array('h', ---- [1, 2] [2.0, 3.0] [e, l, l] +query ??? +select arrow_cast([1, 2, 3], 'LargeList(Int64)')[1:2], + arrow_cast([1.0, 2.0, 3.0], 'LargeList(Int64)')[2:3], + arrow_cast(['h', 'e', 'l', 'l', 'o'], 'LargeList(Utf8)')[2:4] +; +---- +[1, 2] [2, 3] [e, l, l] + # multiple index with columns #2 (zero index) query ??? select make_array(1, 2, 3)[0:0], make_array(1.0, 2.0, 3.0)[0:2], make_array('h', 'e', 'l', 'l', 'o')[0:6]; ---- [] [1.0, 2.0] [h, e, l, l, o] +query ??? +select arrow_cast([1, 2, 3], 'LargeList(Int64)')[0:0], + arrow_cast([1.0, 2.0, 3.0], 'LargeList(Int64)')[0:2], + arrow_cast(['h', 'e', 'l', 'l', 'o'], 'LargeList(Utf8)')[0:6] +; +---- +[] [1, 2] [h, e, l, l, o] + # TODO: support multiple negative index # multiple index with columns #3 (negative index) # query II