From 91872705091582da650ad64e625845664f390659 Mon Sep 17 00:00:00 2001
From: eldenmoon <15605149486@163.com>
Date: Wed, 21 Aug 2024 20:30:17 +0800
Subject: [PATCH] [Fix](Variant) fix element_at should return nullable if
 result type is nullable

---
 .../functions/function_variant_element.cpp    | 24 ++++++++++++--
 .../data/variant_p0/element_function.out      |  4 +++
 .../suites/variant_p0/element_function.groovy | 32 +++++++++++++++++++
 3 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 regression-test/data/variant_p0/element_function.out
 create mode 100644 regression-test/suites/variant_p0/element_function.groovy

diff --git a/be/src/vec/functions/function_variant_element.cpp b/be/src/vec/functions/function_variant_element.cpp
index 76076a498d07c0..53340cd61ea8d8 100644
--- a/be/src/vec/functions/function_variant_element.cpp
+++ b/be/src/vec/functions/function_variant_element.cpp
@@ -57,7 +57,7 @@ class FunctionVariantElement : public IFunction {
     // Get function name.
     String get_name() const override { return name; }
 
-    bool use_default_implementation_for_nulls() const override { return true; }
+    bool use_default_implementation_for_nulls() const override { return false; }
 
     size_t get_number_of_arguments() const override { return 2; }
 
@@ -77,10 +77,27 @@ class FunctionVariantElement : public IFunction {
         return make_nullable(std::make_shared<DataTypeObject>());
     }
 
+    // wrap variant column with nullable
+    // 1. if variant is null root(empty or nothing as root), then nullable map is all null
+    // 2. if variant is scalar variant, then use the root's nullable map
+    // 3. if variant is hierarchical variant, then create a nullable map with all none null
+    ColumnPtr wrap_variant_nullable(ColumnPtr col) const {
+        const auto& var = assert_cast<const ColumnObject&>(*col);
+        if (var.is_null_root()) {
+            return make_nullable(col, true);
+        }
+        if (var.is_scalar_variant() && var.get_root()->is_nullable()) {
+            const auto* nullable = assert_cast<const ColumnNullable*>(var.get_root().get());
+            return ColumnNullable::create(
+                    col, nullable->get_null_map_column_ptr()->clone_resized(col->size()));
+        }
+        return make_nullable(col);
+    }
+
     Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                         size_t result, size_t input_rows_count) const override {
         const auto* variant_col = check_and_get_column<ColumnObject>(
-                block.get_by_position(arguments[0]).column.get());
+                remove_nullable(block.get_by_position(arguments[0]).column).get());
         if (!variant_col) {
             return Status::RuntimeError(
                     fmt::format("unsupported types for function {}({}, {})", get_name(),
@@ -95,6 +112,9 @@ class FunctionVariantElement : public IFunction {
         auto index_column = block.get_by_position(arguments[1]).column;
         ColumnPtr result_column;
         RETURN_IF_ERROR(get_element_column(*variant_col, index_column, &result_column));
+        if (block.get_by_position(result).type->is_nullable()) {
+            result_column = wrap_variant_nullable(result_column);
+        }
         block.replace_by_position(result, result_column);
         return Status::OK();
     }
diff --git a/regression-test/data/variant_p0/element_function.out b/regression-test/data/variant_p0/element_function.out
new file mode 100644
index 00000000000000..095c7b2035676a
--- /dev/null
+++ b/regression-test/data/variant_p0/element_function.out
@@ -0,0 +1,4 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !sql --
+1
+
diff --git a/regression-test/suites/variant_p0/element_function.groovy b/regression-test/suites/variant_p0/element_function.groovy
new file mode 100644
index 00000000000000..7b5e55ea53bdad
--- /dev/null
+++ b/regression-test/suites/variant_p0/element_function.groovy
@@ -0,0 +1,32 @@
+// 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("regression_test_variant_element_at", "p0")  {
+      sql """
+        CREATE TABLE IF NOT EXISTS element_fn_test(
+            k bigint,
+            v variant,
+            v1 variant not null,
+        )
+        UNIQUE KEY(`k`)
+        DISTRIBUTED BY HASH(k) BUCKETS 4
+        properties("replication_num" = "1");
+    """
+
+    sql """insert into element_fn_test values (1, '{"arr1" : [1, 2, 3]}', '{"arr2" : [4, 5, 6]}')"""
+    qt_sql """select array_first((x,y) -> (x - y) < 0, cast(v['arr1'] as array<int>), cast(v1['arr2'] as array<int>)) from element_fn_test"""
+}
\ No newline at end of file