SQL version of unnest native druid function

processing/src/main/java/org/apache/druid/query/UnnestDataSource.java

@@ -215,6 +215,18 @@ public int hashCode()
   {
     return Objects.hash(base, column, outputName);
   }
+
+  @Override
+  public String toString()
+  {
+    return "UnnestDataSource{" +
+           "base=" + base +
+           ", column='" + column + '\'' +
+           ", outputName='" + outputName + '\'' +
+           ", allowList=" + allowList +
+           '}';
+  }
+
 }
 
 

processing/src/main/java/org/apache/druid/segment/UnnestColumnValueSelectorCursor.java

@@ -100,7 +100,31 @@ public DimensionSelector makeDimensionSelector(DimensionSpec dimensionSpec)
         if (!outputName.equals(dimensionSpec.getDimension())) {
           return baseColumnSelectorFactory.makeDimensionSelector(dimensionSpec);
         }
-        throw new UOE("Unsupported dimension selector while using column value selector for column [%s]", outputName);
+        // this is done to support virtual columns
+        // In future a developer should move towards making sure that
+        // for all dictionary encoded cases we only get the dimension selector
+        return new BaseSingleValueDimensionSelector()
+        {
+          final ColumnValueSelector colSelector = makeColumnValueSelector(dimensionSpec.getDimension());
+
+          @Nullable
+          @Override
+          protected String getValue()
+          {
+            final Object returnedObj = colSelector.getObject();
+            if (returnedObj == null) {
+              return null;
+            } else {
+              return String.valueOf(returnedObj);
+            }
+          }
+
+          @Override
+          public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+          {
+            colSelector.inspectRuntimeShape(inspector);
+          }
+        };
       }
 
       @Override

processing/src/main/java/org/apache/druid/segment/UnnestDimensionCursor.java

@@ -83,7 +83,7 @@ public class UnnestDimensionCursor implements Cursor
   private final BitSet allowedBitSet;
   private final ColumnSelectorFactory baseColumnSelectorFactory;
   private int index;
-  private IndexedInts indexedIntsForCurrentRow;
+  @Nullable private IndexedInts indexedIntsForCurrentRow;
   private boolean needInitialization;
   private SingleIndexInts indexIntsForRow;
 
@@ -181,7 +181,7 @@ public void inspectRuntimeShape(RuntimeShapeInspector inspector)
           @Override
           public Object getObject()
           {
-            if (indexedIntsForCurrentRow == null) {
+            if (indexedIntsForCurrentRow == null || indexedIntsForCurrentRow.size() == 0) {
               return null;
             }
             if (allowedBitSet.isEmpty()) {
@@ -319,6 +319,7 @@ public void reset()
    * This would also create a bitset for dictonary encoded columns to
    * check for matching values specified in allowedList of UnnestDataSource.
    */
+  @Nullable
   private void initialize()
   {
     IdLookup idLookup = dimSelector.idLookup();
@@ -409,7 +410,12 @@ public int size()
     @Override
     public int get(int idx)
     {
-      return indexedIntsForCurrentRow.get(index);
+      // need to get value from the indexed ints
+      // only if it is non null and has at least 1 value
+      if (indexedIntsForCurrentRow != null && indexedIntsForCurrentRow.size() > 0) {
+        return indexedIntsForCurrentRow.get(index);
+      }
+      return 0;
     }
   }
 }

processing/src/test/java/org/apache/druid/segment/UnnestColumnValueSelectorCursorTest.java

@@ -19,7 +19,6 @@
 
 package org.apache.druid.segment;
 
-import org.apache.druid.java.util.common.UOE;
 import org.apache.druid.query.dimension.DefaultDimensionSpec;
 import org.apache.druid.testing.InitializedNullHandlingTest;
 import org.junit.Assert;
@@ -62,7 +61,7 @@ public void test_list_unnest_cursors()
     int j = 0;
     while (!unnestCursor.isDone()) {
       Object colSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(colSelectorVal.toString(), String.valueOf(j));
+      Assert.assertEquals(String.valueOf(j), colSelectorVal.toString());
       j++;
       unnestCursor.advance();
     }
@@ -96,7 +95,7 @@ public void test_list_unnest_cursors_user_supplied_list()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+      Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       k++;
       unnestCursor.advance();
     }
@@ -112,8 +111,6 @@ public void test_list_unnest_cursors_user_supplied_list_only_nulls()
         Collections.singletonList(null)
     );
 
-    List<String> expectedResults = Arrays.asList(null, null, null, null);
-
     //Create base cursor
     ListCursor listCursor = new ListCursor(inputList);
 
