From 261c1f271f1b06fed79ed24e11adadd5fecf18c6 Mon Sep 17 00:00:00 2001
From: Jihoon Son <jihoonson@apache.org>
Date: Tue, 27 Apr 2021 18:45:23 -0700
Subject: [PATCH] Keep traitSet of logicalValues (#11138)

---
 .../druid/sql/calcite/rel/DruidQueryRel.java  |  4 +-
 .../druid/sql/calcite/CalciteQueryTest.java   | 70 +++++++++++++++++--
 2 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQueryRel.java b/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQueryRel.java
index dfbd6be0683..02273ce22e3 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQueryRel.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQueryRel.java
@@ -76,7 +76,7 @@ public class DruidQueryRel extends DruidRel<DruidQueryRel>
     return new DruidQueryRel(
         scanRel.getCluster(),
         scanRel.getCluster().traitSetOf(Convention.NONE),
-        table,
+        Preconditions.checkNotNull(table, "table"),
         druidTable,
         queryMaker,
         PartialDruidQuery.create(scanRel)
@@ -91,7 +91,7 @@ public class DruidQueryRel extends DruidRel<DruidQueryRel>
   {
     return new DruidQueryRel(
         valuesRel.getCluster(),
-        valuesRel.getCluster().traitSetOf(Convention.NONE),
+        valuesRel.getTraitSet(), // the traitSet of valuesRel should be kept
         null,
         druidTable,
         queryMaker,
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 e47a89a5f78..0ce2d2bd901 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
@@ -109,6 +109,7 @@ import org.apache.druid.query.topn.TopNQueryBuilder;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.segment.join.JoinType;
+import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
 import org.apache.druid.server.QueryLifecycle;
 import org.apache.druid.server.QueryLifecycleFactory;
 import org.apache.druid.server.security.Access;
@@ -150,7 +151,29 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
     // from a table.
     testQuery(
         "SELECT REGEXP_EXTRACT('foo', '^(.)')",
-        ImmutableList.of(),
+        ImmutableList.of(
+            Druids.newScanQueryBuilder()
+                  .dataSource(
+                      InlineDataSource.fromIterable(
+                          ImmutableList.of(new Object[]{0L}),
+                          RowSignature.builder().add("ZERO", ValueType.LONG).build()
+                      )
+                  )
+                  .intervals(querySegmentSpec(Filtration.eternity()))
+                  .virtualColumns(
+                      new ExpressionVirtualColumn(
+                          "v0",
+                          "'f'",
+                          ValueType.STRING,
+                          ExprMacroTable.nil()
+                      )
+                  )
+                  .columns("v0")
+                  .resultFormat(ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
+                  .legacy(false)
+                  .context(QUERY_CONTEXT_DEFAULT)
+                  .build()
+        ),
         ImmutableList.of(
             new Object[]{"f"}
         )
@@ -160,13 +183,32 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
   @Test
   public void testExpressionContainingNull() throws Exception
   {
-    List<String> expectedResult = new ArrayList<>();
-    expectedResult.add("Hello");
-    expectedResult.add(null);
     testQuery(
         "SELECT ARRAY ['Hello', NULL]",
-        ImmutableList.of(),
-        ImmutableList.of(new Object[]{expectedResult})
+        ImmutableList.of(
+            Druids.newScanQueryBuilder()
+                  .dataSource(
+                      InlineDataSource.fromIterable(
+                          ImmutableList.of(new Object[]{0L}),
+                          RowSignature.builder().add("ZERO", ValueType.LONG).build()
+                      )
+                  )
+                  .intervals(querySegmentSpec(Filtration.eternity()))
+                  .virtualColumns(
+                      new ExpressionVirtualColumn(
+                          "v0",
+                          "array('Hello',null)",
+                          ValueType.STRING,
+                          ExprMacroTable.nil()
+                      )
+                  )
+                  .columns("v0")
+                  .resultFormat(ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
+                  .legacy(false)
+                  .context(QUERY_CONTEXT_DEFAULT)
+                  .build()
+        ),
+        ImmutableList.of(new Object[]{"[\"Hello\",null]"})
     );
   }
 
@@ -7543,7 +7585,21 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
 
     testQuery(
         "SELECT dim2 FROM druid.foo ORDER BY dim2 LIMIT 0",
-        ImmutableList.of(),
+        ImmutableList.of(
+            Druids.newScanQueryBuilder()
+                  .dataSource(
+                      InlineDataSource.fromIterable(
+                          ImmutableList.of(),
+                          RowSignature.builder().add("dim2", ValueType.STRING).build()
+                      )
+                  )
+                  .intervals(querySegmentSpec(Filtration.eternity()))
+                  .columns("dim2")
+                  .resultFormat(ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
+                  .legacy(false)
+                  .context(QUERY_CONTEXT_DEFAULT)
+                  .build()
+        ),
         ImmutableList.of()
     );
   }