diff --git a/e2e_test/batch/functions/array_concat.slt.part b/e2e_test/batch/functions/array_concat.slt.part new file mode 100644 index 0000000000000..c5a6129097f59 --- /dev/null +++ b/e2e_test/batch/functions/array_concat.slt.part @@ -0,0 +1,134 @@ +query T +select array_cat(array[66], array[123]); +---- +{66,123} + +query T +select array_cat(array[66], null::int[]); +---- +{66} + +query T +select array_cat(null::int[], array[123]); +---- +{123} + +query T +select array_cat(array[array[66]], array[233]); +---- +{{66},{233}} + +query T +select array_cat(array[array[66]], null::int[]); +---- +{{66}} + +query T +select array_cat(null::int[][], array[233]); +---- +{{233}} + +query T +select array_cat(null::int[][], null::int[]); +---- +NULL + +query T +select array_append(array[66], 123); +---- +{66,123} + +query T +select array_append(array[66], null::int); +---- +{66,NULL} + +query T +select array_append(null::int[], 233); +---- +{233} + +query T +select array_append(null::int[], null::int); +---- +{NULL} + +query T +select array_cat(array[233], array[array[66]]); +---- +{{233},{66}} + +query T +select array_cat(null::int[], array[array[66]]); +---- +{{66}} + +query T +select array_cat(array[233], null::int[][]); +---- +{{233}} + +query T +select array_cat(null::int[], null::int[][]); +---- +NULL + +query T +select array_prepend(123, array[66]); +---- +{123,66} + +query T +select array_prepend(null::int, array[66]); +---- +{NULL,66} + +query T +select array_prepend(233, null::int[]); +---- +{233} + +query T +select array_prepend(null::int, null::int[]); +---- +{NULL} + +query T +select array[1,2,3] || array[4,5,6]; +---- +{1,2,3,4,5,6} + +query T +select array[1,2,3] || 4; +---- +{1,2,3,4} + +query T +select 6 || array[7,8]; +---- +{6,7,8} + +query T +select array[array[1,2]] || array[array[3,4]]; +---- +{{1,2},{3,4}} + +query T +select array[array[1,2]] || array[3,4]; +---- +{{1,2},{3,4}} + +query T +select array[1,2] || array[array[3,4]]; +---- +{{1,2},{3,4}} + +query T +select '123'::varchar || array['abc']; +---- +{123,abc} + +query T +select array['abc'] || '123'::varchar; +---- +{abc,123} diff --git a/e2e_test/batch/types/array_ty.slt.part b/e2e_test/batch/types/array_ty.slt.part index 91f0911453e5e..2663bbde986cd 100644 --- a/e2e_test/batch/types/array_ty.slt.part +++ b/e2e_test/batch/types/array_ty.slt.part @@ -81,3 +81,8 @@ select max(ARRAY[1, v1*2]) from t; statement ok drop table t; + +# Now we don't disallow arrays with unmatching dimensions in multidimensional arrays. +# This is different from PostgreSQL, we may want to change this in the future. +statement ok +select array[array[1,2], array[3]]; diff --git a/proto/expr.proto b/proto/expr.proto index e34392e9a0982..c131fbe09d201 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -95,6 +95,10 @@ message ExprNode { ARRAY = 521; ARRAY_ACCESS = 522; ROW = 523; + // Array functions + ARRAY_CAT = 531; + ARRAY_APPEND = 532; + ARRAY_PREPEND = 533; // Search operator and Search ARGument SEARCH = 998; SARG = 999; diff --git a/src/expr/src/expr/expr_array_concat.rs b/src/expr/src/expr/expr_array_concat.rs new file mode 100644 index 0000000000000..9f5ee27d05801 --- /dev/null +++ b/src/expr/src/expr/expr_array_concat.rs @@ -0,0 +1,502 @@ +// 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 std::sync::Arc; + +use itertools::Itertools; +use risingwave_common::array::{ArrayRef, DataChunk, ListValue, Row}; +use risingwave_common::types::{to_datum_ref, DataType, Datum, DatumRef, ScalarRefImpl}; +use risingwave_pb::expr::expr_node::{RexNode, Type}; +use risingwave_pb::expr::ExprNode; + +use crate::expr::{build_from_prost as expr_build_from_prost, BoxedExpression, Expression}; +use crate::{bail, ensure, ExprError, Result}; + +#[derive(Debug, Copy, Clone)] +enum Operation { + ConcatArray, + AppendArray, + PrependArray, + AppendValue, + PrependValue, +} + +pub struct ArrayConcatExpression { + return_type: DataType, + left: BoxedExpression, + right: BoxedExpression, + op: Operation, + op_func: fn(DatumRef, DatumRef) -> Datum, +} + +impl std::fmt::Debug for ArrayConcatExpression { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ArrayConcatExpression") + .field("return_type", &self.return_type) + .field("left", &self.left) + .field("right", &self.right) + .field("op", &self.op) + .finish() + } +} + +impl ArrayConcatExpression { + fn new( + return_type: DataType, + left: BoxedExpression, + right: BoxedExpression, + op: Operation, + ) -> Self { + Self { + return_type, + left, + right, + op, + op_func: match op { + Operation::ConcatArray => Self::concat_array, + Operation::AppendArray => Self::append_array, + Operation::PrependArray => Self::prepend_array, + Operation::AppendValue => Self::append_value, + Operation::PrependValue => Self::prepend_value, + }, + } + } + + /// Concatenates two arrays with same data type. + /// The behavior is the same as PG. + /// + /// Examples: + /// - `select array_cat(array[66], array[123]);` => `[66,123]` + /// - `select array_cat(array[66], null::int[]);` => `[66]` + /// - `select array_cat(null::int[], array[123]);` => `[123]` + fn concat_array(left: DatumRef, right: DatumRef) -> Datum { + match (left, right) { + (None, right) => right.map(ScalarRefImpl::into_scalar_impl), + (left, None) => left.map(ScalarRefImpl::into_scalar_impl), + (Some(ScalarRefImpl::List(left)), Some(ScalarRefImpl::List(right))) => Some( + ListValue::new( + left.values_ref() + .into_iter() + .chain(right.values_ref().into_iter()) + .map(|x| x.map(ScalarRefImpl::into_scalar_impl)) + .collect(), + ) + .into(), + ), + _ => { + panic!("the operands must be two arrays with the same data type"); + } + } + } + + /// Appends an array as the back element of an array of array. + /// Note the behavior is slightly different from PG. + /// + /// Examples: + /// - `select array_cat(array[array[66]], array[233]);` => `[[66], [233]]` + /// - `select array_cat(array[array[66]], null::int[]);` => `[[66]]` ignore NULL, same as PG + /// - `select array_cat(null::int[][], array[233]);` => `[[233]]` different from PG + /// - `select array_cat(null::int[][], null::int[]);` => `null` same as PG + fn append_array(left: DatumRef, right: DatumRef) -> Datum { + match (left, right) { + (None, None) => None, + (None, right) => { + Some(ListValue::new(vec![right.map(ScalarRefImpl::into_scalar_impl)]).into()) + } + (left, None) => left.map(ScalarRefImpl::into_scalar_impl), + (Some(ScalarRefImpl::List(left)), right) => Some( + ListValue::new( + left.values_ref() + .into_iter() + .chain(std::iter::once(right)) + .map(|x| x.map(ScalarRefImpl::into_scalar_impl)) + .collect(), + ) + .into(), + ), + _ => { + panic!("the rhs must be compatible to append to lhs"); + } + } + } + + /// Appends a value as the back element of an array. + /// The behavior is the same as PG. + /// + /// Examples: + /// - `select array_append(array[66], 123);` => `[66, 123]` + /// - `select array_append(array[66], null::int);` => `[66, null]` + /// - `select array_append(null::int[], 233);` => `[233]` + /// - `select array_append(null::int[], null::int);` => `[null]` + fn append_value(left: DatumRef, right: DatumRef) -> Datum { + match (left, right) { + (None, right) => { + Some(ListValue::new(vec![right.map(ScalarRefImpl::into_scalar_impl)]).into()) + } + (Some(ScalarRefImpl::List(left)), right) => Some( + ListValue::new( + left.values_ref() + .into_iter() + .chain(std::iter::once(right)) + .map(|x| x.map(ScalarRefImpl::into_scalar_impl)) + .collect(), + ) + .into(), + ), + _ => { + panic!("the rhs must be compatible to append to lhs"); + } + } + } + + /// Prepends an array as the front element of an array of array. + /// Note the behavior is slightly different from PG. + /// + /// Examples: + /// - `select array_cat(array[233], array[array[66]]);` => `[[233], [66]]` + /// - `select array_cat(null::int[], array[array[66]]);` => `[[66]]` ignore NULL, same as PG + /// - `select array_cat(array[233], null::int[][]);` => `[[233]]` different from PG + /// - `select array_cat(null::int[], null::int[][]);` => `null` same as PG + fn prepend_array(left: DatumRef, right: DatumRef) -> Datum { + match (left, right) { + (None, None) => None, + (left, None) => { + Some(ListValue::new(vec![left.map(ScalarRefImpl::into_scalar_impl)]).into()) + } + (None, right) => right.map(ScalarRefImpl::into_scalar_impl), + (left, Some(ScalarRefImpl::List(right))) => Some( + ListValue::new( + std::iter::once(left) + .chain(right.values_ref().into_iter()) + .map(|x| x.map(ScalarRefImpl::into_scalar_impl)) + .collect(), + ) + .into(), + ), + _ => { + panic!("the lhs must be compatible to prepend to rhs"); + } + } + } + + /// Prepends a value as the front element of an array. + /// The behavior is the same as PG. + /// + /// Examples: + /// - `select array_prepend(123, array[66]);` => `[123, 66]` + /// - `select array_prepend(null::int, array[66]);` => `[null, 66]` + /// - `select array_prepend(233, null::int[]);` => `[233]` + /// - `select array_prepend(null::int, null::int[]);` => `[null]` + fn prepend_value(left: DatumRef, right: DatumRef) -> Datum { + match (left, right) { + (left, None) => { + Some(ListValue::new(vec![left.map(ScalarRefImpl::into_scalar_impl)]).into()) + } + (left, Some(ScalarRefImpl::List(right))) => Some( + ListValue::new( + std::iter::once(left) + .chain(right.values_ref().into_iter()) + .map(|x| x.map(ScalarRefImpl::into_scalar_impl)) + .collect(), + ) + .into(), + ), + _ => { + panic!("the lhs must be compatible to prepend to rhs"); + } + } + } + + fn evaluate(&self, left: DatumRef, right: DatumRef) -> Datum { + (self.op_func)(left, right) + } +} + +impl Expression for ArrayConcatExpression { + fn return_type(&self) -> DataType { + self.return_type.clone() + } + + fn eval(&self, input: &DataChunk) -> Result { + let left_array = self.left.eval_checked(input)?; + let right_array = self.right.eval_checked(input)?; + let mut builder = self + .return_type + .create_array_builder(left_array.len() + right_array.len()); + for (vis, (left, right)) in input + .vis() + .iter() + .zip_eq(left_array.iter().zip_eq(right_array.iter())) + { + if !vis { + builder.append_null()?; + } else { + builder.append_datum(&self.evaluate(left, right))?; + } + } + Ok(Arc::new(builder.finish()?)) + } + + fn eval_row(&self, input: &Row) -> Result { + let left_data = self.left.eval_row(input)?; + let right_data = self.right.eval_row(input)?; + Ok(self.evaluate(to_datum_ref(&left_data), to_datum_ref(&right_data))) + } +} + +impl<'a> TryFrom<&'a ExprNode> for ArrayConcatExpression { + type Error = ExprError; + + fn try_from(prost: &'a ExprNode) -> Result { + let RexNode::FuncCall(func_call_node) = prost.get_rex_node()? else { + bail!("expects a RexNode::FuncCall"); + }; + let children = func_call_node.get_children(); + ensure!(children.len() == 2); + let left = expr_build_from_prost(&children[0])?; + let right = expr_build_from_prost(&children[1])?; + let left_type = left.return_type(); + let right_type = right.return_type(); + let ret_type = DataType::from(prost.get_return_type()?); + let op = match prost.get_expr_type()? { + // the types are checked in frontend, so no need for type checking here + Type::ArrayCat => { + if left_type == right_type { + Operation::ConcatArray + } else if left_type == ret_type { + Operation::AppendArray + } else if right_type == ret_type { + Operation::PrependArray + } else { + bail!("function call node invalid"); + } + } + Type::ArrayAppend => Operation::AppendValue, + Type::ArrayPrepend => Operation::PrependValue, + _ => bail!("expects `ArrayCat`|`ArrayAppend`|`ArrayPrepend`"), + }; + Ok(Self::new(ret_type, left, right, op)) + } +} + +#[cfg(test)] +mod tests { + use risingwave_common::array::DataChunk; + use risingwave_common::types::ScalarImpl; + use risingwave_pb::expr::expr_node::{RexNode, Type as ProstType}; + use risingwave_pb::expr::{ConstantValue, ExprNode, FunctionCall}; + + use super::*; + use crate::expr::{Expression, LiteralExpression}; + + fn make_i64_expr_node(value: i64) -> ExprNode { + ExprNode { + expr_type: ProstType::ConstantValue as i32, + return_type: Some(DataType::Int64.to_protobuf()), + rex_node: Some(RexNode::Constant(ConstantValue { + body: value.to_be_bytes().to_vec(), + })), + } + } + + fn make_i64_array_expr_node(values: Vec) -> ExprNode { + ExprNode { + expr_type: ProstType::Array as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: values.into_iter().map(make_i64_expr_node).collect(), + })), + } + } + + fn make_i64_array_array_expr_node(values: Vec>) -> ExprNode { + ExprNode { + expr_type: ProstType::Array as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::List { + datatype: Box::new(DataType::Int64), + }), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: values.into_iter().map(make_i64_array_expr_node).collect(), + })), + } + } + + #[test] + fn test_array_concat_try_from() { + { + let left = make_i64_array_expr_node(vec![42]); + let right = make_i64_array_expr_node(vec![43]); + let expr = ExprNode { + expr_type: ProstType::ArrayCat as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: vec![left, right], + })), + }; + assert!(ArrayConcatExpression::try_from(&expr).is_ok()); + } + + { + let left = make_i64_array_array_expr_node(vec![vec![42]]); + let right = make_i64_array_array_expr_node(vec![vec![43]]); + let expr = ExprNode { + expr_type: ProstType::ArrayCat as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: vec![left, right], + })), + }; + assert!(ArrayConcatExpression::try_from(&expr).is_ok()); + } + + { + let left = make_i64_array_expr_node(vec![42]); + let right = make_i64_expr_node(43); + let expr = ExprNode { + expr_type: ProstType::ArrayAppend as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: vec![left, right], + })), + }; + assert!(ArrayConcatExpression::try_from(&expr).is_ok()); + } + + { + let left = make_i64_array_array_expr_node(vec![vec![42]]); + let right = make_i64_array_expr_node(vec![43]); + let expr = ExprNode { + expr_type: ProstType::ArrayAppend as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: vec![left, right], + })), + }; + assert!(ArrayConcatExpression::try_from(&expr).is_ok()); + } + + { + let left = make_i64_expr_node(43); + let right = make_i64_array_expr_node(vec![42]); + let expr = ExprNode { + expr_type: ProstType::ArrayPrepend as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: vec![left, right], + })), + }; + assert!(ArrayConcatExpression::try_from(&expr).is_ok()); + } + + { + let left = make_i64_array_expr_node(vec![43]); + let right = make_i64_array_array_expr_node(vec![vec![42]]); + let expr = ExprNode { + expr_type: ProstType::ArrayPrepend as i32, + return_type: Some( + DataType::List { + datatype: Box::new(DataType::Int64), + } + .to_protobuf(), + ), + rex_node: Some(RexNode::FuncCall(FunctionCall { + children: vec![left, right], + })), + }; + assert!(ArrayConcatExpression::try_from(&expr).is_ok()); + } + } + + fn make_i64_array_expr(values: Vec) -> BoxedExpression { + LiteralExpression::new( + DataType::List { + datatype: Box::new(DataType::Int64), + }, + Some(ListValue::new(values.into_iter().map(|x| Some(x.into())).collect()).into()), + ) + .boxed() + } + + #[test] + fn test_array_concat_array_of_primitives() { + let left = make_i64_array_expr(vec![42]); + let right = make_i64_array_expr(vec![43, 44]); + let expr = ArrayConcatExpression::new( + DataType::List { + datatype: Box::new(DataType::Int64), + }, + left, + right, + Operation::ConcatArray, + ); + + let chunk = DataChunk::new_dummy(4) + .with_visibility([true, false, true, true].into_iter().collect()); + let expected_array = Some(ScalarImpl::List(ListValue::new(vec![ + Some(42i64.into()), + Some(43i64.into()), + Some(44i64.into()), + ]))); + let expected = vec![ + expected_array.clone(), + None, + expected_array.clone(), + expected_array, + ]; + let actual = expr + .eval(&chunk) + .unwrap() + .iter() + .map(|v| v.map(|s| s.into_scalar_impl())) + .collect_vec(); + assert_eq!(actual, expected); + } + + // More test cases are in e2e tests. +} diff --git a/src/expr/src/expr/mod.rs b/src/expr/src/expr/mod.rs index 3996345c500de..efb26c7022cf6 100644 --- a/src/expr/src/expr/mod.rs +++ b/src/expr/src/expr/mod.rs @@ -15,6 +15,7 @@ mod agg; pub mod build_expr_from_prost; pub mod data_types; +mod expr_array_concat; mod expr_binary_bytes; pub mod expr_binary_nonnull; pub mod expr_binary_nullable; @@ -46,6 +47,7 @@ use risingwave_pb::expr::ExprNode; use super::Result; use crate::expr::build_expr_from_prost::*; +use crate::expr::expr_array_concat::ArrayConcatExpression; use crate::expr::expr_case::CaseExpression; use crate::expr::expr_coalesce::CoalesceExpression; use crate::expr::expr_concat_ws::ConcatWsExpression; @@ -133,6 +135,12 @@ pub fn build_from_prost(prost: &ExprNode) -> Result { Array => NestedConstructExpression::try_from(prost).map(Expression::boxed), Row => NestedConstructExpression::try_from(prost).map(Expression::boxed), RegexpMatch => RegexpMatchExpression::try_from(prost).map(Expression::boxed), + ArrayCat | ArrayAppend | ArrayPrepend => { + // Now we implement these three functions as a single expression for the + // sake of simplicity. If performance matters at some time, we can split + // the implementation to improve performance. + ArrayConcatExpression::try_from(prost).map(Expression::boxed) + } Vnode => VnodeExpression::try_from(prost).map(Expression::boxed), _ => Err(ExprError::UnsupportedFunction(format!( "{:?}", diff --git a/src/frontend/src/binder/expr/binary_op.rs b/src/frontend/src/binder/expr/binary_op.rs index a4f66215fa8f4..a189415e34b6e 100644 --- a/src/frontend/src/binder/expr/binary_op.rs +++ b/src/frontend/src/binder/expr/binary_op.rs @@ -71,26 +71,22 @@ impl Binder { /// Bind `||`. Based on the types of the inputs, this can be string concat or array concat. fn bind_concat_op(&mut self, left: ExprImpl, right: ExprImpl) -> Result { - let types = [left.return_type(), right.return_type()]; - - let has_string = types.iter().any(|t| matches!(t, DataType::Varchar)); - let has_array = types.iter().any(|t| matches!(t, DataType::List { .. })); - - // StringConcat - if has_string && !has_array { - Ok(FunctionCall::new(ExprType::ConcatOp, vec![left, right])?.into()) - } - // ArrayConcat - else if has_array { - Err(ErrorCode::NotImplemented("array concat operator".into(), None.into()).into()) - } - // Invalid types - else { - Err(ErrorCode::BindError(format!( - "operator does not exist: {:?} || {:?}", - &types[0], &types[1] - )) - .into()) - } + let func_type = match (left.return_type(), right.return_type()) { + // array concatenation + (DataType::List { .. }, DataType::List { .. }) => ExprType::ArrayCat, + (DataType::List { .. }, _) => ExprType::ArrayAppend, + (_, DataType::List { .. }) => ExprType::ArrayPrepend, + // string concatenation + (DataType::Varchar, _) | (_, DataType::Varchar) => ExprType::ConcatOp, + // invalid + (left_type, right_type) => { + return Err(ErrorCode::BindError(format!( + "operator does not exist: {} || {}", + left_type, right_type + )) + .into()); + } + }; + Ok(FunctionCall::new(func_type, vec![left, right])?.into()) } } diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index e75e5883a6dc2..238fdd14393ec 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -148,6 +148,10 @@ impl Binder { "octet_length" => ExprType::OctetLength, "bit_length" => ExprType::BitLength, "regexp_match" => ExprType::RegexpMatch, + // array + "array_cat" => ExprType::ArrayCat, + "array_append" => ExprType::ArrayAppend, + "array_prepend" => ExprType::ArrayPrepend, // System information operations. "pg_typeof" if inputs.len() == 1 => { let input = &inputs[0]; diff --git a/src/frontend/src/expr/type_inference/func.rs b/src/frontend/src/expr/type_inference/func.rs index 5f896c16e4f3e..84ef51303e160 100644 --- a/src/frontend/src/expr/type_inference/func.rs +++ b/src/frontend/src/expr/type_inference/func.rs @@ -75,6 +75,17 @@ macro_rules! ensure_arity { .into()); } }; + ($func:literal, $num:literal == | $inputs:ident |) => { + if !($inputs.len() == $num) { + return Err(ErrorCode::BindError(format!( + "Function `{}` takes {} arguments ({} given)", + $func, + $num, + $inputs.len(), + )) + .into()); + } + }; } /// Special exprs that cannot be handled by [`infer_type_name`] and [`FuncSigMap`] are handled here. @@ -143,6 +154,72 @@ fn infer_type_for_special( datatype: Box::new(DataType::Varchar), })) } + ExprType::ArrayCat => { + ensure_arity!("array_cat", 2 == | inputs |); + let left_type = inputs[0].return_type(); + let right_type = inputs[1].return_type(); + let return_type = match (&left_type, &right_type) { + ( + DataType::List { + datatype: left_elem_type, + }, + DataType::List { + datatype: right_elem_type, + }, + ) => { + if **left_elem_type == **right_elem_type || **left_elem_type == right_type { + Some(left_type.clone()) + } else if left_type == **right_elem_type { + Some(right_type.clone()) + } else { + None + } + } + _ => None, + }; + Ok(Some(return_type.ok_or_else(|| { + ErrorCode::BindError(format!( + "Cannot concatenate {} and {}", + left_type, right_type + )) + })?)) + } + ExprType::ArrayAppend => { + ensure_arity!("array_append", 2 == | inputs |); + let left_type = inputs[0].return_type(); + let right_type = inputs[1].return_type(); + let return_type = match (&left_type, &right_type) { + (DataType::List { .. }, DataType::List { .. }) => None, + ( + DataType::List { + datatype: left_elem_type, + }, + _, // non-array + ) if **left_elem_type == right_type => Some(left_type.clone()), + _ => None, + }; + Ok(Some(return_type.ok_or_else(|| { + ErrorCode::BindError(format!("Cannot append {} to {}", right_type, left_type)) + })?)) + } + ExprType::ArrayPrepend => { + ensure_arity!("array_prepend", 2 == | inputs |); + let left_type = inputs[0].return_type(); + let right_type = inputs[1].return_type(); + let return_type = match (&left_type, &right_type) { + (DataType::List { .. }, DataType::List { .. }) => None, + ( + _, // non-array + DataType::List { + datatype: right_elem_type, + }, + ) if left_type == **right_elem_type => Some(right_type.clone()), + _ => None, + }; + Ok(Some(return_type.ok_or_else(|| { + ErrorCode::BindError(format!("Cannot prepend {} to {}", left_type, right_type)) + })?)) + } ExprType::Vnode => { ensure_arity!("vnode", 1 <= | inputs |); Ok(Some(DataType::Int16)) diff --git a/src/frontend/test_runner/tests/testdata/array.yaml b/src/frontend/test_runner/tests/testdata/array.yaml index f68de6bfa6a3c..62e945b13eb68 100644 --- a/src/frontend/test_runner/tests/testdata/array.yaml +++ b/src/frontend/test_runner/tests/testdata/array.yaml @@ -26,3 +26,64 @@ logical_plan: | LogicalProject { exprs: [Array(null:Varchar)] } LogicalValues { rows: [[]], schema: Schema { fields: [] } } +- sql: | + select array_cat(array[66], array[123]); + logical_plan: | + LogicalProject { exprs: [ArrayCat(Array(66:Int32), Array(123:Int32))] } + LogicalValues { rows: [[]], schema: Schema { fields: [] } } + batch_plan: | + BatchProject { exprs: [ArrayCat(Array(66:Int32), Array(123:Int32))] } + BatchValues { rows: [[]] } +- sql: | + select array_cat(array[array[66]], array[233]); + logical_plan: | + LogicalProject { exprs: [ArrayCat(Array(Array(66:Int32)), Array(233:Int32))] } + LogicalValues { rows: [[]], schema: Schema { fields: [] } } + batch_plan: | + BatchProject { exprs: [ArrayCat(Array(Array(66:Int32)), Array(233:Int32))] } + BatchValues { rows: [[]] } +- sql: | + select array_cat(array[233], array[array[66]]); + logical_plan: | + LogicalProject { exprs: [ArrayCat(Array(233:Int32), Array(Array(66:Int32)))] } + LogicalValues { rows: [[]], schema: Schema { fields: [] } } + batch_plan: | + BatchProject { exprs: [ArrayCat(Array(233:Int32), Array(Array(66:Int32)))] } + BatchValues { rows: [[]] } +- sql: | + select array_cat(array[233], array[array[array[66]]]); + binder_error: 'Bind error: Cannot concatenate integer[] and integer[][][]' +- sql: | + select array_cat(array[233], 123); + binder_error: 'Bind error: Cannot concatenate integer[] and integer' +- sql: | + select array_cat(123, array[233]); + binder_error: 'Bind error: Cannot concatenate integer and integer[]' +- sql: | + select array_append(array[66], 123); + logical_plan: | + LogicalProject { exprs: [ArrayAppend(Array(66:Int32), 123:Int32)] } + LogicalValues { rows: [[]], schema: Schema { fields: [] } } + batch_plan: | + BatchProject { exprs: [ArrayAppend(Array(66:Int32), 123:Int32)] } + BatchValues { rows: [[]] } +- sql: | + select array_append(123, 234); + binder_error: 'Bind error: Cannot append integer to integer' +- sql: | + select array_append(array[array[66]], array[233]); + binder_error: 'Bind error: Cannot append integer[] to integer[][]' +- sql: | + select array_prepend(123, array[66]); + logical_plan: | + LogicalProject { exprs: [ArrayPrepend(123:Int32, Array(66:Int32))] } + LogicalValues { rows: [[]], schema: Schema { fields: [] } } + batch_plan: | + BatchProject { exprs: [ArrayPrepend(123:Int32, Array(66:Int32))] } + BatchValues { rows: [[]] } +- sql: | + select array_prepend(123, 234); + binder_error: 'Bind error: Cannot prepend integer to integer' +- sql: | + select array_prepend(array[233], array[array[66]]); + binder_error: 'Bind error: Cannot prepend integer[] to integer[][]'