From 51cc729fd17fb8d5687fac34833914dc540ee5a9 Mon Sep 17 00:00:00 2001
From: Soumyava <93540295+somu-imply@users.noreply.github.com>
Date: Mon, 26 Feb 2024 21:53:49 -0800
Subject: [PATCH] Enforcing type checking for flatten concat (#15903)

---
 .../sql/calcite/rule/FlattenConcatRule.java   |  5 +++-
 .../druid/sql/calcite/CalciteQueryTest.java   | 30 +++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/FlattenConcatRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/FlattenConcatRule.java
index 589c590b107..48aee4cf355 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/FlattenConcatRule.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/FlattenConcatRule.java
@@ -32,6 +32,7 @@ import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.tools.RelBuilder;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.math.expr.Function;
 import org.apache.druid.sql.calcite.expression.builtin.ConcatOperatorConversion;
@@ -59,7 +60,9 @@ public class FlattenConcatRule extends RelOptRule implements SubstitutionRule
 
     //noinspection ObjectEquality
     if (newNode != oldNode) {
-      call.transformTo(newNode);
+      final RelBuilder relBuilder = call.builder().push(newNode);
+      relBuilder.convert(oldNode.getRowType(), false);
+      call.transformTo(relBuilder.build());
       call.getPlanner().prune(oldNode);
     }
   }
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
index 68174e7757a..4d9614132d8 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
@@ -15312,6 +15312,36 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
         .run();
   }
 
+  @Test
+  public void testCastCharToVarcharInFlattenConcat()
+  {
+    cannotVectorize();
+    testQuery(
+        "select 'A'||cast(col as char)||'B' from (values(1)) as t(col)",
+        ImmutableList.of(
+            Druids.newScanQueryBuilder()
+                  .dataSource(
+                      InlineDataSource.fromIterable(
+                          ImmutableList.of(
+                              new Object[]{"A1B"}
+                          ),
+                          RowSignature.builder().add("EXPR$0", ColumnType.STRING).build()
+                      )
+                  )
+                  .intervals(querySegmentSpec(Filtration.eternity()))
+                  .resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
+                  .legacy(false)
+                  .columns(ImmutableList.of(
+                      "EXPR$0"
+                  ))
+                  .build()
+        ),
+        ImmutableList.of(
+            new Object[]{"A1B"}
+        )
+    );
+  }
+
   @Test
   public void testFilterParseLongNullable()
   {