[feat](array) support trim_array function

be/src/exprs/function/array/function_array_pop.cpp

@@ -23,7 +23,9 @@
 #include <ostream>
 #include <utility>
 
+#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"
@@ -104,9 +106,74 @@ class FunctionArrayPopfront : public FunctionArrayPop<FunctionArrayPopfront> {
     static constexpr int start_offset = 2;
 };
 
+class FunctionArrayTrim : public IFunction {
+public:
+    static constexpr auto name = "trim_array";
+    static FunctionPtr create() { return std::make_shared<FunctionArrayTrim>(); }
+
+    String get_name() const override { return name; }
+
+    bool is_variadic() const override { return false; }
+
+    size_t get_number_of_arguments() const override { return 2; }
+
+    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
+        DCHECK(arguments[0]->get_primitive_type() == TYPE_ARRAY)
+                << "First argument for function: " << name
+                << " should be DataTypeArray but it has type " << arguments[0]->get_name() << ".";
+        DCHECK(arguments[1]->get_primitive_type() == TYPE_BIGINT)
+                << "Second argument for function: " << name << " should be BigInt but it has type "
+                << arguments[1]->get_name() << ".";
+        return arguments[0];
+    }
+
+    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
+                        uint32_t result, size_t input_rows_count) const override {
+        auto array_column =
+                block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
+        auto size_column =
+                block.get_by_position(arguments[1]).column->convert_to_full_column_if_const();
+
+        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()));
+        }
+
+        auto length_column = ColumnInt64::create();
+        length_column->reserve(input_rows_count);
+        const auto& sizes = assert_cast<const ColumnInt64&>(*size_column).get_data();
+        for (size_t row = 0; row < input_rows_count; ++row) {
+            const auto size = sizes[row];
+            const size_t offset = (*src.offsets_ptr)[row - 1];
+            const size_t cardinality = (*src.offsets_ptr)[row] - offset;
+            if (UNLIKELY(size < 0)) {
+                return Status::InvalidArgument("size must not be negative: {}", size);
+            }
+            if (UNLIKELY(static_cast<size_t>(size) > cardinality)) {
+                return Status::InvalidArgument("size must not exceed array cardinality {}: {}",
+                                               cardinality, size);
+            }
+            length_column->insert_value(static_cast<Int64>(cardinality) - size);
+        }
+
+        const bool nested_is_nullable = src.nested_nullmap_data != nullptr;
+        ColumnArrayMutableData dst = create_mutable_data(src.nested_col.get(), nested_is_nullable);
+        dst.offsets_ptr->reserve(input_rows_count);
+        auto offset_column = ColumnInt64::create(input_rows_count, 1);
+        slice_array(dst, src, *offset_column, length_column.get());
+
+        block.replace_by_position(result, assemble_column_array(dst));
+        return Status::OK();
+    }
+};
+
 void register_function_array_pop(SimpleFunctionFactory& factory) {
     factory.register_function<FunctionArrayPopback>();
     factory.register_function<FunctionArrayPopfront>();
+    factory.register_function<FunctionArrayTrim>();
 }
 
 } // namespace doris

be/test/exprs/function/function_array_trim_test.cpp

@@ -0,0 +1,50 @@
+// 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 <memory>
+#include <string>
+
+#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"
+
+namespace doris {
+
+TEST(FunctionArrayTrimTest, all_argument_combinations) {
+    const std::string function_name = "trim_array";
+    const InputArgTypeSet input_types = {{PrimitiveType::TYPE_ARRAY, PrimitiveType::TYPE_INT},
+                                         {PrimitiveType::TYPE_BIGINT}};
+
+    TestArray empty;
+    TestArray values = {Int32(1), Int32(2), Int32(3), Int32(4)};
+    TestArray values_with_null = {Int32(1), Null(), Int32(3)};
+    DataSet data_set = {
+            {{AnyType(values), Int64(0)}, AnyType(values)},
+            {{AnyType(values), Int64(2)}, AnyType(TestArray {Int32(1), Int32(2)})},
+            {{AnyType(values), Int64(4)}, AnyType(empty)},
+            {{AnyType(empty), Int64(0)}, AnyType(empty)},
+            {{AnyType(values_with_null), Int64(1)}, AnyType(TestArray {Int32(1), Null()})},
+            {{Null(), Int64(0)}, Null()},
+            {{AnyType(values), Null()}, Null()}};
+
+    auto result_nested_type = make_nullable(std::make_shared<DataTypeInt32>());
+    check_function_all_arg_comb<DataTypeArray, true>(function_name, input_types, data_set,
+                                                     result_nested_type);
+}
+
+} // namespace doris

be/test/exprs/function/function_test_util.h

@@ -79,6 +79,9 @@ using Row = std::pair<InputCell, Expect>;
 using DataSet = std::vector<Row>;
 // to represent Array<Int64>: {PrimitiveType::TYPE_ARRAY, PrimitiveType::TYPE_BIGINT}
 using InputTypeSet = std::vector<AnyType>;
