diff --git a/be/src/exprs/function/array/function_array_pop.cpp b/be/src/exprs/function/array/function_array_pop.cpp index aa415a59449327..90c95ba6b6b5e0 100644 --- a/be/src/exprs/function/array/function_array_pop.cpp +++ b/be/src/exprs/function/array/function_array_pop.cpp @@ -19,17 +19,23 @@ #include #include +#include #include #include #include +#include "common/compiler_util.h" #include "common/status.h" +#include "core/assert_cast.h" #include "core/block/block.h" #include "core/block/column_numbers.h" #include "core/block/column_with_type_and_name.h" #include "core/column/column.h" +#include "core/column/column_array.h" +#include "core/column/column_nullable.h" #include "core/column/column_vector.h" #include "core/data_type/data_type.h" +#include "core/data_type/data_type_nullable.h" #include "core/types.h" #include "exprs/aggregate/aggregate_function.h" #include "exprs/function/array/function_array_utils.h" @@ -104,9 +110,124 @@ class FunctionArrayPopfront : public FunctionArrayPop { static constexpr int start_offset = 2; }; +class FunctionArrayTrim : public IFunction { +public: + static constexpr auto name = "trim_array"; + static FunctionPtr create() { return std::make_shared(); } + + String get_name() const override { return name; } + + bool is_variadic() const override { return false; } + + size_t get_number_of_arguments() const override { return 2; } + + bool use_default_implementation_for_nulls() const override { return false; } + + DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { + const auto array_type = remove_nullable(arguments[0]); + DCHECK(array_type->get_primitive_type() == TYPE_ARRAY) + << "First argument for function: " << name + << " should be DataTypeArray but it has type " << arguments[0]->get_name() << "."; + DCHECK(remove_nullable(arguments[1])->get_primitive_type() == TYPE_BIGINT) + << "Second argument for function: " << name << " should be BigInt but it has type " + << arguments[1]->get_name() << "."; + if (arguments[0]->is_nullable() || arguments[1]->is_nullable()) { + return make_nullable(array_type); + } + return array_type; + } + + Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, + uint32_t result, size_t input_rows_count) const override { + const auto& [array_column, array_is_const] = + unpack_if_const(block.get_by_position(arguments[0]).column); + const auto& [size_column, size_is_const] = + unpack_if_const(block.get_by_position(arguments[1]).column); + + ColumnArrayExecutionData src; + if (!extract_column_array_info(*array_column, src)) { + return Status::RuntimeError( + fmt::format("execute failed, unsupported types for function {}({}, {})", + get_name(), block.get_by_position(arguments[0]).type->get_name(), + block.get_by_position(arguments[1]).type->get_name())); + } + + const UInt8* size_null_map = nullptr; + const IColumn* size_data_column = size_column.get(); + if (const auto* nullable_size = check_and_get_column(size_data_column)) { + size_null_map = nullable_size->get_null_map_data().data(); + size_data_column = &nullable_size->get_nested_column(); + } + const auto& sizes = assert_cast(*size_data_column).get_data(); + + auto result_array = ColumnArray::create(src.array_col->get_data_ptr()->clone_empty(), + ColumnArray::ColumnOffsets::create()); + auto& result_data = result_array->get_data(); + auto& result_offsets = result_array->get_offsets(); + result_offsets.resize(input_rows_count); + + auto result_null_map = ColumnUInt8::create(input_rows_count, 0); + auto& result_null_map_data = result_null_map->get_data(); + size_t result_offset = 0; + for (size_t row = 0; row < input_rows_count; ++row) { + const size_t array_row = index_check_const(row, array_is_const); + const size_t size_row = index_check_const(row, size_is_const); + const bool is_null = (src.array_nullmap_data && src.array_nullmap_data[array_row]) || + (size_null_map && size_null_map[size_row]); + result_null_map_data[row] = is_null; + if (is_null) { + result_offsets[row] = result_offset; + continue; + } + + const auto size = sizes[size_row]; + const size_t offset = (*src.offsets_ptr)[array_row - 1]; + const size_t cardinality = (*src.offsets_ptr)[array_row] - offset; + if (UNLIKELY(size < 0)) { + return Status::InvalidArgument("size must not be negative: {}", size); + } + if (UNLIKELY(static_cast(size) > cardinality)) { + return Status::InvalidArgument("size must not exceed array cardinality {}: {}", + cardinality, size); + } + + const size_t keep = cardinality - size; + if (UNLIKELY(result_offset > std::numeric_limits::max() - keep)) { + return Status::InvalidArgument("result array size overflows"); + } + result_offset += keep; + result_offsets[row] = result_offset; + } + + result_data.reserve(result_offset); + for (size_t row = 0; row < input_rows_count; ++row) { + if (result_null_map_data[row]) { + continue; + } + + const size_t array_row = index_check_const(row, array_is_const); + const size_t offset = (*src.offsets_ptr)[array_row - 1]; + const size_t previous_result_offset = row == 0 ? 0 : result_offsets[row - 1]; + const size_t keep = result_offsets[row] - previous_result_offset; + if (keep > 0) { + result_data.insert_range_from(src.array_col->get_data(), offset, keep); + } + } + + if (block.get_by_position(result).type->is_nullable()) { + block.replace_by_position(result, ColumnNullable::create(std::move(result_array), + std::move(result_null_map))); + } else { + block.replace_by_position(result, std::move(result_array)); + } + return Status::OK(); + } +}; + void register_function_array_pop(SimpleFunctionFactory& factory) { factory.register_function(); factory.register_function(); + factory.register_function(); } } // namespace doris diff --git a/be/test/exprs/function/function_array_trim_test.cpp b/be/test/exprs/function/function_array_trim_test.cpp new file mode 100644 index 00000000000000..ead082e9ed11b4 --- /dev/null +++ b/be/test/exprs/function/function_array_trim_test.cpp @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +#include +#include +#include + +#include "core/block/block.h" +#include "core/column/column_const.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_nullable.h" +#include "core/data_type/data_type_number.h" +#include "exprs/function/function_test_util.h" +#include "exprs/function/simple_function_factory.h" + +namespace doris { + +static void check_array_trim_case(const DataSet& data_set, bool const_array, bool const_size) { + auto element_type = make_nullable(std::make_shared()); + DataTypePtr array_type = + make_nullable(std::make_shared(std::move(element_type))); + DataTypePtr size_type = make_nullable(std::make_shared()); + const size_t row_size = data_set.size(); + + MutableColumnPtr array_column = array_type->create_column(); + MutableColumnPtr size_column = size_type->create_column(); + for (size_t row = 0; row < row_size; ++row) { + if (!const_array || row == 0) { + ASSERT_TRUE(insert_cell(array_column, array_type, data_set[row].first[0])); + } + if (!const_size || row == 0) { + ASSERT_TRUE(insert_cell(size_column, size_type, data_set[row].first[1])); + } + } + + if (const_array) { + array_column = ColumnConst::create(std::move(array_column), row_size); + } + if (const_size) { + size_column = ColumnConst::create(std::move(size_column), row_size); + } + + Block block; + block.insert({std::move(array_column), array_type, "array"}); + block.insert({std::move(size_column), size_type, "size"}); + + FunctionBasePtr function = SimpleFunctionFactory::instance().get_function( + "trim_array", block.get_columns_with_type_and_name(), array_type); + ASSERT_NE(function, nullptr); + + std::vector> constant_columns(2); + if (const_array) { + constant_columns[0] = std::make_shared(block.get_by_position(0).column); + } + if (const_size) { + constant_columns[1] = std::make_shared(block.get_by_position(1).column); + } + + FunctionUtils function_utils(array_type, {array_type, size_type}, false); + auto* function_context = function_utils.get_fn_ctx(); + function_context->set_constant_cols(constant_columns); + ASSERT_TRUE(function->open(function_context, FunctionContext::FRAGMENT_LOCAL).ok()); + ASSERT_TRUE(function->open(function_context, FunctionContext::THREAD_LOCAL).ok()); + + block.insert({nullptr, array_type, "result"}); + const auto result = block.columns() - 1; + ASSERT_TRUE(function->execute(function_context, block, {0, 1}, result, row_size).ok()); + + static_cast(function->close(function_context, FunctionContext::THREAD_LOCAL)); + static_cast(function->close(function_context, FunctionContext::FRAGMENT_LOCAL)); + + MutableColumnPtr expected_column = array_type->create_column(); + for (const auto& row : data_set) { + ASSERT_TRUE(insert_cell(expected_column, array_type, row.second)); + } + const auto result_column = + block.get_by_position(result).column->convert_to_full_column_if_const(); + for (size_t row = 0; row < row_size; ++row) { + EXPECT_EQ(0, result_column->compare_at(row, row, *expected_column, 1)) + << "row " << row << ", result: " << array_type->to_string(*result_column, row) + << ", expected: " << array_type->to_string(*expected_column, row); + } +} + +TEST(FunctionArrayTrimTest, all_argument_combinations) { + const TestArray empty; + const TestArray values = {Int32(1), Int32(2), Int32(3), Int32(4)}; + const TestArray short_values = {Int32(5), Int32(6)}; + const TestArray values_with_null = {Int32(1), Null(), Int32(3)}; + + check_array_trim_case( + {{{AnyType(values), Int64(0)}, AnyType(values)}, + {{AnyType(values), Int64(2)}, AnyType(TestArray {Int32(1), Int32(2)})}, + {{AnyType(empty), Int64(0)}, AnyType(empty)}, + {{AnyType(values_with_null), Int64(1)}, AnyType(TestArray {Int32(1), Null()})}, + {{Null(), Int64(1)}, Null()}, + {{AnyType(values), Null()}, Null()}}, + false, false); + + check_array_trim_case({{{AnyType(values), Int64(0)}, AnyType(values)}, + {{AnyType(values), Int64(2)}, AnyType(TestArray {Int32(1), Int32(2)})}, + {{AnyType(values), Int64(4)}, AnyType(empty)}}, + true, false); + + check_array_trim_case( + {{{AnyType(values), Int64(1)}, AnyType(TestArray {Int32(1), Int32(2), Int32(3)})}, + {{AnyType(short_values), Int64(1)}, AnyType(TestArray {Int32(5)})}, + {{AnyType(values_with_null), Int64(1)}, AnyType(TestArray {Int32(1), Null()})}, + {{Null(), Int64(1)}, Null()}}, + false, true); + + check_array_trim_case({{{AnyType(values), Int64(2)}, AnyType(TestArray {Int32(1), Int32(2)})}}, + true, true); +} + +} // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java index e55e1008762e68..6289835dd22d43 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java @@ -551,6 +551,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformValues; import org.apache.doris.nereids.trees.expressions.functions.scalar.Translate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Trim; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TrimArray; import org.apache.doris.nereids.trees.expressions.functions.scalar.TrimIn; import org.apache.doris.nereids.trees.expressions.functions.scalar.Truncate; import org.apache.doris.nereids.trees.expressions.functions.scalar.TryParseToVariant; @@ -1149,6 +1150,7 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(ToSeconds.class, "to_seconds"), scalar(Translate.class, "translate"), scalar(Trim.class, "trim"), + scalar(TrimArray.class, "trim_array"), scalar(TrimIn.class, "trim_in"), scalar(Truncate.class, "truncate"), scalar(Unhex.class, "unhex"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmetic.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmetic.java index c1201807f487af..38fda5dd233f26 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmetic.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmetic.java @@ -50,6 +50,22 @@ public static Expression cardinality(MapLiteral map) { return new BigIntLiteral(map.getValue().size()); } + /** Remove a number of elements from the end of an array. */ + @ExecFunction(name = "trim_array") + public static Expression trimArray(ArrayLiteral array, BigIntLiteral size) { + long trimSize = size.getValue(); + int cardinality = array.getValue().size(); + if (trimSize < 0) { + throw new AnalysisException("size must not be negative: " + trimSize); + } + if (trimSize > cardinality) { + throw new AnalysisException("size must not exceed array cardinality " + + cardinality + ": " + trimSize); + } + return new ArrayLiteral(array.getValue().subList(0, cardinality - (int) trimSize), + array.getDataType()); + } + /** * Compute the cross product between two 3D float arrays. */ diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TrimArray.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TrimArray.java new file mode 100644 index 00000000000000..9f89ba82c37c5a --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TrimArray.java @@ -0,0 +1,70 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.coercion.AnyDataType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * Remove the specified number of elements from the end of an array. + */ +public class TrimArray extends ScalarFunction + implements BinaryExpression, ExplicitlyCastableSignature, PropagateNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.retArgType(0) + .args(ArrayType.of(AnyDataType.INSTANCE_WITHOUT_INDEX), BigIntType.INSTANCE) + ); + + public TrimArray(Expression array, Expression size) { + super("trim_array", array, size); + } + + private TrimArray(ScalarFunctionParams functionParams) { + super(functionParams); + } + + @Override + public TrimArray withChildren(List children) { + Preconditions.checkArgument(children.size() == 2, + "trim_array accepts 2 arguments, but got %s (%s)", children.size(), children); + return new TrimArray(getFunctionParams(children)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitTrimArray(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index ec86766df02f2e..a130ee7be37f73 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -568,6 +568,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.TransformValues; import org.apache.doris.nereids.trees.expressions.functions.scalar.Translate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Trim; +import org.apache.doris.nereids.trees.expressions.functions.scalar.TrimArray; import org.apache.doris.nereids.trees.expressions.functions.scalar.TrimIn; import org.apache.doris.nereids.trees.expressions.functions.scalar.Truncate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Uncompress; @@ -2660,6 +2661,10 @@ default R visitTrim(Trim trim, C context) { return visitScalarFunction(trim, context); } + default R visitTrimArray(TrimArray trimArray, C context) { + return visitScalarFunction(trimArray, context); + } + default R visitTrimIn(TrimIn trimIn, C context) { return visitScalarFunction(trimIn, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmeticTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmeticTest.java new file mode 100644 index 00000000000000..386c9c11572682 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/executable/ArrayArithmeticTest.java @@ -0,0 +1,58 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +package org.apache.doris.nereids.trees.expressions.functions.executable; + +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; +import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class ArrayArithmeticTest { + + @Test + void testTrimArray() { + ArrayLiteral input = new ArrayLiteral(ImmutableList.of( + new IntegerLiteral(1), new IntegerLiteral(2), new IntegerLiteral(3))); + + ArrayLiteral trimmed = (ArrayLiteral) ArrayArithmetic.trimArray(input, new BigIntLiteral(1)); + Assertions.assertEquals(ImmutableList.of(new IntegerLiteral(1), new IntegerLiteral(2)), + trimmed.getValue()); + Assertions.assertEquals(input.getDataType(), trimmed.getDataType()); + + ArrayLiteral unchanged = (ArrayLiteral) ArrayArithmetic.trimArray(input, new BigIntLiteral(0)); + Assertions.assertEquals(input.getValue(), unchanged.getValue()); + + ArrayLiteral empty = (ArrayLiteral) ArrayArithmetic.trimArray(input, new BigIntLiteral(3)); + Assertions.assertTrue(empty.getValue().isEmpty()); + Assertions.assertEquals(input.getDataType(), empty.getDataType()); + } + + @Test + void testTrimArrayRejectsInvalidSize() { + ArrayLiteral input = new ArrayLiteral(ImmutableList.of(new IntegerLiteral(1))); + + Assertions.assertThrows(AnalysisException.class, + () -> ArrayArithmetic.trimArray(input, new BigIntLiteral(-1))); + Assertions.assertThrows(AnalysisException.class, + () -> ArrayArithmetic.trimArray(input, new BigIntLiteral(2))); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TrimArrayTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TrimArrayTest.java new file mode 100644 index 00000000000000..fb7909c6993d84 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TrimArrayTest.java @@ -0,0 +1,75 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral; +import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TrimArrayTest { + + @Test + void testPropertiesAndWithChildren() { + ArrayLiteral array = new ArrayLiteral(ImmutableList.of( + new IntegerLiteral(1), new IntegerLiteral(2), new IntegerLiteral(3))); + BigIntLiteral size = new BigIntLiteral(1); + TrimArray trimArray = new TrimArray(array, size); + + Assertions.assertEquals("trim_array", trimArray.getName()); + Assertions.assertEquals(2, trimArray.arity()); + Assertions.assertEquals(TrimArray.SIGNATURES, trimArray.getSignatures()); + Assertions.assertEquals(array, trimArray.child(0)); + Assertions.assertEquals(size, trimArray.child(1)); + Assertions.assertFalse(trimArray.nullable()); + + ArrayLiteral newArray = new ArrayLiteral(ImmutableList.of(new IntegerLiteral(4))); + BigIntLiteral newSize = new BigIntLiteral(0); + TrimArray rewritten = trimArray.withChildren(ImmutableList.of(newArray, newSize)); + Assertions.assertNotSame(trimArray, rewritten); + Assertions.assertEquals(newArray, rewritten.child(0)); + Assertions.assertEquals(newSize, rewritten.child(1)); + + Assertions.assertThrows(IllegalArgumentException.class, + () -> trimArray.withChildren(ImmutableList.of(array))); + } + + @Test + void testVisitorDispatch() { + TrimArray trimArray = new TrimArray( + new ArrayLiteral(ImmutableList.of(new IntegerLiteral(1))), new BigIntLiteral(0)); + ExpressionVisitor visitor = new ExpressionVisitor() { + @Override + public String visit(Expression expression, Void context) { + return "expression"; + } + + @Override + public String visitTrimArray(TrimArray function, Void context) { + return function.getName(); + } + }; + + Assertions.assertEquals("trim_array", trimArray.accept(visitor, null)); + } +} diff --git a/regression-test/data/query_p0/sql_functions/array_functions/test_trim_array.out b/regression-test/data/query_p0/sql_functions/array_functions/test_trim_array.out new file mode 100644 index 00000000000000..89e0babf57a310 --- /dev/null +++ b/regression-test/data/query_p0/sql_functions/array_functions/test_trim_array.out @@ -0,0 +1,69 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !trim_two -- +[1, 2] + +-- !trim_zero -- +[1, 2, 3, 4] + +-- !trim_one -- +[1, 2, 3] + +-- !trim_all -- +[] + +-- !trim_string -- +["a", "b", "c"] + +-- !trim_null_element -- +["a", "b", null] + +-- !trim_nested -- +[[1, 2, 3]] + +-- !trim_empty -- +[] + +-- !trim_boolean -- +[1, 0] + +-- !trim_tinyint -- +[-128, 0] + +-- !trim_bigint -- +[-9223372036854775808, 0] + +-- !trim_double -- +[-1.7976931348623157e+308, 0] + +-- !trim_decimal -- +[-99999999.99, 0.00] + +-- !trim_date -- +["0000-01-01", "2024-02-29"] + +-- !trim_null_array -- +\N + +-- !trim_null_array_invalid_size -- +\N + +-- !trim_null_size -- +\N + +-- !trim_columns -- +1 [1, 2] +2 [5, 6] +3 [] +4 \N +5 \N + +-- !trim_const_array -- +1 [1, 2] +2 [1, 2, 3, 4] +3 [1, 2, 3, 4] +4 [1, 2, 3, 4] + +-- !trim_const_size -- +1 [1, 2, 3] +2 [5] +4 \N diff --git a/regression-test/suites/query_p0/sql_functions/array_functions/test_trim_array.groovy b/regression-test/suites/query_p0/sql_functions/array_functions/test_trim_array.groovy new file mode 100644 index 00000000000000..8f62ec03bbafd0 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/array_functions/test_trim_array.groovy @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +suite("test_trim_array") { + testFoldConst("select trim_array([1, 2, 3, 4], 2)") + testFoldConst("select trim_array([1, 2, 3, 4], 0)") + testFoldConst("select trim_array([1, 2, 3, 4], 1)") + testFoldConst("select trim_array([1, 2, 3, 4], 4)") + testFoldConst("select trim_array(['a', 'b', 'c', 'd'], 1)") + testFoldConst("select trim_array(['a', 'b', null, 'd'], 1)") + testFoldConst("select trim_array([[1, 2, 3], [4, 5, 6]], 1)") + testFoldConst("select trim_array(cast([] as array), 0)") + testFoldConst("select trim_array(cast([true, false, true] as array), 1)") + testFoldConst("select trim_array(cast([-128, 0, 127] as array), 1)") + testFoldConst("select trim_array(cast([-9223372036854775808, 0, 9223372036854775807] as array), 1)") + testFoldConst("select trim_array(cast([-1.7976931348623157E308, 0.0, 1.7976931348623157E308] as array), 1)") + testFoldConst("select trim_array(cast([-99999999.99, 0.00, 99999999.99] as array), 1)") + testFoldConst("select trim_array(cast(['0000-01-01', '2024-02-29', '9999-12-31'] as array), 1)") + testFoldConst("select trim_array(cast(null as array), 0)") + testFoldConst("select trim_array(cast(null as array), 9223372036854775807)") + testFoldConst("select trim_array([1, 2, 3], cast(null as bigint))") + + qt_trim_two "select trim_array([1, 2, 3, 4], 2)" + qt_trim_zero "select trim_array([1, 2, 3, 4], 0)" + qt_trim_one "select trim_array([1, 2, 3, 4], 1)" + qt_trim_all "select trim_array([1, 2, 3, 4], 4)" + qt_trim_string "select trim_array(['a', 'b', 'c', 'd'], 1)" + qt_trim_null_element "select trim_array(['a', 'b', null, 'd'], 1)" + qt_trim_nested "select trim_array([[1, 2, 3], [4, 5, 6]], 1)" + qt_trim_empty "select trim_array(cast([] as array), 0)" + qt_trim_boolean "select trim_array(cast([true, false, true] as array), 1)" + qt_trim_tinyint "select trim_array(cast([-128, 0, 127] as array), 1)" + qt_trim_bigint "select trim_array(cast([-9223372036854775808, 0, 9223372036854775807] as array), 1)" + qt_trim_double "select trim_array(cast([-1.7976931348623157E308, 0.0, 1.7976931348623157E308] as array), 1)" + qt_trim_decimal "select trim_array(cast([-99999999.99, 0.00, 99999999.99] as array), 1)" + qt_trim_date "select trim_array(cast(['0000-01-01', '2024-02-29', '9999-12-31'] as array), 1)" + qt_trim_null_array "select trim_array(cast(null as array), 0)" + qt_trim_null_array_invalid_size "select trim_array(cast(null as array), 9223372036854775807)" + qt_trim_null_size "select trim_array([1, 2, 3], cast(null as bigint))" + + test { + sql "select trim_array([1, 2, 3, 4], 5)" + exception "size must not exceed array cardinality 4: 5" + } + test { + sql "select trim_array([1, 2, 3, 4], -1)" + exception "size must not be negative: -1" + } + test { + sql "select trim_array([1, 2, 3, 4], 9223372036854775807)" + exception "size must not exceed array cardinality 4: 9223372036854775807" + } + test { + sql "select trim_array([1, 2, 3, 4], -9223372036854775808)" + exception "size must not be negative: -9223372036854775808" + } + + sql "drop table if exists trim_array_test" + sql """ + create table trim_array_test ( + id int, + items array, + trim_size bigint + ) distributed by hash(id) buckets 1 + properties('replication_num' = '1') + """ + sql """ + insert into trim_array_test values + (1, [1, 2, 3, 4], 2), + (2, [5, 6], 0), + (3, [], 0), + (4, null, 0), + (5, [7, 8], null) + """ + order_qt_trim_columns "select id, trim_array(items, trim_size) from trim_array_test" + order_qt_trim_const_array """ + select id, trim_array([1, 2, 3, 4], trim_size) + from trim_array_test where trim_size is not null order by id + """ + order_qt_trim_const_size "select id, trim_array(items, 1) from trim_array_test where id in (1, 2, 4) order by id" +}