diff --git a/be/src/exprs/function/function_fake.cpp b/be/src/exprs/function/function_fake.cpp index 0966ed9af613c1..5e048d15feed71 100644 --- a/be/src/exprs/function/function_fake.cpp +++ b/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 { +public: + static FunctionPtr create() { return std::make_shared(); } + + bool skip_return_type_check() const override { return true; } + + ColumnNumbers get_arguments_that_are_always_constant() const override { return {0}; } +}; + template void register_function(SimpleFunctionFactory& factory, const std::string& name) { factory.register_function>(name); @@ -254,6 +263,7 @@ void register_table_function_with_impl(SimpleFunctionFactory& factory, const std void register_function_fake(SimpleFunctionFactory& factory) { register_function(factory, "esquery"); + factory.register_function("stack"); register_table_function_expand_outer(factory, "explode"); register_table_alternative_function_expand_outer(factory, "explode"); diff --git a/be/src/exprs/table_function/table_function_factory.cpp b/be/src/exprs/table_function/table_function_factory.cpp index f4c705c183c18a..57bd58d2b25c66 100644 --- a/be/src/exprs/table_function/table_function_factory.cpp +++ b/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 {}}, {"json_each_text", TableFunctionCreator {}}, {"posexplode", TableFunctionCreator {}}, + {"stack", TableFunctionCreator {}}, {"explode", TableFunctionCreator {}}, {"explode_variant_array_old", TableFunctionCreator()}, {"explode_old", TableFunctionCreator {}}}; diff --git a/be/src/exprs/table_function/vstack.cpp b/be/src/exprs/table_function/vstack.cpp new file mode 100644 index 00000000000000..a0eb24417e6265 --- /dev/null +++ b/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 + +#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(*num_rows_column).get_int(0); + DORIS_CHECK_GT(num_rows, 0); + _num_rows = static_cast(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(_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(&destination); + DORIS_CHECK(nullable_destination != nullptr); + + if (const auto* nullable_source = check_and_get_column(&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(*output); + nullable_output.get_null_map_data().push_back(0); + output = &nullable_output.get_nested_column(); + } + + auto& struct_output = assert_cast(*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(_cur_offset)); + } +} + +int VStackTableFunction::get_value(MutableColumnPtr& column, int max_step) { + max_step = std::min(max_step, static_cast(_cur_size - _cur_offset)); + for (int i = 0; i < max_step; ++i) { + _insert_output_row(column, static_cast(_cur_offset + i)); + } + forward(max_step); + return max_step; +} + +} // namespace doris diff --git a/be/src/exprs/table_function/vstack.h b/be/src/exprs/table_function/vstack.h new file mode 100644 index 00000000000000..c71f0438469edd --- /dev/null +++ b/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 +#include + +#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 _value_columns; + size_t _row_idx = 0; + size_t _num_rows = 0; + size_t _num_fields = 0; +}; + +} // namespace doris diff --git a/be/test/exprs/function/function_test_util.cpp b/be/test/exprs/function/function_test_util.cpp index b37e452b69a52a..b24b8846b30fe2 100644 --- a/be/test/exprs/function/function_test_util.cpp +++ b/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; } diff --git a/be/test/exprs/function/table_function_test.cpp b/be/test/exprs/function/table_function_test.cpp index 326a2d98f98c89..06ae08af71354d 100644 --- a/be/test/exprs/function/table_function_test.cpp +++ b/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; diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinTableGeneratingFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinTableGeneratingFunctions.java index f764da5718c389..165f3c5f966ecc 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinTableGeneratingFunctions.java +++ b/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 getReturnManyColumnFunctions() { return RETURN_MULTI_COLUMNS_FUNCTIONS; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java index a650d17d93408a..74019610b15a8e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindExpression.java +++ b/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> 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 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())), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java index 5338bf7df9b2ef..27d5a509609e5b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/FoldConstantRuleOnFE.java @@ -121,6 +121,8 @@ public class FoldConstantRuleOnFE extends AbstractExpressionRewriteRule public static final FoldConstantRuleOnFE VISITOR_INSTANCE = new FoldConstantRuleOnFE(true); public static final FoldConstantRuleOnFE PATTERN_MATCH_INSTANCE = new FoldConstantRuleOnFE(false); + private static final FoldConstantRuleOnFE CONTEXT_FREE_VISITOR_INSTANCE + = new FoldConstantRuleOnFE(true, false); // record whether current expression is in an aggregate function with distinct, // if is, we will skip to fold constant @@ -128,15 +130,26 @@ public class FoldConstantRuleOnFE extends AbstractExpressionRewriteRule private static final CheckWhetherUnderAggDistinct NOT_UNDER_AGG_DISTINCT = new CheckWhetherUnderAggDistinct(); private final boolean deepRewrite; + private final boolean foldContextDependentExpressions; public FoldConstantRuleOnFE(boolean deepRewrite) { + this(deepRewrite, true); + } + + private FoldConstantRuleOnFE(boolean deepRewrite, boolean foldContextDependentExpressions) { this.deepRewrite = deepRewrite; + this.foldContextDependentExpressions = foldContextDependentExpressions; } public static Expression evaluate(Expression expression, ExpressionRewriteContext expressionRewriteContext) { return VISITOR_INSTANCE.rewrite(expression, expressionRewriteContext); } + /** Evaluate expressions that do not require a rewrite or connection context. */ + public static Expression evaluateWithoutContext(Expression expression) { + return CONTEXT_FREE_VISITOR_INSTANCE.rewrite(expression, null); + } + @Override public List> buildListeners() { return ImmutableList.of( @@ -229,12 +242,18 @@ public Expression visitMatch(Match match, ExpressionRewriteContext context) { @Override public Expression visitUnboundVariable(UnboundVariable unboundVariable, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return unboundVariable; + } Variable variable = ExpressionAnalyzer.resolveUnboundVariable(unboundVariable); return variable.getRealExpression(); } @Override public Expression visitEncryptKeyRef(EncryptKeyRef encryptKeyRef, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return encryptKeyRef; + } String dbName = encryptKeyRef.getDbName(); ConnectContext connectContext = context.cascadesContext.getConnectContext(); if (Strings.isNullOrEmpty(dbName)) { @@ -357,36 +376,54 @@ public Expression visitNot(Not not, ExpressionRewriteContext context) { @Override public Expression visitDatabase(Database database, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return database; + } String res = context.cascadesContext.getConnectContext().getDatabase(); return new VarcharLiteral(res); } @Override public Expression visitCurrentUser(CurrentUser currentUser, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return currentUser; + } String res = context.cascadesContext.getConnectContext().getCurrentUserIdentity().toString(); return new VarcharLiteral(res); } @Override public Expression visitCurrentCatalog(CurrentCatalog currentCatalog, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return currentCatalog; + } String res = context.cascadesContext.getConnectContext().getDefaultCatalog(); return new VarcharLiteral(res); } @Override public Expression visitUser(User user, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return user; + } String res = context.cascadesContext.getConnectContext().getUserWithLoginRemoteIpString(); return new VarcharLiteral(res); } @Override public Expression visitSessionUser(SessionUser user, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return user; + } String res = context.cascadesContext.getConnectContext().getUserWithLoginRemoteIpString(); return new VarcharLiteral(res); } @Override public Expression visitLastQueryId(LastQueryId queryId, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return queryId; + } String res = "Not Available"; TUniqueId id = context.cascadesContext.getConnectContext().getLastQueryId(); if (id != null) { @@ -397,6 +434,9 @@ public Expression visitLastQueryId(LastQueryId queryId, ExpressionRewriteContext @Override public Expression visitConnectionId(ConnectionId connectionId, ExpressionRewriteContext context) { + if (!foldContextDependentExpressions) { + return connectionId; + } return new BigIntLiteral(context.cascadesContext.getConnectContext().getConnectionId()); } 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 1f3f9ffaf8d83e..c1201807f487af 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 @@ -21,8 +21,10 @@ import org.apache.doris.nereids.trees.expressions.ExecFunction; 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.FloatLiteral; import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.literal.MapLiteral; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.types.ArrayType; import org.apache.doris.nereids.types.FloatType; @@ -36,6 +38,18 @@ */ public class ArrayArithmetic { + /** Return the number of elements in an array. */ + @ExecFunction(name = "cardinality") + public static Expression cardinality(ArrayLiteral array) { + return new BigIntLiteral(array.getValue().size()); + } + + /** Return the number of entries in a map. */ + @ExecFunction(name = "cardinality") + public static Expression cardinality(MapLiteral map) { + return new BigIntLiteral(map.getValue().size()); + } + /** * 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/generator/Stack.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/generator/Stack.java new file mode 100644 index 00000000000000..f47005c52b5914 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/generator/Stack.java @@ -0,0 +1,155 @@ +// 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.generator; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.rules.expression.rules.FoldConstantRuleOnFE; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable; +import org.apache.doris.nereids.trees.expressions.functions.ComputePrecision; +import org.apache.doris.nereids.trees.expressions.functions.CustomSignature; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLikeLiteral; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.DataType; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.NullType; +import org.apache.doris.nereids.types.StructField; +import org.apache.doris.nereids.types.StructType; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.ArrayList; +import java.util.List; + +/** + * stack(n, expr1, ..., exprk) separates the expressions into n rows in row-major order. + * Missing values in the last row are padded with nulls. + */ +public class Stack extends TableGeneratingFunction implements CustomSignature, ComputePrecision, AlwaysNullable { + + /** constructor with two or more arguments. */ + public Stack(Expression numRows, Expression argument, Expression... otherArguments) { + super("stack", ExpressionUtils.mergeArguments(numRows, argument, otherArguments)); + } + + /** constructor for withChildren and reuse signature. */ + private Stack(GeneratorFunctionParams functionParams) { + super(functionParams); + } + + @Override + public Stack withChildren(List children) { + Preconditions.checkArgument(children.size() >= 2); + return new Stack(getFunctionParams(children)); + } + + @Override + public void checkLegalityBeforeTypeCoercion() { + getColumnTypes(); + } + + @Override + public FunctionSignature computePrecision(FunctionSignature signature) { + return signature; + } + + @Override + public FunctionSignature searchSignature(List signatures) { + return super.searchSignature(signatures); + } + + @Override + public FunctionSignature customSignature() { + List columnTypes = getColumnTypes(); + List argumentTypes = new ArrayList<>(arity()); + argumentTypes.add(IntegerType.INSTANCE); + for (int i = 1; i < arity(); i++) { + argumentTypes.add(columnTypes.get((i - 1) % columnTypes.size())); + } + + if (columnTypes.size() == 1) { + return FunctionSignature.of(columnTypes.get(0), argumentTypes); + } + ImmutableList.Builder fields = ImmutableList.builder(); + for (int i = 0; i < columnTypes.size(); i++) { + fields.add(new StructField("col" + i, columnTypes.get(i), true, "")); + } + return FunctionSignature.of(new StructType(fields.build()), argumentTypes); + } + + private int getNumRows() { + Expression numRowsArgument = getArgument(0); + if (!numRowsArgument.isConstant()) { + throw new AnalysisException("The first argument of stack must be a positive constant integer, but got: " + + numRowsArgument.toSql()); + } + Expression evaluated = FoldConstantRuleOnFE.evaluateWithoutContext(numRowsArgument); + if (!(evaluated instanceof IntegerLikeLiteral)) { + throw new AnalysisException("The first argument of stack must be a positive constant integer, but got: " + + numRowsArgument.toSql()); + } + long numRows = ((IntegerLikeLiteral) evaluated).getLongValue(); + if (numRows <= 0 || numRows > Integer.MAX_VALUE) { + throw new AnalysisException("The first argument of stack must be in (0, " + Integer.MAX_VALUE + + "], but got: " + numRows); + } + return (int) numRows; + } + + /** Return the number of logical output columns derived from the row count and value arguments. */ + public int getOutputColumnCount() { + int numRows = getNumRows(); + return (arity() - 2) / numRows + 1; + } + + private List getColumnTypes() { + int numFields = getOutputColumnCount(); + List columnTypes = new ArrayList<>(numFields); + for (int columnIndex = 0; columnIndex < numFields; columnIndex++) { + DataType referenceType = NullType.INSTANCE; + int referenceArgumentIndex = -1; + for (int argumentIndex = columnIndex + 1; argumentIndex < arity(); argumentIndex += numFields) { + DataType fieldType = getArgument(argumentIndex).getDataType(); + if (fieldType.isNullType()) { + continue; + } + if (referenceType.isNullType()) { + referenceType = fieldType; + referenceArgumentIndex = argumentIndex; + continue; + } + if (!referenceType.equals(fieldType)) { + throw new AnalysisException("The expressions for stack output column " + columnIndex + + " must have compatible types, but argument " + referenceArgumentIndex + " is " + + referenceType.toSql() + " while argument " + argumentIndex + " is " + + fieldType.toSql()); + } + } + columnTypes.add(referenceType); + } + return columnTypes; + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitStack(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/TableGeneratingFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/TableGeneratingFunctionVisitor.java index 384b3209ee5c52..9bcdae6351fd67 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/TableGeneratingFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/TableGeneratingFunctionVisitor.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.TableGeneratingFunction; import org.apache.doris.nereids.trees.expressions.functions.generator.Unnest; import org.apache.doris.nereids.trees.expressions.functions.udf.JavaUdtf; @@ -171,6 +172,10 @@ default R visitPosExplodeOuter(PosExplodeOuter posExplodeOuter, C context) { return visitTableGeneratingFunction(posExplodeOuter, context); } + default R visitStack(Stack stack, C context) { + return visitTableGeneratingFunction(stack, context); + } + default R visitUnnest(Unnest unnest, C context) { return visitTableGeneratingFunction(unnest, context); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/generator/StackTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/generator/StackTest.java new file mode 100644 index 00000000000000..d5330665e3dfac --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/generator/StackTest.java @@ -0,0 +1,170 @@ +// 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.generator; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.Subtract; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Array; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Cardinality; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ConnectionId; +import org.apache.doris.nereids.trees.expressions.functions.scalar.CurrentCatalog; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.NullType; +import org.apache.doris.nereids.types.StringType; +import org.apache.doris.nereids.types.StructType; +import org.apache.doris.nereids.util.MemoTestUtils; +import org.apache.doris.nereids.util.PlanChecker; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class StackTest { + + @Test + public void testMultiColumnSignature() { + Stack stack = new Stack(new IntegerLiteral(2), new IntegerLiteral(1), + new StringLiteral("a"), new IntegerLiteral(2), new StringLiteral("b")); + + FunctionSignature signature = stack.getSignatures().get(0); + Assertions.assertEquals(5, signature.argumentsTypes.size()); + Assertions.assertEquals(IntegerType.INSTANCE, signature.argumentsTypes.get(0)); + Assertions.assertEquals(IntegerType.INSTANCE, signature.argumentsTypes.get(1)); + Assertions.assertEquals(StringType.INSTANCE, signature.argumentsTypes.get(2)); + Assertions.assertTrue(signature.returnType.isStructType()); + StructType returnType = (StructType) signature.returnType; + Assertions.assertEquals(2, returnType.getFields().size()); + Assertions.assertEquals("col0", returnType.getFields().get(0).getName()); + Assertions.assertEquals(IntegerType.INSTANCE, returnType.getFields().get(0).getDataType()); + Assertions.assertEquals("col1", returnType.getFields().get(1).getName()); + Assertions.assertEquals(StringType.INSTANCE, returnType.getFields().get(1).getDataType()); + } + + @Test + public void testNullUsesOutputColumnType() { + Stack stack = new Stack(new IntegerLiteral(2), new IntegerLiteral(1), + new StringLiteral("a"), new NullLiteral(), new StringLiteral("b")); + + FunctionSignature signature = stack.getSignatures().get(0); + Assertions.assertEquals(IntegerType.INSTANCE, signature.argumentsTypes.get(3)); + StructType returnType = (StructType) signature.returnType; + Assertions.assertEquals(IntegerType.INSTANCE, returnType.getFields().get(0).getDataType()); + Assertions.assertEquals(StringType.INSTANCE, returnType.getFields().get(1).getDataType()); + } + + @Test + public void testSingleColumnSignature() { + Stack stack = new Stack(new IntegerLiteral(4), new IntegerLiteral(1), + new IntegerLiteral(2), new IntegerLiteral(3)); + + FunctionSignature signature = stack.getSignatures().get(0); + Assertions.assertEquals(IntegerType.INSTANCE, signature.returnType); + Assertions.assertEquals(4, signature.argumentsTypes.size()); + } + + @Test + public void testFoldableNumRowsExpression() { + Stack stack = new Stack( + new Cast(new Subtract(new IntegerLiteral(3), new IntegerLiteral(1)), IntegerType.INSTANCE), + new IntegerLiteral(1), new IntegerLiteral(2), new IntegerLiteral(3)); + + FunctionSignature signature = stack.getSignatures().get(0); + Assertions.assertTrue(signature.returnType.isStructType()); + Assertions.assertEquals(2, ((StructType) signature.returnType).getFields().size()); + } + + @Test + public void testCardinalityNumRowsExpression() { + Stack stack = new Stack(new Cardinality(new Array(new IntegerLiteral(1), new IntegerLiteral(2))), + new IntegerLiteral(1), new IntegerLiteral(2), new IntegerLiteral(3)); + + FunctionSignature signature = stack.getSignatures().get(0); + Assertions.assertTrue(signature.returnType.isStructType()); + Assertions.assertEquals(2, ((StructType) signature.returnType).getFields().size()); + } + + @Test + public void testAllNullOutputColumn() { + Stack stack = new Stack(new IntegerLiteral(2), new NullLiteral(), new NullLiteral()); + + FunctionSignature signature = stack.getSignatures().get(0); + Assertions.assertEquals(NullType.INSTANCE, signature.returnType); + } + + @Test + public void testInvalidArguments() { + Assertions.assertThrows(AnalysisException.class, + () -> new Stack(new IntegerLiteral(0), new IntegerLiteral(1)).getSignatures()); + Assertions.assertThrows(AnalysisException.class, + () -> new Stack(SlotReference.of("n", IntegerType.INSTANCE), + new IntegerLiteral(1)).getSignatures()); + Assertions.assertThrows(AnalysisException.class, + () -> new Stack(new ConnectionId(), new IntegerLiteral(1)).getSignatures()); + Assertions.assertThrows(AnalysisException.class, + () -> new Stack(new CurrentCatalog(), new IntegerLiteral(1)).getSignatures()); + Assertions.assertThrows(AnalysisException.class, + () -> new Stack(new IntegerLiteral(2), new IntegerLiteral(1), + new StringLiteral("a")).getSignatures()); + } + + @Test + public void testContextDependentNumRowsExpression() { + AnalysisException keyException = Assertions.assertThrows(AnalysisException.class, + () -> PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select c1 from (select 1) t " + + "lateral view stack(KEY test_db.test_key, 1) s as c1")); + Assertions.assertTrue(keyException.getMessage().contains( + "The first argument of stack must be a positive constant integer"), keyException.getMessage()); + + AnalysisException castKeyException = Assertions.assertThrows(AnalysisException.class, + () -> PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select c1 from (select 1) t " + + "lateral view stack(CAST(KEY test_db.test_key AS INT), 1) s as c1")); + Assertions.assertTrue(castKeyException.getMessage().contains( + "The first argument of stack must be a positive constant integer"), castKeyException.getMessage()); + } + + @Test + public void testMultiColumnAliasCount() { + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select c1 from (select 1) t " + + "lateral view stack(2, 1, 2, 3, 4, 5) s as c1")); + Assertions.assertTrue(exception.getMessage().contains( + "table s has 3 columns available but 1 columns specified")); + + PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select c1, c2, c3 from (select 1) t " + + "lateral view stack(2, 1, 2, 3, 4, 5) s as c1, c2, c3"); + } + + @Test + public void testSingleStructOutputColumn() { + PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select c.a from (select 1) t " + + "lateral view stack(2, named_struct('a', 1), named_struct('a', 2)) s as c"); + PlanChecker.from(MemoTestUtils.createConnectContext()).analyze( + "select c.a, c.b from (select 1) t lateral view stack(2, " + + "named_struct('a', 1, 'b', 2), named_struct('a', 3, 'b', 4)) s as c"); + } +} diff --git a/regression-test/data/query_p0/sql_functions/table_function/stack.out b/regression-test/data/query_p0/sql_functions/table_function/stack.out new file mode 100644 index 00000000000000..fbb0d2844386f5 --- /dev/null +++ b/regression-test/data/query_p0/sql_functions/table_function/stack.out @@ -0,0 +1,52 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !two_rows -- +1 2 +3 \N + +-- !three_rows -- +1 a +2 b +3 c + +-- !single_column -- +\N +1 +2 +3 + +-- !null_type_coercion -- +\N b +1 a + +-- !all_null_column -- +\N +\N + +-- !constant_expression -- +1 2 +3 \N + +-- !multi_row_constants -- +1 1 a +1 2 b +2 1 a +2 2 b + +-- !cardinality_num_rows -- +1 2 +3 \N + +-- !single_field_struct_value -- +{"a":1} +{"a":2} + +-- !multi_field_struct_value -- +{"a":1, "b":2} +{"a":3, "b":4} + +-- !column_arguments -- +1 10 x +1 20 y +2 \N \N +2 30 m + diff --git a/regression-test/suites/query_p0/sql_functions/table_function/stack.groovy b/regression-test/suites/query_p0/sql_functions/table_function/stack.groovy new file mode 100644 index 00000000000000..97b54c79ef4ab7 --- /dev/null +++ b/regression-test/suites/query_p0/sql_functions/table_function/stack.groovy @@ -0,0 +1,147 @@ +// 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("stack") { + order_qt_two_rows """ + select c1, c2 + from (select 1) t lateral view stack(2, 1, 2, 3) s as c1, c2 + order by c1, c2 + """ + + order_qt_three_rows """ + select c1, c2 + from (select 1) t lateral view stack(3, 1, 'a', 2, 'b', 3, 'c') s as c1, c2 + order by c1, c2 + """ + + order_qt_single_column """ + select c1 + from (select 1) t lateral view stack(4, 1, 2, 3) s as c1 + order by c1 + """ + + order_qt_null_type_coercion """ + select c1, c2 + from (select 1) t lateral view stack(2, 1, 'a', null, 'b') s as c1, c2 + order by c1, c2 + """ + + order_qt_all_null_column """ + select c1 + from (select 1) t lateral view stack(2, null, null) s as c1 + order by c1 + """ + + order_qt_constant_expression """ + select c1, c2 + from (select 1) t lateral view stack(3 - 1, 1, 2, 3) s as c1, c2 + order by c1, c2 + """ + + order_qt_multi_row_constants """ + select id, c1, c2 + from (select 1 as id union all select 2 as id) t + lateral view stack(2, 1, 'a', 2, 'b') s as c1, c2 + order by id, c1, c2 + """ + + order_qt_cardinality_num_rows """ + select c1, c2 + from (select 1) t lateral view stack(cardinality([1, 2]), 1, 2, 3) s as c1, c2 + order by c1, c2 + """ + + order_qt_single_field_struct_value """ + select c + from (select 1) t lateral view stack(2, + named_struct('a', 1), named_struct('a', 2)) s as c + order by c.a + """ + + order_qt_multi_field_struct_value """ + select c + from (select 1) t lateral view stack(2, + named_struct('a', 1, 'b', 2), named_struct('a', 3, 'b', 4)) s as c + order by c.a, c.b + """ + + sql "drop table if exists test_stack" + sql """ + create table test_stack ( + id int, + a int, + b int, + s1 string, + s2 string + ) + duplicate key(id) + distributed by hash(id) buckets 1 + properties("replication_num" = "1") + """ + sql "insert into test_stack values (1, 10, 20, 'x', 'y'), (2, 30, null, 'm', null)" + + order_qt_column_arguments """ + select id, c1, c2 + from test_stack lateral view stack(2, a, s1, b, s2) s as c1, c2 + order by id, c1, c2 + """ + + test { + sql "select c1 from (select 1) t lateral view stack(0, 1) s as c1" + exception "The first argument of stack must be in" + } + + test { + sql "select c1 from (select 1) t lateral view stack(1.5, 1) s as c1" + exception "The first argument of stack must be a positive constant integer" + } + + test { + sql "select c1 from test_stack lateral view stack(id, a) s as c1" + exception "The first argument of stack must be a positive constant integer" + } + + test { + sql "select c1 from (select 1) t lateral view stack(2, 1, 'a') s as c1" + exception "must have compatible types" + } + + test { + sql "select c1 from (select 1) t lateral view stack(connection_id(), 1) s as c1" + exception "The first argument of stack must be a positive constant integer" + } + + test { + sql "select c1 from (select 1) t lateral view stack(current_catalog(), 1) s as c1" + exception "The first argument of stack must be a positive constant integer" + } + + test { + sql "select c1 from (select 1) t lateral view stack(2, 1, 2, 3, 4, 5) s as c1" + exception "has 3 columns available but 1 columns specified" + } + + test { + sql "select c1 from (select 1) t lateral view stack(2, 1, 2, 3, 4, 5) s as c1, c2" + exception "has 3 columns available but 2 columns specified" + } + + test { + sql "select c1 from (select 1) t lateral view stack(2, 1, 2, 3, 4, 5) s as c1, c2, c3, c4" + exception "has 3 columns available but 4 columns specified" + } +}