diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index ac3b7ebaeac1..d383ddce9242 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -209,10 +209,15 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result { let key = create_physical_name(key, false)?; format!("{expr}[{key}]") } - GetFieldAccess::ListRange { start, stop } => { + GetFieldAccess::ListRange { + start, + stop, + stride, + } => { let start = create_physical_name(start, false)?; let stop = create_physical_name(stop, false)?; - format!("{expr}[{start}:{stop}]") + let stride = create_physical_name(stride, false)?; + format!("{expr}[{start}:{stop}:{stride}]") } }; diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index c5d158d87638..9da1f4bb4df7 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -421,8 +421,12 @@ pub enum GetFieldAccess { NamedStructField { name: ScalarValue }, /// Single list index, for example: `list[i]` ListIndex { key: Box }, - /// List range, for example `list[i:j]` - ListRange { start: Box, stop: Box }, + /// List stride, for example `list[i:j:k]` + ListRange { + start: Box, + stop: Box, + stride: Box, + }, } /// Returns the field of a [`arrow::array::ListArray`] or @@ -1209,7 +1213,7 @@ impl Expr { /// # use datafusion_expr::{lit, col}; /// let expr = col("c1") /// .range(lit(2), lit(4)); - /// assert_eq!(expr.display_name().unwrap(), "c1[Int32(2):Int32(4)]"); + /// assert_eq!(expr.display_name().unwrap(), "c1[Int32(2):Int32(4):Int64(1)]"); /// ``` pub fn range(self, start: Expr, stop: Expr) -> Self { Expr::GetIndexedField(GetIndexedField { @@ -1217,6 +1221,7 @@ impl Expr { field: GetFieldAccess::ListRange { start: Box::new(start), stop: Box::new(stop), + stride: Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))), }, }) } @@ -1530,8 +1535,12 @@ impl fmt::Display for Expr { write!(f, "({expr})[{name}]") } GetFieldAccess::ListIndex { key } => write!(f, "({expr})[{key}]"), - GetFieldAccess::ListRange { start, stop } => { - write!(f, "({expr})[{start}:{stop}]") + GetFieldAccess::ListRange { + start, + stop, + stride, + } => { + write!(f, "({expr})[{start}:{stop}:{stride}]") } }, Expr::GroupingSet(grouping_sets) => match grouping_sets { @@ -1732,10 +1741,15 @@ fn create_name(e: &Expr) -> Result { let key = create_name(key)?; Ok(format!("{expr}[{key}]")) } - GetFieldAccess::ListRange { start, stop } => { + GetFieldAccess::ListRange { + start, + stop, + stride, + } => { let start = create_name(start)?; let stop = create_name(stop)?; - Ok(format!("{expr}[{start}:{stop}]")) + let stride = create_name(stride)?; + Ok(format!("{expr}[{start}:{stop}:{stride}]")) } } } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index ba21d09f0619..4967e66fed40 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -374,9 +374,14 @@ fn field_for_index( GetFieldAccess::ListIndex { key } => GetFieldAccessSchema::ListIndex { key_dt: key.get_type(schema)?, }, - GetFieldAccess::ListRange { start, stop } => GetFieldAccessSchema::ListRange { + GetFieldAccess::ListRange { + start, + stop, + stride, + } => GetFieldAccessSchema::ListRange { start_dt: start.get_type(schema)?, stop_dt: stop.get_type(schema)?, + stride_dt: stride.get_type(schema)?, }, } .get_accessed_field(&expr_dt) diff --git a/datafusion/expr/src/field_util.rs b/datafusion/expr/src/field_util.rs index 3829a2086b26..c46ec50234dd 100644 --- a/datafusion/expr/src/field_util.rs +++ b/datafusion/expr/src/field_util.rs @@ -28,10 +28,11 @@ pub enum GetFieldAccessSchema { NamedStructField { name: ScalarValue }, /// Single list index, for example: `list[i]` ListIndex { key_dt: DataType }, - /// List range, for example `list[i:j]` + /// List stride, for example `list[i:j:k]` ListRange { start_dt: DataType, stop_dt: DataType, + stride_dt: DataType, }, } @@ -85,13 +86,13 @@ impl GetFieldAccessSchema { (other, _) => plan_err!("The expression to get an indexed field is only valid for `List` or `Struct` types, got {other}"), } } - Self::ListRange{ start_dt, stop_dt } => { - match (data_type, start_dt, stop_dt) { - (DataType::List(_), DataType::Int64, DataType::Int64) => Ok(Field::new("list", data_type.clone(), true)), - (DataType::List(_), _, _) => plan_err!( + 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!( "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` or `Struct` types, got {other}"), } } } diff --git a/datafusion/expr/src/tree_node/expr.rs b/datafusion/expr/src/tree_node/expr.rs index 05464c96d05e..8b38d1cf01d6 100644 --- a/datafusion/expr/src/tree_node/expr.rs +++ b/datafusion/expr/src/tree_node/expr.rs @@ -52,8 +52,8 @@ impl TreeNode for Expr { let expr = expr.as_ref(); match field { GetFieldAccess::ListIndex {key} => vec![key.as_ref(), expr], - GetFieldAccess::ListRange {start, stop} => { - vec![start.as_ref(), stop.as_ref(), expr] + GetFieldAccess::ListRange {start, stop, stride} => { + vec![start.as_ref(), stop.as_ref(),stride.as_ref(), expr] } GetFieldAccess::NamedStructField { .. } => vec![expr], } diff --git a/datafusion/physical-expr/src/expressions/get_indexed_field.rs b/datafusion/physical-expr/src/expressions/get_indexed_field.rs index 43fd5a812a16..58fe4728543d 100644 --- a/datafusion/physical-expr/src/expressions/get_indexed_field.rs +++ b/datafusion/physical-expr/src/expressions/get_indexed_field.rs @@ -21,6 +21,7 @@ use crate::PhysicalExpr; use datafusion_common::exec_err; use crate::array_expressions::{array_element, array_slice}; +use crate::expressions::Literal; use crate::physical_expr::down_cast_any_ref; use arrow::{ array::{Array, Scalar, StringArray}, @@ -43,10 +44,11 @@ pub enum GetFieldAccessExpr { NamedStructField { name: ScalarValue }, /// Single list index, for example: `list[i]` ListIndex { key: Arc }, - /// List range, for example `list[i:j]` + /// List stride, for example `list[i:j:k]` ListRange { start: Arc, stop: Arc, + stride: Arc, }, } @@ -55,8 +57,12 @@ impl std::fmt::Display for GetFieldAccessExpr { match self { GetFieldAccessExpr::NamedStructField { name } => write!(f, "[{}]", name), GetFieldAccessExpr::ListIndex { key } => write!(f, "[{}]", key), - GetFieldAccessExpr::ListRange { start, stop } => { - write!(f, "[{}:{}]", start, stop) + GetFieldAccessExpr::ListRange { + start, + stop, + stride, + } => { + write!(f, "[{}:{}:{}]", start, stop, stride) } } } @@ -76,12 +82,18 @@ impl PartialEq for GetFieldAccessExpr { ListRange { start: start_lhs, stop: stop_lhs, + stride: stride_lhs, }, ListRange { start: start_rhs, stop: stop_rhs, + stride: stride_rhs, }, - ) => start_lhs.eq(start_rhs) && stop_lhs.eq(stop_rhs), + ) => { + start_lhs.eq(start_rhs) + && stop_lhs.eq(stop_rhs) + && stride_lhs.eq(stride_rhs) + } (NamedStructField { .. }, ListIndex { .. } | ListRange { .. }) => false, (ListIndex { .. }, NamedStructField { .. } | ListRange { .. }) => false, (ListRange { .. }, NamedStructField { .. } | ListIndex { .. }) => false, @@ -126,7 +138,32 @@ impl GetIndexedFieldExpr { start: Arc, stop: Arc, ) -> Self { - Self::new(arg, GetFieldAccessExpr::ListRange { start, stop }) + Self::new( + arg, + GetFieldAccessExpr::ListRange { + start, + stop, + stride: Arc::new(Literal::new(ScalarValue::Int64(Some(1)))) + as Arc, + }, + ) + } + + /// Create a new [`GetIndexedFieldExpr`] for accessing the stride + pub fn new_stride( + arg: Arc, + start: Arc, + stop: Arc, + stride: Arc, + ) -> Self { + Self::new( + arg, + GetFieldAccessExpr::ListRange { + start, + stop, + stride, + }, + ) } /// Get the description of what field should be accessed @@ -147,12 +184,15 @@ impl GetIndexedFieldExpr { GetFieldAccessExpr::ListIndex { key } => GetFieldAccessSchema::ListIndex { key_dt: key.data_type(input_schema)?, }, - GetFieldAccessExpr::ListRange { start, stop } => { - GetFieldAccessSchema::ListRange { - start_dt: start.data_type(input_schema)?, - stop_dt: stop.data_type(input_schema)?, - } - } + GetFieldAccessExpr::ListRange { + start, + stop, + stride, + } => GetFieldAccessSchema::ListRange { + start_dt: start.data_type(input_schema)?, + stop_dt: stop.data_type(input_schema)?, + stride_dt: stride.data_type(input_schema)?, + }, }) } } @@ -223,21 +263,24 @@ impl PhysicalExpr for GetIndexedFieldExpr { with utf8 indexes. Tried {dt:?} with {key:?} index"), } }, - GetFieldAccessExpr::ListRange{start, stop} => { + GetFieldAccessExpr::ListRange { start, stop, stride } => { let start = start.evaluate(batch)?.into_array(batch.num_rows())?; let stop = stop.evaluate(batch)?.into_array(batch.num_rows())?; - match (array.data_type(), start.data_type(), stop.data_type()) { - (DataType::List(_), DataType::Int64, DataType::Int64) => Ok(ColumnarValue::Array(array_slice(&[ - array, start, stop - ])?)), - (DataType::List(_), start, stop) => exec_err!( + 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) => { + 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. \ - Tried with {start:?} and {stop:?} indices"), - (dt, start, stop) => exec_err!( + 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 \ - with utf8 indexes. Tried {dt:?} with {start:?} and {stop:?} indices"), + with utf8 indexes. Tried {dt:?} with {start:?}, {stop:?} and {stride:?} indices"), } - }, + } } } diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 09b8da836c30..ee5da05d1151 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -238,20 +238,19 @@ pub fn create_physical_expr( GetFieldAccess::ListIndex { key } => GetFieldAccessExpr::ListIndex { key: create_physical_expr(key, input_dfschema, execution_props)?, }, - GetFieldAccess::ListRange { start, stop } => { - GetFieldAccessExpr::ListRange { - start: create_physical_expr( - start, - input_dfschema, - execution_props, - )?, - stop: create_physical_expr( - stop, - input_dfschema, - execution_props, - )?, - } - } + GetFieldAccess::ListRange { + start, + stop, + stride, + } => GetFieldAccessExpr::ListRange { + start: create_physical_expr(start, input_dfschema, execution_props)?, + stop: create_physical_expr(stop, input_dfschema, execution_props)?, + stride: create_physical_expr( + stride, + input_dfschema, + execution_props, + )?, + }, }; Ok(Arc::new(GetIndexedFieldExpr::new( create_physical_expr(expr, input_dfschema, execution_props)?, diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 66c1271e65c1..c8468e1709c3 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -466,6 +466,7 @@ message ListIndex { message ListRange { LogicalExprNode start = 1; LogicalExprNode stop = 2; + LogicalExprNode stride = 3; } message GetIndexedField { @@ -1773,6 +1774,7 @@ message ListIndexExpr { message ListRangeExpr { PhysicalExprNode start = 1; PhysicalExprNode stop = 2; + PhysicalExprNode stride = 3; } message PhysicalGetIndexedFieldExprNode { diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index 39a8678ef250..47667fb68c43 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -12542,6 +12542,9 @@ impl serde::Serialize for ListRange { if self.stop.is_some() { len += 1; } + if self.stride.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ListRange", len)?; if let Some(v) = self.start.as_ref() { struct_ser.serialize_field("start", v)?; @@ -12549,6 +12552,9 @@ impl serde::Serialize for ListRange { if let Some(v) = self.stop.as_ref() { struct_ser.serialize_field("stop", v)?; } + if let Some(v) = self.stride.as_ref() { + struct_ser.serialize_field("stride", v)?; + } struct_ser.end() } } @@ -12561,12 +12567,14 @@ impl<'de> serde::Deserialize<'de> for ListRange { const FIELDS: &[&str] = &[ "start", "stop", + "stride", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Start, Stop, + Stride, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -12590,6 +12598,7 @@ impl<'de> serde::Deserialize<'de> for ListRange { match value { "start" => Ok(GeneratedField::Start), "stop" => Ok(GeneratedField::Stop), + "stride" => Ok(GeneratedField::Stride), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -12611,6 +12620,7 @@ impl<'de> serde::Deserialize<'de> for ListRange { { let mut start__ = None; let mut stop__ = None; + let mut stride__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Start => { @@ -12625,11 +12635,18 @@ impl<'de> serde::Deserialize<'de> for ListRange { } stop__ = map_.next_value()?; } + GeneratedField::Stride => { + if stride__.is_some() { + return Err(serde::de::Error::duplicate_field("stride")); + } + stride__ = map_.next_value()?; + } } } Ok(ListRange { start: start__, stop: stop__, + stride: stride__, }) } } @@ -12650,6 +12667,9 @@ impl serde::Serialize for ListRangeExpr { if self.stop.is_some() { len += 1; } + if self.stride.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ListRangeExpr", len)?; if let Some(v) = self.start.as_ref() { struct_ser.serialize_field("start", v)?; @@ -12657,6 +12677,9 @@ impl serde::Serialize for ListRangeExpr { if let Some(v) = self.stop.as_ref() { struct_ser.serialize_field("stop", v)?; } + if let Some(v) = self.stride.as_ref() { + struct_ser.serialize_field("stride", v)?; + } struct_ser.end() } } @@ -12669,12 +12692,14 @@ impl<'de> serde::Deserialize<'de> for ListRangeExpr { const FIELDS: &[&str] = &[ "start", "stop", + "stride", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Start, Stop, + Stride, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -12698,6 +12723,7 @@ impl<'de> serde::Deserialize<'de> for ListRangeExpr { match value { "start" => Ok(GeneratedField::Start), "stop" => Ok(GeneratedField::Stop), + "stride" => Ok(GeneratedField::Stride), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -12719,6 +12745,7 @@ impl<'de> serde::Deserialize<'de> for ListRangeExpr { { let mut start__ = None; let mut stop__ = None; + let mut stride__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Start => { @@ -12733,11 +12760,18 @@ impl<'de> serde::Deserialize<'de> for ListRangeExpr { } stop__ = map_.next_value()?; } + GeneratedField::Stride => { + if stride__.is_some() { + return Err(serde::de::Error::duplicate_field("stride")); + } + stride__ = map_.next_value()?; + } } } Ok(ListRangeExpr { start: start__, stop: stop__, + stride: stride__, }) } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 7bf1d8ed0450..a5582cc2dc64 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -731,6 +731,8 @@ pub struct ListRange { pub start: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, boxed, tag = "2")] pub stop: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "3")] + pub stride: ::core::option::Option<::prost::alloc::boxed::Box>, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] @@ -2538,6 +2540,8 @@ pub struct ListRangeExpr { pub start: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(message, optional, boxed, tag = "2")] pub stop: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "3")] + pub stride: ::core::option::Option<::prost::alloc::boxed::Box>, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 42d39b5c5139..eb72d1f9c3e8 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -1067,6 +1067,11 @@ pub fn parse_expr( registry, "stop", )?), + stride: Box::new(parse_required_expr( + list_range.stride.as_deref(), + registry, + "stride", + )?), } } None => return Err(proto_error("Field must not be None")), diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index dbb52eced36c..e1fc3f0c8525 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -1033,14 +1033,17 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode { }, )) } - GetFieldAccess::ListRange { start, stop } => { - protobuf::get_indexed_field::Field::ListRange(Box::new( - protobuf::ListRange { - start: Some(Box::new(start.as_ref().try_into()?)), - stop: Some(Box::new(stop.as_ref().try_into()?)), - }, - )) - } + GetFieldAccess::ListRange { + start, + stop, + stride, + } => protobuf::get_indexed_field::Field::ListRange(Box::new( + protobuf::ListRange { + start: Some(Box::new(start.as_ref().try_into()?)), + stop: Some(Box::new(stop.as_ref().try_into()?)), + stride: Some(Box::new(stride.as_ref().try_into()?)), + }, + )), }; Self { diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index dc827d02bf25..454f74dfd132 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -411,8 +411,15 @@ pub fn parse_physical_expr( "stop", input_schema )?, + stride: parse_required_physical_expr( + list_range_expr.stride.as_deref(), + registry, + "stride", + input_schema + )?, }, - None => return Err(proto_error( + None => + return Err(proto_error( "Field must not be None", )), }; diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index cff32ca2f8c9..a67410da57f4 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -562,12 +562,15 @@ impl TryFrom> for protobuf::PhysicalExprNode { key: Some(Box::new(key.to_owned().try_into()?)) })) ), - GetFieldAccessExpr::ListRange{start, stop} => Some( - protobuf::physical_get_indexed_field_expr_node::Field::ListRangeExpr(Box::new(protobuf::ListRangeExpr { - start: Some(Box::new(start.to_owned().try_into()?)), - stop: Some(Box::new(stop.to_owned().try_into()?)), - })) - ), + GetFieldAccessExpr::ListRange { start, stop, stride } => { + Some( + protobuf::physical_get_indexed_field_expr_node::Field::ListRangeExpr(Box::new(protobuf::ListRangeExpr { + start: Some(Box::new(start.to_owned().try_into()?)), + stop: Some(Box::new(stop.to_owned().try_into()?)), + stride: Some(Box::new(stride.to_owned().try_into()?)), + })) + ) + } }; Ok(protobuf::PhysicalExprNode { diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 38eb39000317..eba3db298f84 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -35,6 +35,7 @@ use datafusion::logical_expr::{ create_udf, BuiltinScalarFunction, JoinType, Operator, Volatility, }; use datafusion::parquet::file::properties::WriterProperties; +use datafusion::physical_expr::expressions::Literal; use datafusion::physical_expr::window::SlidingAggregateWindowExpr; use datafusion::physical_expr::{PhysicalSortRequirement, ScalarFunctionExpr}; use datafusion::physical_plan::aggregates::{ @@ -750,6 +751,8 @@ fn roundtrip_get_indexed_field_list_range() -> Result<()> { GetFieldAccessExpr::ListRange { start: col_start, stop: col_stop, + stride: Arc::new(Literal::new(ScalarValue::Int64(Some(1)))) + as Arc, }, )); diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index 9fded63af3fc..b22c458b6db6 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -753,18 +753,47 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { operator: JsonOperator::Colon, right, } => { - let start = Box::new(self.sql_expr_to_logical_expr( - *left, - schema, - planner_context, - )?); - let stop = Box::new(self.sql_expr_to_logical_expr( - *right, - schema, - planner_context, - )?); - - GetFieldAccess::ListRange { start, stop } + let (start, stop, stride) = if let SQLExpr::JsonAccess { + left: l, + operator: JsonOperator::Colon, + right: r, + } = *left + { + let start = Box::new(self.sql_expr_to_logical_expr( + *l, + schema, + planner_context, + )?); + let stop = Box::new(self.sql_expr_to_logical_expr( + *r, + schema, + planner_context, + )?); + let stride = Box::new(self.sql_expr_to_logical_expr( + *right, + schema, + planner_context, + )?); + (start, stop, stride) + } else { + let start = Box::new(self.sql_expr_to_logical_expr( + *left, + schema, + planner_context, + )?); + let stop = Box::new(self.sql_expr_to_logical_expr( + *right, + schema, + planner_context, + )?); + let stride = Box::new(Expr::Literal(ScalarValue::Int64(Some(1)))); + (start, stop, stride) + }; + GetFieldAccess::ListRange { + start, + stop, + stride, + } } _ => GetFieldAccess::ListIndex { key: Box::new(self.sql_expr_to_logical_expr( diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index b7d92aec88e6..e072e4146f13 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/datafusion/sqllogictest/test_files/array.slt @@ -759,6 +759,51 @@ select column1[0:5], column2[0:3], column3[0:9] from arrays; # select column1[column2:column3] from arrays_with_repeating_elements; # ---- +# array[i:j:k] + +# multiple index with columns #1 (positive index) +query ??? +select make_array(1, 2, 3)[1:2:2], make_array(1.0, 2.0, 3.0)[2:3:2], make_array('h', 'e', 'l', 'l', 'o')[2:4:2]; +---- +[1] [2.0] [e, l] + +# multiple index with columns #2 (zero index) +query ??? +select make_array(1, 2, 3)[0:0:2], make_array(1.0, 2.0, 3.0)[0:2:2], make_array('h', 'e', 'l', 'l', 'o')[0:6:2]; +---- +[] [1.0] [h, l, o] + +#TODO: sqlparser does not support negative index +## multiple index with columns #3 (negative index) +#query ??? +#select make_array(1, 2, 3)[-1:-2:-2], make_array(1.0, 2.0, 3.0)[-2:-3:-2], make_array('h', 'e', 'l', 'l', 'o')[-2:-4:-2]; +#---- +#[1] [2.0] [e, l] + +# multiple index with columns #1 (positive index) +query ??? +select column1[2:4:2], column2[1:4:2], column3[3:4:2] from arrays; +---- +[[3, ]] [1.1, 3.3] [r] +[[5, 6]] [, 6.6] [] +[[7, 8]] [7.7, 9.9] [l] +[[9, 10]] [10.1, 12.2] [t] +[] [13.3, 15.5] [e] +[[13, 14]] [] [] +[[, 18]] [16.6, 18.8] [] + +# multiple index with columns #2 (zero index) +query ??? +select column1[0:5:2], column2[0:3:2], column3[0:9:2] from arrays; +---- +[[, 2]] [1.1, 3.3] [L, r, m] +[[3, 4]] [, 6.6] [i, , m] +[[5, 6]] [7.7, 9.9] [d, l, r] +[[7, ]] [10.1, 12.2] [s, t] +[] [13.3, 15.5] [a, e] +[[11, 12]] [] [,] +[[15, 16]] [16.6, 18.8] [] + ### Array function tests