SQL + Expressions = Best friends forever.

extensions-core/histogram/src/main/java/io/druid/query/aggregation/histogram/sql/QuantileSqlAggregator.java

@@ -19,7 +19,6 @@
 
 package io.druid.query.aggregation.histogram.sql;
 
-import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import io.druid.java.util.common.StringUtils;
 import io.druid.query.aggregation.AggregatorFactory;
@@ -28,13 +27,14 @@
 import io.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory;
 import io.druid.query.aggregation.histogram.QuantilePostAggregator;
 import io.druid.query.filter.DimFilter;
+import io.druid.segment.VirtualColumn;
 import io.druid.segment.column.ValueType;
+import io.druid.segment.virtual.ExpressionVirtualColumn;
 import io.druid.sql.calcite.aggregation.Aggregation;
 import io.druid.sql.calcite.aggregation.Aggregations;
 import io.druid.sql.calcite.aggregation.SqlAggregator;
+import io.druid.sql.calcite.expression.DruidExpression;
 import io.druid.sql.calcite.expression.Expressions;
-import io.druid.sql.calcite.expression.RowExtraction;
-import io.druid.sql.calcite.planner.DruidOperatorTable;
 import io.druid.sql.calcite.planner.PlannerContext;
 import io.druid.sql.calcite.table.RowSignature;
 import org.apache.calcite.rel.core.AggregateCall;
@@ -49,6 +49,7 @@
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
 
+import java.util.ArrayList;
 import java.util.List;
 
 public class QuantileSqlAggregator implements SqlAggregator
