feat(expr): add array access support (#2883)

nanderstabel · Tao Wu · web-flow · commit a3e618ed5fda · 2022-06-02T14:31:34.000Z
* Implement array access support for  arrays

* Add comment

* Add array_access test

* Add value_at method

* Delete ArrayAccessExpression

* Add array access

* Add array_access

* Expand value_at and add test

* Make array_access a BinaryNullableExpression

* Add test with negative index

* Refactor array_access

* Rewrite test for array access as nullable binary expression

* Fix return_type

* Clean

* Clean

* Add nested test and cargo check

* Add nl

* Fix index out of bounds error

* Add test

* Add ArrayAccess

* Update src/expr/src/expr/expr_binary_nullable.rs

returning DataType::List is not implemented yet.

Co-authored-by: Tao Wu &lt;wutao@singularity-data.com&gt;

* Add comma

* Fix test

* Delete unnecessary 'vec'

* Fix cargo.lock file

Co-authored-by: Tao Wu &lt;wutao@singularity-data.com&gt;
diff --git a/e2e_test/batch/array_access.slt b/e2e_test/batch/array_access.slt
@@ -0,0 +1,25 @@
+statement ok
+SET RW_IMPLICIT_FLUSH TO true;
+
+query T
+select (ARRAY['foo', 'bar'])[-1];
+----
+NULL
+
+query T
+select (ARRAY['foo', 'bar'])[0];
+----
+NULL
+
+query T
+select (ARRAY['foo', 'bar'])[1];
+----
+foo
+
+query T
+select (ARRAY['foo', 'bar'])[3];
+----
+NULL
+
+statement error
+select (ARRAY['foo', 'bar'])[];
diff --git a/src/common/src/array/list_array.rs b/src/common/src/array/list_array.rs
@@ -239,7 +239,7 @@ impl ListArray {
         Ok(arr.into())
     }
 
-    #[cfg(test)]
+    // Used for testing purposes
     pub fn from_slices(
         null_bitmap: &[bool],
         values: Vec<Option<ArrayImpl>>,
@@ -352,6 +352,25 @@ impl<'a> ListRef<'a> {
             ListRef::ValueRef { val } => val.values.iter().map(to_datum_ref).collect(),
         }
     }
+
+    pub fn value_at(&self, index: usize) -> Result<DatumRef<'a>> {
+        match self {
+            ListRef::Indexed { arr, .. } => {
+                if index <= arr.value.len() {
+                    Ok(arr.value.value_at(index - 1))
+                } else {
+                    Ok(None)
+                }
+            }
+            ListRef::ValueRef { val } => {
+                if let Some(datum) = val.values().iter().nth(index - 1) {
+                    Ok(to_datum_ref(datum))
+                } else {
+                    Ok(None)
+                }
+            }
+        }
+    }
 }
 
 impl Hash for ListRef<'_> {