@@ -168,9 +165,9 @@ public void test_list_unnest_cursors_user_supplied_list_mixed_with_nulls()
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
       if (valueSelectorVal == null) {
-        Assert.assertEquals(null, expectedResults.get(k));
+        Assert.assertEquals(expectedResults.get(k), null);
       } else {
-        Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+        Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       }
       k++;
       unnestCursor.advance();
@@ -201,7 +198,7 @@ public void test_list_unnest_cursors_user_supplied_strings_and_no_lists()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+      Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       k++;
       unnestCursor.advance();
     }
@@ -231,7 +228,7 @@ public void test_list_unnest_cursors_user_supplied_strings_mixed_with_list()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+      Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       k++;
       unnestCursor.advance();
     }
@@ -265,7 +262,7 @@ public void test_list_unnest_cursors_user_supplied_lists_three_levels()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k).toString());
+      Assert.assertEquals(expectedResults.get(k).toString(), valueSelectorVal.toString());
       k++;
       unnestCursor.advance();
     }
@@ -306,7 +303,7 @@ public void test_list_unnest_of_unnest_cursors_user_supplied_list_three_levels()
     int k = 0;
     while (!parentCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k).toString());
+      Assert.assertEquals(expectedResults.get(k).toString(), valueSelectorVal.toString());
       k++;
       parentCursor.advance();
     }
@@ -344,12 +341,12 @@ public void test_list_unnest_cursors_user_supplied_list_with_nulls()
       if (valueSelectorVal == null) {
         Assert.assertEquals(null, expectedResults.get(k));
       } else {
-        Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+        Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       }
       k++;
       unnestCursor.advance();
     }
-    Assert.assertEquals(k, expectedResults.size());
+    Assert.assertEquals(expectedResults.size(), k);
   }
 
   @Test
@@ -382,7 +379,7 @@ public void test_list_unnest_cursors_user_supplied_list_with_dups()
       if (valueSelectorVal == null) {
         Assert.assertEquals(null, expectedResults.get(k));
       } else {
-        Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+        Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       }
       k++;
       unnestCursor.advance();
@@ -420,7 +417,7 @@ public void test_list_unnest_cursors_user_supplied_list_with_ignore_set()
       if (valueSelectorVal == null) {
         Assert.assertEquals(null, expectedResults.get(k));
       } else {
-        Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k));
+        Assert.assertEquals(expectedResults.get(k), valueSelectorVal.toString());
       }
       k++;
       unnestCursor.advance();
@@ -455,7 +452,7 @@ public void test_list_unnest_cursors_user_supplied_list_double()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Double valueSelectorVal = unnestColumnValueSelector.getDouble();
-      Assert.assertEquals(valueSelectorVal, expectedResults.get(k));
+      Assert.assertEquals(expectedResults.get(k), valueSelectorVal);
       k++;
       unnestCursor.advance();
     }
@@ -489,7 +486,7 @@ public void test_list_unnest_cursors_user_supplied_list_float()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Float valueSelectorVal = unnestColumnValueSelector.getFloat();
-      Assert.assertEquals(valueSelectorVal, expectedResults.get(k));
+      Assert.assertEquals(expectedResults.get(k), valueSelectorVal);
       k++;
       unnestCursor.advance();
     }
@@ -526,7 +523,7 @@ public void test_list_unnest_cursors_user_supplied_list_long()
       Object obj = unnestColumnValueSelector.getObject();
       Assert.assertNotNull(obj);
       Long valueSelectorVal = unnestColumnValueSelector.getLong();
-      Assert.assertEquals(valueSelectorVal, expectedResults.get(k));
+      Assert.assertEquals(expectedResults.get(k), valueSelectorVal);
       k++;
       unnestCursor.advance();
     }
@@ -561,7 +558,7 @@ public void test_list_unnest_cursors_user_supplied_list_three_level_arrays_and_m
     int k = 0;
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k).toString());
+      Assert.assertEquals(expectedResults.get(k).toString(), valueSelectorVal.toString());
       k++;
       unnestCursor.advance();
     }
@@ -570,17 +567,16 @@ public void test_list_unnest_cursors_user_supplied_list_three_level_arrays_and_m
     Assert.assertFalse(unnestCursor.isDoneOrInterrupted());
   }
 
-  @Test(expected = UOE.class)
+  @Test
   public void test_list_unnest_cursors_dimSelector()
   {
     List<Object> inputList = Arrays.asList(
         Arrays.asList("a", "b", "c"),
         Arrays.asList("e", "f", "g", "h", "i"),
-        Collections.singletonList("j")
+        Collections.singletonList(null)
     );
 
-    List<String> expectedResults = Arrays.asList("a", "b", "c", "e", "f", "g", "h", "i", "j");
-
+    List<Object> expectedResults = Arrays.asList("a", "b", "c", "e", "f", "g", "h", "i");
     //Create base cursor
     ListCursor listCursor = new ListCursor(inputList);
 
@@ -592,7 +588,25 @@ public void test_list_unnest_cursors_dimSelector()
         OUTPUT_NAME,
         IGNORE_SET
     );