+// Each entry represents one logical argument. This form supports parametric types, for example:
+// {{PrimitiveType::TYPE_ARRAY, PrimitiveType::TYPE_INT}, {PrimitiveType::TYPE_BIGINT}}.
+using InputArgTypeSet = std::vector<InputTypeSet>;
 
 struct Nullable {
     PrimitiveType tp;
@@ -322,7 +325,7 @@ template <typename ResultType, bool ResultNullable = false, bool datetime_is_str
 Status check_function(const std::string& func_name, const InputTypeSet& input_types,
                       const DataSet& data_set, int result_scale = -1, int result_precision = -1,
                       bool expect_execute_fail = false, bool expect_result_ne = false,
-                      bool is_strict_mode = false) {
+                      bool is_strict_mode = false, DataTypePtr result_nested_type = nullptr) {
     TestCaseInfo::arg_size = static_cast<int>(input_types.size());
     TestCaseInfo::func_call_index++;
     // 1.0 create data type
@@ -384,6 +387,10 @@ Status check_function(const std::string& func_name, const InputTypeSet& input_ty
             }
             return ResultNullable ? make_nullable(std::make_shared<ResultType>(real_scale))
                                   : std::make_shared<ResultType>(real_scale);
+        } else if constexpr (std::is_same_v<ResultType, DataTypeArray>) {
+            EXPECT_NE(result_nested_type, nullptr);
+            DataTypePtr array_type = std::make_shared<DataTypeArray>(result_nested_type);
+            return ResultNullable ? make_nullable(array_type) : array_type;
         } else {
             return ResultNullable ? make_nullable(std::make_shared<ResultType>())
                                   : std::make_shared<ResultType>();
@@ -394,8 +401,14 @@ Status check_function(const std::string& func_name, const InputTypeSet& input_ty
     assert(func.get() != nullptr);
 
     // this may be useless now. for some type like array, it's wrong. TODO: need more details explainations
-    auto fn_ctx_return = get_return_type_descriptor<ResultType>(std::max(0, result_scale),
-                                                                std::max(0, result_precision));
+    auto fn_ctx_return = [&]() {
+        if constexpr (std::is_same_v<ResultType, DataTypeArray>) {
+            return remove_nullable(return_type);
+        } else {
+            return get_return_type_descriptor<ResultType>(std::max(0, result_scale),
+                                                          std::max(0, result_precision));
+        }
+    }();
 
     FunctionUtils fn_utils(fn_ctx_return, arg_types, is_strict_mode);
     auto* fn_ctx = fn_utils.get_fn_ctx();
@@ -423,7 +436,9 @@ Status check_function(const std::string& func_name, const InputTypeSet& input_ty
 
     // 3.0. create expected result column in block
     DataTypePtr result_type_ptr;
-    if constexpr (IsDataTypeDecimal<ResultType>) { // decimal
+    if constexpr (std::is_same_v<ResultType, DataTypeArray>) {
+        result_type_ptr = return_type;
+    } else if constexpr (IsDataTypeDecimal<ResultType>) { // decimal
         result_type_ptr = ResultNullable
                                   ? make_nullable(std::make_shared<ResultType>(result_precision,
                                                                                result_scale))
@@ -532,4 +547,41 @@ void check_function_all_arg_comb(const std::string& func_name, const InputTypeSe
         }
     }
 }
+
+// Variant for parametric argument or result types. Each entry in base_arg_types is one logical
+// argument and may contain multiple type descriptors, such as ARRAY followed by its item type.
+template <typename ReturnType, bool nullable = false>
+void check_function_all_arg_comb(const std::string& func_name,
+                                 const InputArgTypeSet& base_arg_types, const DataSet& data_set,
+                                 DataTypePtr result_nested_type) {
+    TestCaseInfo::func_call_index++;
+    const size_t arg_cnt = base_arg_types.size();
+    for (int combination = 0; combination < (1 << arg_cnt); ++combination) {
+        InputTypeSet input_types;
+        for (size_t arg = 0; arg < arg_cnt; ++arg) {
+            const bool is_const = (1 << arg) & combination;
+            const auto& arg_types = base_arg_types[arg];
+            DCHECK(!arg_types.empty());
+            const auto base_type = any_cast<PrimitiveType>(arg_types.front());
+            input_types.emplace_back(is_const ? AnyType(Consted {base_type}) : arg_types.front());
+            input_types.insert(input_types.end(), arg_types.begin() + 1, arg_types.end());
+        }
+
+        TestCaseInfo::arg_const_info = combination;
+        TestCaseInfo::error_line_number = -1;
+        if (combination != 0) {
+            for (const auto& line : data_set) {
+                TestCaseInfo::func_call_index--;
+                static_cast<void>(check_function<ReturnType, nullable>(
+                        func_name, input_types, DataSet {line}, -1, -1, false, false, false,
+                        result_nested_type));
+            }
+        } else {
+            TestCaseInfo::func_call_index--;
+            static_cast<void>(check_function<ReturnType, nullable>(func_name, input_types, data_set,
+                                                                   -1, -1, false, false, false,
+                                                                   result_nested_type));
+        }
+    }
+}
 } // namespace doris

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"),

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.
      */

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<FunctionSignature> 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<Expression> 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, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
+        return visitor.visitTrimArray(this, context);
+    }
+
+    @Override
+    public List<FunctionSignature> getSignatures() {
+        return SIGNATURES;
+    }
+}

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);
     }