diff --git a/src/expr/src/expr/build_expr_from_prost.rs b/src/expr/src/expr/build_expr_from_prost.rs
@@ -239,13 +239,84 @@ pub fn build_to_char_expr(prost: &ExprNode) -> Result<BoxedExpression> {
 mod tests {
     use std::vec;
 
+    use risingwave_common::array::{ArrayImpl, Utf8Array};
     use risingwave_pb::data::data_type::TypeName;
     use risingwave_pb::data::DataType as ProstDataType;
     use risingwave_pb::expr::expr_node::{RexNode, Type};
     use risingwave_pb::expr::{ConstantValue, ExprNode, FunctionCall, InputRefExpr};
 
     use super::*;
 
+    #[test]
+    fn test_array_access_expr() {
+        let values = FunctionCall {
+            children: vec![
+                ExprNode {
+                    expr_type: Type::ConstantValue as i32,
+                    return_type: Some(ProstDataType {
+                        type_name: TypeName::Varchar as i32,
+                        ..Default::default()
+                    }),
+                    rex_node: Some(RexNode::Constant(ConstantValue {
+                        body: "foo".as_bytes().to_vec(),
+                    })),
+                },
+                ExprNode {
+                    expr_type: Type::ConstantValue as i32,
+                    return_type: Some(ProstDataType {
+                        type_name: TypeName::Varchar as i32,
+                        ..Default::default()
+                    }),
+                    rex_node: Some(RexNode::Constant(ConstantValue {
+                        body: "bar".as_bytes().to_vec(),
+                    })),
+                },
+            ],
+        };
+        let array_index = FunctionCall {
+            children: vec![
+                ExprNode {
+                    expr_type: Type::Array as i32,
+                    return_type: Some(ProstDataType {
+                        type_name: TypeName::List as i32,
+                        field_type: vec![ProstDataType {
+                            type_name: TypeName::Varchar as i32,
+                            ..Default::default()
+                        }],
+                        ..Default::default()
+                    }),
+                    rex_node: Some(RexNode::FuncCall(values)),
+                },
+                ExprNode {
+                    expr_type: Type::ConstantValue as i32,
+                    return_type: Some(ProstDataType {
+                        type_name: TypeName::Int32 as i32,
+                        ..Default::default()
+                    }),
+                    rex_node: Some(RexNode::Constant(ConstantValue {
+                        body: vec![0, 0, 0, 1],
+                    })),
+                },
+            ],
+        };
+        let access = ExprNode {
+            expr_type: Type::ArrayAccess as i32,
+            return_type: Some(ProstDataType {
+                type_name: TypeName::Varchar as i32,
+                ..Default::default()
+            }),
+            rex_node: Some(RexNode::FuncCall(array_index)),
+        };
+        let expr = build_nullable_binary_expr_prost(&access);
+        assert!(expr.is_ok());
+
+        let res = expr.unwrap().eval(&DataChunk::new_dummy(1)).unwrap();
+        assert_eq!(
+            *res,
+            ArrayImpl::Utf8(Utf8Array::from_slice(&[Some("foo")]).unwrap())
+        );
+    }
+
     #[test]
     fn test_build_in_expr() {
         let input_ref = InputRefExpr { column_idx: 0 };
diff --git a/src/expr/src/expr/expr_binary_nullable.rs b/src/expr/src/expr/expr_binary_nullable.rs
@@ -14,13 +14,14 @@
 
 //! For expression that only accept two nullable arguments as input.
 
-use risingwave_common::array::{Array, BoolArray, Utf8Array};
+use risingwave_common::array::*;
 use risingwave_common::types::DataType;
 use risingwave_pb::expr::expr_node::Type;
 
 use super::BoxedExpression;
 use crate::expr::template::BinaryNullableExpression;
 use crate::for_all_cmp_variants;
+use crate::vector_op::array_access::array_access;
 use crate::vector_op::cmp::{general_is_distinct_from, str_is_distinct_from};
 use crate::vector_op::conjunction::{and, or};
 
@@ -62,6 +63,7 @@ pub fn new_nullable_binary_expr(
     r: BoxedExpression,
 ) -> BoxedExpression {
     match expr_type {
+        Type::ArrayAccess => build_array_access_expr(ret, l, r),
         Type::And => Box::new(
             BinaryNullableExpression::<BoolArray, BoolArray, BoolArray, _>::new(l, r, ret, and),
         ),
@@ -78,6 +80,43 @@ pub fn new_nullable_binary_expr(
     }
 }
 
+fn build_array_access_expr(
+    ret: DataType,
+    l: BoxedExpression,
+    r: BoxedExpression,
+) -> BoxedExpression {
+    macro_rules! array_access_expression {
+        ($array:ty) => {
+            Box::new(
+                BinaryNullableExpression::<ListArray, I32Array, $array, _>::new(
+                    l,
+                    r,
+                    ret,
+                    array_access,
+                ),
+            )
+        };
+    }
+
+    match ret {
+        DataType::Boolean => array_access_expression!(BoolArray),
+        DataType::Int16 => array_access_expression!(I16Array),
+        DataType::Int32 => array_access_expression!(I32Array),
+        DataType::Int64 => array_access_expression!(I64Array),
+        DataType::Float32 => array_access_expression!(F32Array),
+        DataType::Float64 => array_access_expression!(F64Array),
+        DataType::Decimal => array_access_expression!(DecimalArray),
+        DataType::Date => array_access_expression!(NaiveDateArray),
+        DataType::Varchar => array_access_expression!(Utf8Array),
+        DataType::Time => array_access_expression!(NaiveTimeArray),
+        DataType::Timestamp => array_access_expression!(NaiveDateTimeArray),
+        DataType::Timestampz => array_access_expression!(PrimitiveArray::<i64>),
+        DataType::Interval => array_access_expression!(IntervalArray),
+        DataType::Struct { .. } => array_access_expression!(StructArray),
+        DataType::List { .. } => array_access_expression!(ListArray),
+    }
+}
+
 pub fn new_distinct_from_expr(
     l: BoxedExpression,
     r: BoxedExpression,
diff --git a/src/expr/src/expr/mod.rs b/src/expr/src/expr/mod.rs
@@ -84,7 +84,7 @@ pub fn build_from_prost(prost: &ExprNode) -> Result<BoxedExpression> {
         Equal | NotEqual | LessThan | LessThanOrEqual | GreaterThan | GreaterThanOrEqual | Add
         | Subtract | Multiply | Divide | Modulus | Extract | RoundDigit | TumbleStart
         | Position => build_binary_expr_prost(prost),
-        And | Or | IsDistinctFrom => build_nullable_binary_expr_prost(prost),
+        And | Or | IsDistinctFrom | ArrayAccess => build_nullable_binary_expr_prost(prost),
         ToChar => build_to_char_expr(prost),
         Coalesce => CoalesceExpression::try_from(prost).map(Expression::boxed),
         Substr => build_substr_expr(prost),
diff --git a/src/expr/src/vector_op/array_access.rs b/src/expr/src/vector_op/array_access.rs
@@ -0,0 +1,111 @@
+// Copyright 2022 Singularity Data
+//
+// Licensed 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.
+
+use risingwave_common::array::ListRef;
+use risingwave_common::error::Result;
+use risingwave_common::types::{Scalar, ToOwnedDatum};
+
+#[inline(always)]
+pub fn array_access<T: Scalar>(l: Option<ListRef>, r: Option<i32>) -> Result<Option<T>> {
+    match (l, r) {
+        // index must be greater than 0 following a one-based numbering convention for arrays
+        (Some(list), Some(index)) if index > 0 => {
+            let datumref = list.value_at(index as usize)?;
+            if let Some(scalar) = datumref.to_owned_datum() {
+                Ok(Some(scalar.try_into()?))
+            } else {
+                Ok(None)
+            }
+        }
+        _ => Ok(None),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use risingwave_common::array::ListValue;
+    use risingwave_common::types::ScalarImpl;
+
+    use super::*;
+
+    #[test]
+    fn test_int32_array_access() {
+        let v1 = ListValue::new(vec![
+            Some(ScalarImpl::Int32(1)),
+            Some(ScalarImpl::Int32(2)),
+            Some(ScalarImpl::Int32(3)),
+        ]);
+        let l1 = ListRef::ValueRef { val: &v1 };
+
+        assert_eq!(array_access::<i32>(Some(l1), Some(1)), Ok(Some(1)));
+        assert_eq!(array_access::<i32>(Some(l1), Some(-1)), Ok(None));
+        assert_eq!(array_access::<i32>(Some(l1), Some(0)), Ok(None));
+        assert_eq!(array_access::<i32>(Some(l1), Some(4)), Ok(None));
+    }
+
+    #[test]
+    fn test_utf8_array_access() {
+        let v1 = ListValue::new(vec![
+            Some(ScalarImpl::Utf8("来自".into())),
+            Some(ScalarImpl::Utf8("foo".into())),
+            Some(ScalarImpl::Utf8("bar".into())),
+        ]);
+        let v2 = ListValue::new(vec![
+            Some(ScalarImpl::Utf8("fizz".into())),
+            Some(ScalarImpl::Utf8("荷兰".into())),
+            Some(ScalarImpl::Utf8("buzz".into())),
+        ]);
+        let v3 = ListValue::new(vec![None, None, Some(ScalarImpl::Utf8("的爱".into()))]);
+
+        let l1 = ListRef::ValueRef { val: &v1 };
+        let l2 = ListRef::ValueRef { val: &v2 };
+        let l3 = ListRef::ValueRef { val: &v3 };
+
+        assert_eq!(
+            array_access::<String>(Some(l1), Some(1)),
+            Ok(Some("来自".into()))
+        );
+        assert_eq!(
+            array_access::<String>(Some(l2), Some(2)),
+            Ok(Some("荷兰".into()))
+        );
+        assert_eq!(
+            array_access::<String>(Some(l3), Some(3)),
+            Ok(Some("的爱".into()))
+        );
+    }
+
+    #[test]
+    fn test_nested_array_access() {
+        let v = ListValue::new(vec![
+            Some(ScalarImpl::List(ListValue::new(vec![
+                Some(ScalarImpl::Utf8("foo".into())),
+                Some(ScalarImpl::Utf8("bar".into())),
+            ]))),
+            Some(ScalarImpl::List(ListValue::new(vec![
+                Some(ScalarImpl::Utf8("fizz".into())),
+                Some(ScalarImpl::Utf8("buzz".into())),
+            ]))),
+        ]);
+        let l = ListRef::ValueRef { val: &v };
+        assert_eq!(
+            array_access::<ListValue>(Some(l), Some(1)),
+            Ok(Some(ListValue::new(vec![
+                Some(ScalarImpl::Utf8("foo".into())),
+                Some(ScalarImpl::Utf8("bar".into())),
+            ])))
+        );
+    }
+}
diff --git a/src/expr/src/vector_op/mod.rs b/src/expr/src/vector_op/mod.rs
@@ -14,6 +14,7 @@
 
 pub mod agg;
 pub mod arithmetic_op;
+pub mod array_access;
 pub mod ascii;
 pub mod cast;
 pub mod cmp;
diff --git a/src/frontend/src/binder/expr/mod.rs b/src/frontend/src/binder/expr/mod.rs
@@ -51,6 +51,7 @@ impl Binder {
             Expr::BinaryOp { left, op, right } => self.bind_binary_op(*left, op, *right),
             Expr::Nested(expr) => self.bind_expr(*expr),
             Expr::Array(exprs) => self.bind_array(exprs),
+            Expr::ArrayIndex { obj, indexs } => self.bind_array_index(*obj, indexs),
             Expr::Function(f) => self.bind_function(f),
             // subquery
             Expr::Subquery(q) => self.bind_subquery_expr(*q, SubqueryKind::Scalar),
diff --git a/src/frontend/src/binder/expr/value.rs b/src/frontend/src/binder/expr/value.rs
diff --git a/src/frontend/test_runner/tests/testdata/array_access.yaml b/src/frontend/test_runner/tests/testdata/array_access.yaml