-    unnestCursor.getColumnSelectorFactory().makeDimensionSelector(DefaultDimensionSpec.of(OUTPUT_NAME));
+    // should return a column value selector for this case
+    BaseSingleValueDimensionSelector unnestDimSelector = (BaseSingleValueDimensionSelector) unnestCursor.getColumnSelectorFactory()
+                                                                                                        .makeDimensionSelector(
+                                                                                                            DefaultDimensionSpec.of(
+                                                                                                                OUTPUT_NAME));
+    unnestDimSelector.inspectRuntimeShape(null);
+    int k = 0;
+    while (!unnestCursor.isDone()) {
+      if (k < 8) {
+        Assert.assertEquals(expectedResults.get(k).toString(), unnestDimSelector.getValue());
+      } else {
+        Assert.assertNull(unnestDimSelector.getValue());
+      }
+      k++;
+      unnestCursor.advance();
+    }
+    Assert.assertEquals(k, 9);
+    unnestCursor.reset();
+    Assert.assertNotNull(unnestDimSelector);
   }
 
   @Test
@@ -622,7 +636,7 @@ public void test_list_unnest_cursors_user_supplied_list_of_integers()
     int k = 0;
     while (!unnestCursor.isDone()) {
       Object valueSelectorVal = unnestColumnValueSelector.getObject();
-      Assert.assertEquals(valueSelectorVal.toString(), expectedResults.get(k).toString());
+      Assert.assertEquals(expectedResults.get(k).toString(), valueSelectorVal.toString());
       k++;
       unnestCursor.advance();
     }

sql/src/main/java/org/apache/druid/sql/calcite/expression/Expressions.java

@@ -24,6 +24,7 @@
 import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexFieldAccess;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
@@ -34,6 +35,7 @@
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.math.expr.Expr;
 import org.apache.druid.math.expr.ExprMacroTable;
@@ -58,6 +60,7 @@
 import org.apache.druid.sql.calcite.filtration.Filtration;
 import org.apache.druid.sql.calcite.planner.Calcites;
 import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.sql.calcite.rel.CannotBuildQueryException;
 import org.apache.druid.sql.calcite.rel.VirtualColumnRegistry;
 import org.apache.druid.sql.calcite.table.RowSignatures;
 import org.joda.time.Interval;
@@ -214,12 +217,51 @@ public static DruidExpression toDruidExpressionWithPostAggOperands(
       return rexCallToDruidExpression(plannerContext, rowSignature, rexNode, postAggregatorVisitor);
     } else if (kind == SqlKind.LITERAL) {
       return literalToDruidExpression(plannerContext, rexNode);
+    } else if (kind == SqlKind.FIELD_ACCESS) {
+      return fieldAccessToDruidExpression(rowSignature, rexNode);
     } else {
       // Can't translate.
       return null;
     }
   }
 
+  private static DruidExpression fieldAccessToDruidExpression(
+      final RowSignature rowSignature,
+      final RexNode rexNode
+  )
+  {
+    // Translate field references.
+    final RexFieldAccess ref = (RexFieldAccess) rexNode;
+    if (ref.getField().getIndex() > rowSignature.size()) {
+      // This case arises in the case of a correlation where the rexNode points to a table from the left subtree
+      // while the underlying datasource is the scan stub created from LogicalValuesRule
+      // In such a case we throw a CannotBuildQueryException so that Calcite does not go ahead with this path
+      // This exception is caught while returning false from isValidDruidQuery() method
+      throw new CannotBuildQueryException(StringUtils.format(
+          "Cannot build query as column name [%s] does not exist in row [%s]", ref.getField().getName(), rowSignature)
+      );
+    }
+
+    final String columnName = ref.getField().getName();
+    final int index = rowSignature.indexOf(columnName);
+
+    // This case arises when the rexNode has a name which is not in the underlying stub created using DruidUnnestDataSourceRule
+    // The column name has name ZERO with rowtype as LONG
+    // causes the index to be -1. In such a case we cannot build the query
+    // and throw an exception while returning false from isValidDruidQuery() method
+    if (index < 0) {
+      throw new CannotBuildQueryException(StringUtils.format(
+          "Expression referred to nonexistent index[%d] in row[%s]",
+          index,
+          rowSignature
+      ));
+    }
+
+    final Optional<ColumnType> columnType = rowSignature.getColumnType(index);
+
+    return DruidExpression.ofColumn(columnType.get(), columnName);
+  }
+
   private static DruidExpression inputRefToDruidExpression(
       final RowSignature rowSignature,
       final RexNode rexNode