[feature](function) Support Spark/Hive-compatible stack table function

be/src/exprs/function/function_fake.cpp

@@ -181,6 +181,15 @@ struct FunctionEsquery {
     static std::string get_error_msg() { return "esquery only supported on es table"; }
 };
 
+class FunctionStack : public FunctionFake<UDTFImpl> {
+public:
+    static FunctionPtr create() { return std::make_shared<FunctionStack>(); }
+
+    bool skip_return_type_check() const override { return true; }
+
+    ColumnNumbers get_arguments_that_are_always_constant() const override { return {0}; }
+};
+
 template <typename FunctionImpl>
 void register_function(SimpleFunctionFactory& factory, const std::string& name) {
     factory.register_function<FunctionFake<FunctionImpl>>(name);
@@ -254,6 +263,7 @@ void register_table_function_with_impl(SimpleFunctionFactory& factory, const std
 
 void register_function_fake(SimpleFunctionFactory& factory) {
     register_function<FunctionEsquery>(factory, "esquery");
+    factory.register_function<FunctionStack>("stack");
 
     register_table_function_expand_outer<FunctionExplodeV2>(factory, "explode");
     register_table_alternative_function_expand_outer<FunctionExplode>(factory, "explode");

be/src/exprs/table_function/table_function_factory.cpp

@@ -35,6 +35,7 @@
 #include "exprs/table_function/vexplode_numbers.h"
 #include "exprs/table_function/vexplode_v2.h"
 #include "exprs/table_function/vjson_each.h"
+#include "exprs/table_function/vstack.h"
 
 namespace doris {
 
@@ -53,6 +54,7 @@ const std::unordered_map<std::string, std::function<std::unique_ptr<TableFunctio
                 {"json_each", TableFunctionCreator<VJsonEachTableFn> {}},
                 {"json_each_text", TableFunctionCreator<VJsonEachTextTableFn> {}},
                 {"posexplode", TableFunctionCreator<VExplodeV2TableFunction> {}},
+                {"stack", TableFunctionCreator<VStackTableFunction> {}},
                 {"explode", TableFunctionCreator<VExplodeV2TableFunction> {}},
                 {"explode_variant_array_old", TableFunctionCreator<VExplodeTableFunction>()},
                 {"explode_old", TableFunctionCreator<VExplodeTableFunction> {}}};

be/src/exprs/table_function/vstack.cpp

@@ -0,0 +1,135 @@
+// 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 "exprs/table_function/vstack.h"
+
+#include <algorithm>
+
+#include "core/assert_cast.h"
+#include "core/block/block.h"
+#include "core/column/column_const.h"
+#include "core/column/column_nullable.h"
+#include "core/column/column_struct.h"
+#include "exprs/vexpr.h"
+
+namespace doris {
+
+VStackTableFunction::VStackTableFunction() {
+    _fn_name = "stack";
+}
+
+Status VStackTableFunction::process_init(Block* block, RuntimeState* /*state*/) {
+    const auto& children = _expr_context->root()->children();
+    DORIS_CHECK_GE(children.size(), 2);
+
+    int column_index = -1;
+    RETURN_IF_ERROR(children[0]->execute(_expr_context.get(), block, &column_index));
+    const auto& num_rows_column = block->get_by_position(column_index).column;
+    DORIS_CHECK(is_column_const(*num_rows_column));
+    const auto num_rows = assert_cast<const ColumnConst&>(*num_rows_column).get_int(0);
+    DORIS_CHECK_GT(num_rows, 0);
+    _num_rows = static_cast<size_t>(num_rows);
+    _num_fields = (children.size() - 2) / _num_rows + 1;
+
+    _value_columns.clear();
+    _value_columns.reserve(children.size() - 1);
+    for (size_t i = 1; i < children.size(); ++i) {
+        RETURN_IF_ERROR(children[i]->execute(_expr_context.get(), block, &column_index));
+        const auto& value_column = block->get_by_position(column_index).column;
+        const auto& [column, is_const] = unpack_if_const(value_column);
+        _value_columns.emplace_back(ValueColumn {.column = column, .is_const = is_const});
+    }
+    return Status::OK();
+}
+
+void VStackTableFunction::process_row(size_t row_idx) {
+    TableFunction::process_row(row_idx);
+    _row_idx = row_idx;
+    _cur_size = static_cast<int64_t>(_num_rows);
+}
+
+void VStackTableFunction::process_close() {
+    _value_columns.clear();
+    _row_idx = 0;
+    _num_rows = 0;
+    _num_fields = 0;
+}
+
+void VStackTableFunction::_insert_value(IColumn& destination, const IColumn& source,
+                                        size_t source_row) {
+    auto* nullable_destination = check_and_get_column<ColumnNullable>(&destination);
+    DORIS_CHECK(nullable_destination != nullptr);
+
+    if (const auto* nullable_source = check_and_get_column<ColumnNullable>(&source)) {
+        nullable_destination->get_nested_column().insert_from(nullable_source->get_nested_column(),
+                                                              source_row);
+        nullable_destination->get_null_map_data().push_back(
+                nullable_source->get_null_map_data()[source_row]);
+    } else {
+        nullable_destination->get_nested_column().insert_from(source, source_row);
+        nullable_destination->get_null_map_data().push_back(0);
+    }
+}
+
+void VStackTableFunction::_insert_output_row(MutableColumnPtr& column, size_t output_row) const {
+    IColumn* output = column.get();
+    if (_num_fields == 1) {
+        const size_t value_index = output_row;
+        if (value_index < _value_columns.size()) {
+            const auto& value_column = _value_columns[value_index];
+            _insert_value(*output, *value_column.column, value_column.is_const ? 0 : _row_idx);
+        } else {
+            output->insert_default();
+        }
+        return;
+    }
+
+    if (_is_nullable) {
+        auto& nullable_output = assert_cast<ColumnNullable&>(*output);
+        nullable_output.get_null_map_data().push_back(0);
+        output = &nullable_output.get_nested_column();
+    }
+
+    auto& struct_output = assert_cast<ColumnStruct&>(*output);
+    for (size_t field_index = 0; field_index < _num_fields; ++field_index) {
+        const size_t value_index = output_row * _num_fields + field_index;
+        auto& field = struct_output.get_column(field_index);
+        if (value_index < _value_columns.size()) {
+            const auto& value_column = _value_columns[value_index];
+            _insert_value(field, *value_column.column, value_column.is_const ? 0 : _row_idx);
+        } else {
+            field.insert_default();
+        }
+    }
+}
+
+void VStackTableFunction::get_same_many_values(MutableColumnPtr& column, int length) {
+    for (int i = 0; i < length; ++i) {
+        _insert_output_row(column, static_cast<size_t>(_cur_offset));
+    }
+}
+
+int VStackTableFunction::get_value(MutableColumnPtr& column, int max_step) {
+    max_step = std::min(max_step, static_cast<int>(_cur_size - _cur_offset));
+    for (int i = 0; i < max_step; ++i) {
+        _insert_output_row(column, static_cast<size_t>(_cur_offset + i));
+    }
+    forward(max_step);
+    return max_step;
+}
+
+} // namespace doris

be/src/exprs/table_function/vstack.h

@@ -0,0 +1,54 @@
+// 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.
+
+#pragma once
+
+#include <cstddef>
+#include <vector>
+
+#include "exprs/table_function/table_function.h"
+
+namespace doris {
+
+class VStackTableFunction : public TableFunction {
+    ENABLE_FACTORY_CREATOR(VStackTableFunction);
+
+public:
+    VStackTableFunction();
+
+    Status process_init(Block* block, RuntimeState* state) override;
+    void process_row(size_t row_idx) override;
+    void process_close() override;
+    void get_same_many_values(MutableColumnPtr& column, int length) override;
+    int get_value(MutableColumnPtr& column, int max_step) override;
+
+private:
+    struct ValueColumn {
+        ColumnPtr column;
+        bool is_const;
+    };
+
+    void _insert_output_row(MutableColumnPtr& column, size_t output_row) const;
+    static void _insert_value(IColumn& destination, const IColumn& source, size_t source_row);
+
+    std::vector<ValueColumn> _value_columns;
+    size_t _row_idx = 0;
+    size_t _num_rows = 0;
+    size_t _num_fields = 0;
+};
+
+} // namespace doris

be/test/exprs/function/function_test_util.cpp

@@ -599,7 +599,8 @@ static Block* create_block_from_inputset(const InputTypeSet& input_types,
                                     ? ((DataTypeNullable*)(desc.data_type.get()))->get_nested_type()
                                     : desc.data_type;
 
-            for (int r = 0; r < row_size; r++) {
+            const size_t rows_to_insert = desc.is_const ? 1 : row_size;
+            for (size_t r = 0; r < rows_to_insert; r++) {
                 if (!insert_cell(column, type_ptr, input_set[r][i * input_col_size + j])) {
                     return nullptr;
                 }

be/test/exprs/function/table_function_test.cpp

@@ -31,6 +31,7 @@
 #include "exprs/table_function/vexplode_numbers.h"
 #include "exprs/table_function/vexplode_v2.h"
 #include "exprs/table_function/vjson_each.h"
+#include "exprs/table_function/vstack.h"
 #include "testutil/any_type.h"
 #include "util/jsonb_parser_simd.h"
 #include "util/jsonb_utils.h"
@@ -304,6 +305,58 @@ TEST_F(TableFunctionTest, vexplode_v2_two_param) {
     }
 }
 
+TEST_F(TableFunctionTest, vstack) {
+    init_expr_context(4);
+    VStackTableFunction stack;
+    stack.set_expr_context(_ctx);
+
+    {
+        InputTypeSet input_types = {ConstedNotnull {PrimitiveType::TYPE_INT},
+                                    PrimitiveType::TYPE_INT, PrimitiveType::TYPE_INT,
+                                    PrimitiveType::TYPE_INT};
+        InputDataSet input_set = {{Int32(2), Int32(1), Int32(2), Int32(3)}};
+
+        InputTypeSet output_types = {PrimitiveType::TYPE_STRUCT, PrimitiveType::TYPE_INT,
+                                     PrimitiveType::TYPE_INT};
+        InputDataSet output_set = {{{TestArray {Int32(1), Int32(2)}}},
+                                   {{TestArray {Int32(3), Null()}}}};
+
+        check_vec_table_function(&stack, input_types, input_set, output_types, output_set, false);
+        check_vec_table_function(&stack, input_types, input_set, output_types, output_set, true);
+    }
+
+    {
+        InputTypeSet input_types = {
+                ConstedNotnull {PrimitiveType::TYPE_INT}, ConstedNotnull {PrimitiveType::TYPE_INT},
+                ConstedNotnull {PrimitiveType::TYPE_INT}, ConstedNotnull {PrimitiveType::TYPE_INT}};
+        InputDataSet input_set = {{Int32(2), Int32(1), Int32(2), Int32(3)},
+                                  {Int32(2), Int32(1), Int32(2), Int32(3)}};
+
+        InputTypeSet output_types = {PrimitiveType::TYPE_STRUCT, PrimitiveType::TYPE_INT,
+                                     PrimitiveType::TYPE_INT};
+        InputDataSet output_set = {{{TestArray {Int32(1), Int32(2)}}},
+                                   {{TestArray {Int32(3), Null()}}},
+                                   {{TestArray {Int32(1), Int32(2)}}},
+                                   {{TestArray {Int32(3), Null()}}}};
+
+        check_vec_table_function(&stack, input_types, input_set, output_types, output_set, false);
+        check_vec_table_function(&stack, input_types, input_set, output_types, output_set, true);
+    }
+
+    {
+        InputTypeSet input_types = {ConstedNotnull {PrimitiveType::TYPE_INT},
+                                    PrimitiveType::TYPE_INT, PrimitiveType::TYPE_INT,
+                                    PrimitiveType::TYPE_INT};
+        InputDataSet input_set = {{Int32(4), Int32(1), Null(), Int32(3)}};
+
+        InputTypeSet output_types = {PrimitiveType::TYPE_INT};
+        InputDataSet output_set = {{Int32(1)}, {Null()}, {Int32(3)}, {Null()}};
+
+        check_vec_table_function(&stack, input_types, input_set, output_types, output_set, false);
+        check_vec_table_function(&stack, input_types, input_set, output_types, output_set, true);
+    }
+}
+
 TEST_F(TableFunctionTest, vexplode_numbers) {
     init_expr_context(1);
     VExplodeNumbersTableFunction tfn;

fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinTableGeneratingFunctions.java

@@ -44,6 +44,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.generator.JsonEachTextOuter;
 import org.apache.doris.nereids.trees.expressions.functions.generator.PosExplode;
 import org.apache.doris.nereids.trees.expressions.functions.generator.PosExplodeOuter;
+import org.apache.doris.nereids.trees.expressions.functions.generator.Stack;
 import org.apache.doris.nereids.trees.expressions.functions.generator.Unnest;
 
 import com.google.common.collect.ImmutableList;
@@ -88,6 +89,7 @@ public class BuiltinTableGeneratingFunctions implements FunctionHelper {
             tableGenerating(JsonEachTextOuter.class, "json_each_text_outer"),
             tableGenerating(PosExplode.class, "posexplode"),
             tableGenerating(PosExplodeOuter.class, "posexplode_outer"),
+            tableGenerating(Stack.class, "stack"),
             tableGenerating(Unnest.class, "unnest")
     );
 
@@ -99,7 +101,7 @@ public class BuiltinTableGeneratingFunctions implements FunctionHelper {
             .add("explode_json_array_json_outer").add("explode_split").add("explode_split_outer")
             .add("json_each").add("json_each_outer")
             .add("json_each_text").add("json_each_text_outer")
-            .add("posexplode").add("posexplode_outer").build();
+            .add("posexplode").add("posexplode_outer").add("stack").build();
 
     public Set<String> getReturnManyColumnFunctions() {
         return RETURN_MULTI_COLUMNS_FUNCTIONS;

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java

@@ -59,6 +59,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.FunctionBuilder;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AnyValue;
+import org.apache.doris.nereids.trees.expressions.functions.generator.Stack;
 import org.apache.doris.nereids.trees.expressions.functions.generator.TableGeneratingFunction;
 import org.apache.doris.nereids.trees.expressions.functions.generator.Unnest;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.Coalesce;
@@ -439,12 +440,29 @@ private LogicalPlan bindGenerate(MatchingContext<LogicalGenerate<Plan>> ctx) {
             // 2. the expandColumnsAlias is empty, we should use origin boundSlot
             if (generate.getExpandColumnAlias() != null && i < generate.getExpandColumnAlias().size()
                     && !CollectionUtils.isEmpty(generate.getExpandColumnAlias().get(i))) {
-                if (boundSlot.getDataType() instanceof StructType
-                        && generate.getExpandColumnAlias().get(i).size() > 1) {
+                int aliasCount = generate.getExpandColumnAlias().get(i).size();
+                boolean shouldExpandStruct = boundSlot.getDataType() instanceof StructType && aliasCount > 1;
+                if (boundGenerator instanceof Stack) {
+                    int outputColumnCount = ((Stack) boundGenerator).getOutputColumnCount();
+                    if (aliasCount != outputColumnCount) {
+                        throw new AnalysisException(String.format(
+                                "table %s has %d columns available but %d columns specified",
+                                slot.getQualifier().get(0), outputColumnCount, aliasCount));
+                    }
+                    shouldExpandStruct = outputColumnCount > 1;
+                }
+                if (shouldExpandStruct) {
+                    Preconditions.checkState(boundSlot.getDataType() instanceof StructType,
+                            "multi-column generator output must use a struct carrier");
                     // if the alias is not empty, we should bind it with struct_element as child expr with alias
                     // element_at(#expand_col#k, #k) as #k
                     // element_at(#expand_col#v, #v) as #v
                     List<StructField> fields = ((StructType) boundSlot.getDataType()).getFields();
+                    if (aliasCount != fields.size()) {
+                        throw new AnalysisException(String.format(
+                                "table %s has %d columns available but %d columns specified",
+                                slot.getQualifier().get(0), fields.size(), aliasCount));
+                    }
                     for (int idx = 0; idx < fields.size(); ++idx) {
                         expandAlias.add(new Alias(new ElementAt(
                                 boundSlot, new StringLiteral(fields.get(idx).getName())),