@@ -66,24 +67,23 @@ public SqlAggFunction calciteFunction()
   public Aggregation toDruidAggregation(
       final String name,
       final RowSignature rowSignature,
-      final DruidOperatorTable operatorTable,
       final PlannerContext plannerContext,
       final List<Aggregation> existingAggregations,
       final Project project,
       final AggregateCall aggregateCall,
       final DimFilter filter
   )
   {
-    final RowExtraction rex = Expressions.toRowExtraction(
+    final DruidExpression input = Expressions.toDruidExpression(
         plannerContext,
-        rowSignature.getRowOrder(),
+        rowSignature,
         Expressions.fromFieldAccess(
             rowSignature,
             project,
             aggregateCall.getArgList().get(0)
         )
     );
-    if (rex == null) {
+    if (input == null) {
       return null;
     }
 
@@ -119,17 +119,32 @@ public Aggregation toDruidAggregation(
             factory,
             filter,
             ApproximateHistogramAggregatorFactory.class,
-            new Predicate<ApproximateHistogramAggregatorFactory>()
-            {
-              @Override
-              public boolean apply(final ApproximateHistogramAggregatorFactory theFactory)
-              {
-                return theFactory.getFieldName().equals(rex.getColumn())
-                       && theFactory.getResolution() == resolution
-                       && theFactory.getNumBuckets() == numBuckets
-                       && theFactory.getLowerLimit() == lowerLimit
-                       && theFactory.getUpperLimit() == upperLimit;
+            theFactory -> {
+              // Check input for equivalence.
+              final boolean inputMatches;
+              final VirtualColumn virtualInput = existing.getVirtualColumns()
+                                                         .stream()
+                                                         .filter(
+                                                             virtualColumn ->
+                                                                 virtualColumn.getOutputName()
+                                                                              .equals(theFactory.getFieldName())
+                                                         )
+                                                         .findFirst()
+                                                         .orElse(null);
+
+              if (virtualInput == null) {
+                inputMatches = input.isDirectColumnAccess()
+                               && input.getDirectColumn().equals(theFactory.getFieldName());
+              } else {
+                inputMatches = ((ExpressionVirtualColumn) virtualInput).getExpression()
+                                                                       .equals(input.getExpression());
               }
+
+              return inputMatches
+                     && theFactory.getResolution() == resolution
+                     && theFactory.getNumBuckets() == numBuckets
+                     && theFactory.getLowerLimit() == lowerLimit
+                     && theFactory.getUpperLimit() == upperLimit;
             }
         );
 
@@ -143,19 +158,39 @@ public boolean apply(final ApproximateHistogramAggregatorFactory theFactory)
       }
     }
 
-    if (rowSignature.getColumnType(rex.getColumn()) == ValueType.COMPLEX) {
-      aggregatorFactory = new ApproximateHistogramFoldingAggregatorFactory(
-          histogramName,
-          rex.getColumn(),
-          resolution,
-          numBuckets,
-          lowerLimit,
-          upperLimit
-      );
+    // No existing match found. Create a new one.
+    final List<VirtualColumn> virtualColumns = new ArrayList<>();
+
+    if (input.isDirectColumnAccess()) {
+      if (rowSignature.getColumnType(input.getDirectColumn()) == ValueType.COMPLEX) {
+        aggregatorFactory = new ApproximateHistogramFoldingAggregatorFactory(
+            histogramName,
+            input.getDirectColumn(),
+            resolution,
+            numBuckets,
+            lowerLimit,
+            upperLimit
+        );
+      } else {
+        aggregatorFactory = new ApproximateHistogramAggregatorFactory(
+            histogramName,
+            input.getDirectColumn(),
+            resolution,
+            numBuckets,
+            lowerLimit,
+            upperLimit
+        );
+      }
     } else {
+      final ExpressionVirtualColumn virtualColumn = input.toVirtualColumn(
+          String.format("%s:v", name),
+          ValueType.FLOAT,
+          plannerContext.getExprMacroTable()
+      );
+      virtualColumns.add(virtualColumn);
       aggregatorFactory = new ApproximateHistogramAggregatorFactory(
           histogramName,
-          rex.getColumn(),
+          virtualColumn.getOutputName(),
           resolution,
           numBuckets,
           lowerLimit,
@@ -164,6 +199,7 @@ public boolean apply(final ApproximateHistogramAggregatorFactory theFactory)
     }
 
     return Aggregation.create(
+        virtualColumns,
         ImmutableList.of(aggregatorFactory),
         new QuantilePostAggregator(name, histogramName, probability)
     ).filter(filter);

extensions-core/histogram/src/test/java/io/druid/query/aggregation/histogram/sql/QuantileSqlAggregatorTest.java

@@ -35,16 +35,19 @@
 import io.druid.query.aggregation.histogram.ApproximateHistogramDruidModule;
 import io.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory;
 import io.druid.query.aggregation.histogram.QuantilePostAggregator;
+import io.druid.query.expression.TestExprMacroTable;
 import io.druid.query.filter.NotDimFilter;
 import io.druid.query.filter.SelectorDimFilter;
 import io.druid.query.spec.MultipleIntervalSegmentSpec;
 import io.druid.segment.IndexBuilder;
 import io.druid.segment.QueryableIndex;
 import io.druid.segment.TestHelper;
+import io.druid.segment.column.ValueType;
 import io.druid.segment.incremental.IncrementalIndexSchema;
+import io.druid.segment.virtual.ExpressionVirtualColumn;
 import io.druid.server.initialization.ServerConfig;
 import io.druid.sql.calcite.aggregation.SqlAggregator;
-import io.druid.sql.calcite.expression.SqlExtractionOperator;
+import io.druid.sql.calcite.expression.SqlOperatorConversion;
 import io.druid.sql.calcite.filtration.Filtration;
 import io.druid.sql.calcite.planner.Calcites;
 import io.druid.sql.calcite.planner.DruidOperatorTable;
@@ -131,7 +134,7 @@ public void setUp() throws Exception
     );
     final DruidOperatorTable operatorTable = new DruidOperatorTable(
         ImmutableSet.<SqlAggregator>of(new QuantileSqlAggregator()),
-        ImmutableSet.<SqlExtractionOperator>of()
+        ImmutableSet.<SqlOperatorConversion>of()
     );
     plannerFactory = new PlannerFactory(
         rootSchema,
@@ -159,6 +162,7 @@ public void testQuantileOnFloatAndLongs() throws Exception
                          + "APPROX_QUANTILE(m1, 0.5, 50),\n"
                          + "APPROX_QUANTILE(m1, 0.98, 200),\n"
                          + "APPROX_QUANTILE(m1, 0.99),\n"
+                         + "APPROX_QUANTILE(m1 * 2, 0.97),\n"
                          + "APPROX_QUANTILE(m1, 0.99) FILTER(WHERE dim1 = 'abc'),\n"
                          + "APPROX_QUANTILE(m1, 0.999) FILTER(WHERE dim1 <> 'abc'),\n"
                          + "APPROX_QUANTILE(m1, 0.999) FILTER(WHERE dim1 = 'abc'),\n"
@@ -170,7 +174,17 @@ public void testQuantileOnFloatAndLongs() throws Exception
       // Verify results
       final List<Object[]> results = Sequences.toList(plannerResult.run(), new ArrayList<Object[]>());
       final List<Object[]> expectedResults = ImmutableList.of(
-          new Object[]{1.0, 3.0, 5.880000114440918, 5.940000057220459, 6.0, 4.994999885559082, 6.0, 1.0}
+          new Object[]{
+              1.0,
+              3.0,
+              5.880000114440918,
+              5.940000057220459,
+              11.640000343322754,
+              6.0,
+              4.994999885559082,
+              6.0,
+              1.0
+          }
       );
       Assert.assertEquals(expectedResults.size(), results.size());
       for (int i = 0; i < expectedResults.size(); i++) {
@@ -183,28 +197,38 @@ public void testQuantileOnFloatAndLongs() throws Exception
                 .dataSource(CalciteTests.DATASOURCE1)
                 .intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(Filtration.eternity())))
                 .granularity(Granularities.ALL)
+                .virtualColumns(
+                    new ExpressionVirtualColumn(
+                        "a4:v",
+                        "(\"m1\" * 2)",
+                        ValueType.FLOAT,
+                        TestExprMacroTable.INSTANCE
+                    )
+                )
                 .aggregators(ImmutableList.of(
                     new ApproximateHistogramAggregatorFactory("a0:agg", "m1", null, null, null, null),
                     new ApproximateHistogramAggregatorFactory("a2:agg", "m1", 200, null, null, null),
+                    new ApproximateHistogramAggregatorFactory("a4:agg", "a4:v", null, null, null, null),
                     new FilteredAggregatorFactory(
-                        new ApproximateHistogramAggregatorFactory("a4:agg", "m1", null, null, null, null),
+                        new ApproximateHistogramAggregatorFactory("a5:agg", "m1", null, null, null, null),
                         new SelectorDimFilter("dim1", "abc", null)
                     ),
                     new FilteredAggregatorFactory(
-                        new ApproximateHistogramAggregatorFactory("a5:agg", "m1", null, null, null, null),
+                        new ApproximateHistogramAggregatorFactory("a6:agg", "m1", null, null, null, null),
                         new NotDimFilter(new SelectorDimFilter("dim1", "abc", null))
                     ),
-                    new ApproximateHistogramAggregatorFactory("a7:agg", "cnt", null, null, null, null)
+                    new ApproximateHistogramAggregatorFactory("a8:agg", "cnt", null, null, null, null)
                 ))
                 .postAggregators(ImmutableList.<PostAggregator>of(
                     new QuantilePostAggregator("a0", "a0:agg", 0.01f),
                     new QuantilePostAggregator("a1", "a0:agg", 0.50f),
                     new QuantilePostAggregator("a2", "a2:agg", 0.98f),
                     new QuantilePostAggregator("a3", "a0:agg", 0.99f),
-                    new QuantilePostAggregator("a4", "a4:agg", 0.99f),
-                    new QuantilePostAggregator("a5", "a5:agg", 0.999f),
-                    new QuantilePostAggregator("a6", "a4:agg", 0.999f),
-                    new QuantilePostAggregator("a7", "a7:agg", 0.50f)
+                    new QuantilePostAggregator("a4", "a4:agg", 0.97f),
+                    new QuantilePostAggregator("a5", "a5:agg", 0.99f),
+                    new QuantilePostAggregator("a6", "a6:agg", 0.999f),
+                    new QuantilePostAggregator("a7", "a5:agg", 0.999f),
+                    new QuantilePostAggregator("a8", "a8:agg", 0.50f)
                 ))
                 .context(ImmutableMap.<String, Object>of(
                     "skipEmptyBuckets", true,

pom.xml

@@ -61,7 +61,7 @@
         <!-- Curator cannot be updated to 2.12.0 yet, see https://github.com/druid-io/druid/pull/4103 -->
         <apache.curator.version>2.11.0</apache.curator.version>
         <avatica.version>1.9.0</avatica.version>
-        <calcite.version>1.11.0</calcite.version>
+        <calcite.version>1.12.0</calcite.version>
         <guava.version>16.0.1</guava.version>
         <guice.version>4.1.0</guice.version>
         <jetty.version>9.3.19.v20170502</jetty.version>