From bac29c0bbafb7a97e033ca08bdc12a9d13914c15 Mon Sep 17 00:00:00 2001
From: Anton Popov <anton@clickhouse.com>
Date: Tue, 6 Feb 2024 15:52:19 +0000
Subject: [PATCH] add test for variant subcolumn

---
 .../RewriteFunctionToSubcolumnVisitor.cpp     |  6 ++++++
 ..._functions_to_subcolumns_variant.reference |  8 +++++++
 .../02971_functions_to_subcolumns_variant.sql | 21 +++++++++++++++++++
 3 files changed, 35 insertions(+)
 create mode 100644 tests/queries/0_stateless/02971_functions_to_subcolumns_variant.reference
 create mode 100644 tests/queries/0_stateless/02971_functions_to_subcolumns_variant.sql

diff --git a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp
index 04451947796..437d46c24b2 100644
--- a/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp
+++ b/src/Interpreters/RewriteFunctionToSubcolumnVisitor.cpp
@@ -124,6 +124,12 @@ void RewriteFunctionToSubcolumnFirstPassMatcher::visit(const ASTFunction & funct
             if (value_type == Field::Types::UInt64 || value_type == Field::Types::String)
                 ++data.optimized_identifiers_count[column->name];
         }
+        else if (function.name == "variantElement" && column_type_id == TypeIndex::Variant)
+        {
+            const auto * literal = arguments[1]->as<ASTLiteral>();
+            if (literal && literal->value.getType() == Field::Types::String)
+                ++data.optimized_identifiers_count[column->name];
+        }
         else if (function.name == "mapContains" && column_type_id == TypeIndex::Map)
         {
             ++data.optimized_identifiers_count[column->name];
diff --git a/tests/queries/0_stateless/02971_functions_to_subcolumns_variant.reference b/tests/queries/0_stateless/02971_functions_to_subcolumns_variant.reference
new file mode 100644
index 00000000000..7a52155fc2d
--- /dev/null
+++ b/tests/queries/0_stateless/02971_functions_to_subcolumns_variant.reference
@@ -0,0 +1,8 @@
+SELECT `v.String` AS `variantElement(v, \'String\')`
+FROM t_func_to_subcolumns_variant
+foo
+\N
+SELECT __table1.`v.String` AS `variantElement(v, \'String\')`
+FROM default.t_func_to_subcolumns_variant AS __table1
+foo
+\N
diff --git a/tests/queries/0_stateless/02971_functions_to_subcolumns_variant.sql b/tests/queries/0_stateless/02971_functions_to_subcolumns_variant.sql
new file mode 100644
index 00000000000..1cedd877289
--- /dev/null
+++ b/tests/queries/0_stateless/02971_functions_to_subcolumns_variant.sql
@@ -0,0 +1,21 @@
+DROP TABLE IF EXISTS t_func_to_subcolumns_variant;
+
+SET allow_experimental_variant_type = 1;
+
+CREATE TABLE t_func_to_subcolumns_variant (id UInt64, v Variant(String, UInt64)) ENGINE = MergeTree ORDER BY id;
+
+INSERT INTO t_func_to_subcolumns_variant VALUES (1, 'foo') (2, 111);
+
+SET optimize_functions_to_subcolumns = 1;
+SET allow_experimental_analyzer = 0;
+
+EXPLAIN SYNTAX SELECT variantElement(v, 'String') FROM t_func_to_subcolumns_variant;
+SELECT variantElement(v, 'String') FROM t_func_to_subcolumns_variant;
+
+SET optimize_functions_to_subcolumns = 1;
+SET allow_experimental_analyzer = 1;
+
+EXPLAIN QUERY TREE dump_tree = 0, dump_ast = 1 SELECT variantElement(v, 'String') FROM t_func_to_subcolumns_variant;
+SELECT variantElement(v, 'String') FROM t_func_to_subcolumns_variant;
+
+DROP TABLE t_func_to_subcolumns_variant;