array_zip and repeat

zhangstar333 · Dec 12, 2023 · b380676 · b380676
1 parent 96e739c
commit b380676
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 9 deletions.
diff --git a/be/src/vec/functions/array/function_array_zip.cpp b/be/src/vec/functions/array/function_array_zip.cpp
@@ -94,12 +94,16 @@ class FunctionArrayZip : public IFunction {
                         size_t result, size_t input_rows_count) const override {
         size_t num_element = arguments.size();
 
-        // all the columns must have the same size as the first column
+        // all the columns must have the same size as the first column, except have NULL literal
         ColumnPtr first_array_column;
         Columns tuple_columns(num_element);
+        bool have_null_literal = false;
+        int nested_column_data_num = 0;
 
         for (size_t i = 0; i < num_element; ++i) {
             auto col = block.get_by_position(arguments[i]).column;
+            auto type = block.get_by_position(arguments[i]).type;
+            const auto& nested_type = assert_cast<const DataTypeArray&>(*type).get_nested_type();
             col = col->convert_to_full_column_if_const();
 
             const auto* column_array = check_and_get_column<ColumnArray>(col.get());
@@ -108,23 +112,39 @@ class FunctionArrayZip : public IFunction {
                         "execute failed, function {}'s {}-th argument should be array bet get {}",
                         get_name(), i + 1, block.get_by_position(arguments[i]).type->get_name()));
             }
+            bool is_null_literal = nested_type->is_null_literal();
+            have_null_literal |= is_null_literal;
 
             if (i == 0) {
                 first_array_column = col;
-            } else if (!column_array->has_equal_offsets(
+            } else if (!have_null_literal &&
+                       !column_array->has_equal_offsets(
                                static_cast<const ColumnArray&>(*first_array_column))) {
                 return Status::RuntimeError(
                         fmt::format("execute failed, function {}'s {}-th argument should have same "
                                     "offsets with first argument",
                                     get_name(), i + 1));
             }
-
-            tuple_columns[i] = column_array->get_data_ptr();
+            if (is_null_literal) { //wants handle: select /*set_var(enable_fold_constant_by_be = true) */ array_zip([1, 2, 3], null, ['foo', 'bar', 'test']);
+                auto nullable_column =
+                        ColumnNullable::create(ColumnUInt8::create(nested_column_data_num, 1),
+                                               ColumnUInt8::create(nested_column_data_num, 1));
+                tuple_columns[i] = std::move(nullable_column);
+            } else {
+                tuple_columns[i] = column_array->get_data_ptr();
+                nested_column_data_num = column_array->get_data_ptr()->size();
+            }
         }
 
         auto tuples = ColumnStruct::create(tuple_columns);
-        auto nullable_tuples =
-                ColumnNullable::create(tuples, ColumnUInt8::create(tuples->size(), 0));
+
+        ColumnPtr null_map = nullptr;
+        if (have_null_literal) {
+            null_map = ColumnUInt8::create(tuples->size(), 1);
+        } else {
+            null_map = ColumnUInt8::create(tuples->size(), 0);
+        }
+        auto nullable_tuples = ColumnNullable::create(tuples, null_map);
         auto res_column = ColumnArray::create(
                 nullable_tuples,
                 static_cast<const ColumnArray&>(*first_array_column).get_offsets_ptr());

diff --git a/...ore/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Repeat.java b/...ore/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Repeat.java
@@ -25,7 +25,6 @@
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
 import org.apache.doris.nereids.types.IntegerType;
 import org.apache.doris.nereids.types.StringType;
-import org.apache.doris.nereids.types.VarcharType;
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
@@ -39,7 +38,6 @@ public class Repeat extends ScalarFunction
         implements BinaryExpression, ExplicitlyCastableSignature, AlwaysNullable {
 
     public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
-            FunctionSignature.ret(VarcharType.SYSTEM_DEFAULT).args(VarcharType.SYSTEM_DEFAULT, IntegerType.INSTANCE),
             FunctionSignature.ret(StringType.INSTANCE).args(StringType.INSTANCE, IntegerType.INSTANCE)
     );
 

diff --git a/gensrc/script/doris_builtins_functions.py b/gensrc/script/doris_builtins_functions.py
@@ -1558,7 +1558,6 @@
         [['null_or_empty'], 'BOOLEAN', ['VARCHAR'], 'ALWAYS_NOT_NULLABLE'],
         [['not_null_or_empty'], 'BOOLEAN', ['VARCHAR'], 'ALWAYS_NOT_NULLABLE'],
         [['space'], 'VARCHAR', ['INT'], ''],
-        [['repeat'], 'VARCHAR', ['VARCHAR', 'INT'], 'ALWAYS_NULLABLE'],
         [['lpad'], 'VARCHAR', ['VARCHAR', 'INT', 'VARCHAR'], 'ALWAYS_NULLABLE'],
         [['rpad'], 'VARCHAR', ['VARCHAR', 'INT', 'VARCHAR'], 'ALWAYS_NULLABLE'],
         [['append_trailing_char_if_absent'], 'VARCHAR', ['VARCHAR', 'VARCHAR'], 'ALWAYS_NULLABLE'],