From 9620c4ed4500e6a8fb9ec5abbaf58ac9bbac09f6 Mon Sep 17 00:00:00 2001
From: Laksh Singla <lakshsingla@gmail.com>
Date: Tue, 3 Oct 2023 10:12:27 +0530
Subject: [PATCH] changes to the new rules

---
 .../rel/DruidFreeUnionDataSourceRel.java      | 22 ++++++++++---------
 .../rule/DruidFreeUnionDataSourceRule.java    |  9 +++++---
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidFreeUnionDataSourceRel.java b/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidFreeUnionDataSourceRel.java
index a59749c5548b..48784469327a 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidFreeUnionDataSourceRel.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidFreeUnionDataSourceRel.java
@@ -137,13 +137,15 @@ public DruidQuery toDruidQuery(final boolean finalizeAggregations)
         signature = query.getOutputRowSignature();
       }
 
-      if (signature.getColumnNames().equals(query.getOutputRowSignature().getColumnNames())) {
-        dataSources.add(dataSource);
-      } else {
-        getPlannerContext().setPlanningError(
-            "There is a mismatch between the output row signature of input tables and the row signature of union output.");
-        throw new CannotBuildQueryException(druidRel);
-      }
+      dataSources.add(dataSource);
+
+//      if (signature.getColumnNames().equals(query.getOutputRowSignature().getColumnNames())) {
+//        dataSources.add(dataSource);
+//      } else {
+//        getPlannerContext().setPlanningError(
+//            "There is a mismatch between the output row signature of input tables and the row signature of union output.");
+//        throw new CannotBuildQueryException(druidRel);
+//      }
     }
 
     if (signature == null) {
@@ -153,9 +155,9 @@ public DruidQuery toDruidQuery(final boolean finalizeAggregations)
 
     // Sanity check: the columns we think we're building off must equal the "unionColumnNames" registered at
     // creation time.
-    if (!signature.getColumnNames().equals(unionColumnNames)) {
-      throw new CannotBuildQueryException(unionRel);
-    }
+//    if (!signature.getColumnNames().equals(unionColumnNames)) {
+//      throw new CannotBuildQueryException(unionRel);
+//    }
 
     return partialQuery.build(
         new UnionDataSource(dataSources),
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java
index 054956c0ccdb..47f41162a084 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidFreeUnionDataSourceRule.java
@@ -80,7 +80,8 @@ public void onMatch(final RelOptRuleCall call)
       call.transformTo(
           DruidFreeUnionDataSourceRel.create(
               (Union) newUnionRel,
-              getColumnNames(firstDruidRel, plannerContext).get(),
+//              getColumnNames(firstDruidRel, plannerContext).get(),
+              unionRel.getRowType().getFieldNames(),
               firstDruidRel.getPlannerContext()
           )
       );
@@ -88,7 +89,8 @@ public void onMatch(final RelOptRuleCall call)
       call.transformTo(
           DruidFreeUnionDataSourceRel.create(
               unionRel,
-              getColumnNames(firstDruidRel, plannerContext).get(),
+//              getColumnNames(firstDruidRel, plannerContext).get(),
+              unionRel.getRowType().getFieldNames(),
               firstDruidRel.getPlannerContext()
           )
       );
@@ -111,7 +113,8 @@ public static boolean isCompatible(
     if (!unionRel.all && null != plannerContext) {
       plannerContext.setPlanningError("SQL requires 'UNION' but only 'UNION ALL' is supported.");
     }
-    return unionRel.all && isUnionCompatible(first, second, plannerContext);
+//    return unionRel.all && isUnionCompatible(first, second, plannerContext);
+    return unionRel.all;
   }
 
   /**