Tuple sketch SQL support

.gitignore

@@ -30,3 +30,9 @@ integration-tests/gen-scripts/
 *.hprof
 **/.ipynb_checkpoints/
 *.pyc
+
+# ignore NetBeans IDE specific files
+nbproject
+nbactions.xml
+nb-configuration.xml
+

docs/querying/sql-aggregations.md

@@ -139,6 +139,17 @@ Load the [DataSketches extension](../development/extensions-core/datasketches-ex
 |`DS_QUANTILES_SKETCH(expr, [k])`|Creates a [Quantiles sketch](../development/extensions-core/datasketches-quantiles.md) on the values of `expr`, which can be a regular column or a column containing quantiles sketches. The `k` parameter is described in the Quantiles sketch documentation.<br/><br/>See the [known issue](sql-translation.md#approximations) with this function.|`'0'` (STRING)|
 
 
+### Tuple sketch functions
+
+Load the [DataSketches extension](../development/extensions-core/datasketches-extension.md) to use the following functions.
+
+|Function|Notes|Default|
+|--------|-----|-------|
+|`ARRAY_OF_DOUBLES_SKETCH(expr, [nominalEntries])`|Creates a [Tuple sketch](../development/extensions-core/datasketches-tuple.md) on the values of `expr` which is a column containing Tuple sketches. The `nominalEntries` override parameter is optional and described in the Tuple sketch documentation.
+|`ARRAY_OF_DOUBLES_SKETCH(dimensionColumnExpr, metricColumnExpr..., [nominalEntries])`|Creates a [Tuple sketch](../development/extensions-core/datasketches-tuple.md) on the dimension value of `dimensionColumnExpr` and the metric values contained in the list of `metricColumnExpr` columns. The `nominalEntries` override parameter is optional and described in the Tuple sketch documentation.
+|`ARRAY_OF_DOUBLES_SKETCH_METRICS_SUM_ESTIMATE(expr)`|Computes approximate sums of the values contained within a [Tuple sketch](../development/extensions-core/datasketches-tuple.md#estimated-metrics-values-for-each-column-of-arrayofdoublessketch) column.
+
+
 ### T-Digest sketch functions
 
 Load the T-Digest extension to use the following functions. See the [T-Digest extension](../development/extensions-contrib/tdigestsketch-quantiles.md) for additional details and for more information on these functions.

docs/querying/sql-functions.md

@@ -136,6 +136,22 @@ Returns an array of all values of the specified expression.
 
 Concatenates array inputs into a single array.
 
+## ARRAY_OF_DOUBLES_SKETCH
+
+`ARRAY_OF_DOUBLES_SKETCH(expr, [nominalEntries])`
+
+`ARRAY_OF_DOUBLES_SKETCH(dimensionColumnExpr, metricColumnExpr..., [nominalEntries])`
+
+Creates a Tuple sketch
+
+**Function type:** [Aggregation](sql-aggregations.md)
+
+## ARRAY_OF_DOUBLES_SKETCH_METRICS_SUM_ESTIMATE
+
+`ARRAY_OF_DOUBLES_SKETCH_METRICS_SUM_ESTIMATE(expr)`
+
+Computes approximate sums of the values contained within a Tuple sketch
+
 ## ASIN
 
 `ASIN(<NUMERIC>)`

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchAggregatorFactory.java

@@ -91,7 +91,7 @@ public ArrayOfDoublesSketchAggregatorFactory(
   public Aggregator factorize(final ColumnSelectorFactory metricFactory)
   {
     if (metricColumns == null) { // input is sketches, use merge aggregator
-      final BaseObjectColumnValueSelector<ArrayOfDoublesSketch> selector = metricFactory
+      final BaseObjectColumnValueSelector<Object> selector = metricFactory
           .makeColumnValueSelector(fieldName);
       if (selector instanceof NilColumnValueSelector) {
         return new NoopArrayOfDoublesSketchAggregator(numberOfValues);

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchMergeAggregator.java

@@ -35,12 +35,12 @@
 public class ArrayOfDoublesSketchMergeAggregator implements Aggregator
 {
 
-  private final BaseObjectColumnValueSelector<ArrayOfDoublesSketch> selector;
+  private final BaseObjectColumnValueSelector<Object> selector;
   @Nullable
   private ArrayOfDoublesUnion union;
 
   public ArrayOfDoublesSketchMergeAggregator(
-      final BaseObjectColumnValueSelector<ArrayOfDoublesSketch> selector,
+      final BaseObjectColumnValueSelector<Object> selector,
       final int nominalEntries,
       final int numberOfValues
   )
@@ -58,12 +58,20 @@ public ArrayOfDoublesSketchMergeAggregator(
   @Override
   public void aggregate()
   {
-    final ArrayOfDoublesSketch update = selector.getObject();
+    final Object update = selector.getObject();
     if (update == null) {
       return;
     }
+    final ArrayOfDoublesSketch sketch;
+    if (update instanceof ArrayOfDoublesSketch) {
+      sketch = (ArrayOfDoublesSketch) update;
+    } else if (update instanceof String) {
+      sketch = ArrayOfDoublesSketchOperations.deserializeFromBase64EncodedStringSafe((String) update);
+    } else {
+      sketch = null;
+    }
     synchronized (this) {
-      union.union(update);
+      union.union(sketch);
     }
   }
 

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchModule.java

@@ -22,11 +22,18 @@
 import com.fasterxml.jackson.databind.Module;
 import com.fasterxml.jackson.databind.jsontype.NamedType;
 import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Binder;
 import org.apache.datasketches.tuple.arrayofdoubles.ArrayOfDoublesSketch;
 import org.apache.druid.initialization.DruidModule;
+import org.apache.druid.query.aggregation.datasketches.tuple.sql.ArrayOfDoublesSketchMetricsSumEstimateOperatorConversion;
+import org.apache.druid.query.aggregation.datasketches.tuple.sql.ArrayOfDoublesSketchSetIntersectOperatorConversion;
+import org.apache.druid.query.aggregation.datasketches.tuple.sql.ArrayOfDoublesSketchSetNotOperatorConversion;
+import org.apache.druid.query.aggregation.datasketches.tuple.sql.ArrayOfDoublesSketchSetUnionOperatorConversion;
+import org.apache.druid.query.aggregation.datasketches.tuple.sql.ArrayOfDoublesSketchSqlAggregator;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.serde.ComplexMetrics;
+import org.apache.druid.sql.guice.SqlBindings;
 
 import java.util.Collections;
 import java.util.List;
@@ -58,9 +65,14 @@ public class ArrayOfDoublesSketchModule implements DruidModule
   @Override
   public void configure(final Binder binder)
   {
-    ComplexMetrics.registerSerde(ARRAY_OF_DOUBLES_SKETCH, new ArrayOfDoublesSketchMergeComplexMetricSerde());
-    ComplexMetrics.registerSerde(ARRAY_OF_DOUBLES_SKETCH_MERGE_AGG, new ArrayOfDoublesSketchMergeComplexMetricSerde());
-    ComplexMetrics.registerSerde(ARRAY_OF_DOUBLES_SKETCH_BUILD_AGG, new ArrayOfDoublesSketchBuildComplexMetricSerde());
+    registerSerde();
+    SqlBindings.addAggregator(binder, ArrayOfDoublesSketchSqlAggregator.class);
+
+    SqlBindings.addOperatorConversion(binder, ArrayOfDoublesSketchMetricsSumEstimateOperatorConversion.class);
+    SqlBindings.addOperatorConversion(binder, ArrayOfDoublesSketchSetIntersectOperatorConversion.class);
+    SqlBindings.addOperatorConversion(binder, ArrayOfDoublesSketchSetUnionOperatorConversion.class);
+    SqlBindings.addOperatorConversion(binder, ArrayOfDoublesSketchSetNotOperatorConversion.class);
+
   }
 
   @Override
@@ -124,4 +136,13 @@ public List<? extends Module> getJacksonModules()
     );
   }
 
+  @VisibleForTesting
+  public static void registerSerde()
+  {
+    ComplexMetrics.registerSerde(ARRAY_OF_DOUBLES_SKETCH, new ArrayOfDoublesSketchMergeComplexMetricSerde());
+    ComplexMetrics.registerSerde(ARRAY_OF_DOUBLES_SKETCH_MERGE_AGG, new ArrayOfDoublesSketchMergeComplexMetricSerde());
+    ComplexMetrics.registerSerde(ARRAY_OF_DOUBLES_SKETCH_BUILD_AGG, new ArrayOfDoublesSketchBuildComplexMetricSerde());
+  }
+
+
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchToMetricsSumEstimatePostAggregator.java

@@ -55,6 +55,10 @@ public ArrayOfDoublesSketchToMetricsSumEstimatePostAggregator(
   public double[] compute(final Map<String, Object> combinedAggregators)
   {
     final ArrayOfDoublesSketch sketch = (ArrayOfDoublesSketch) getField().compute(combinedAggregators);
+    if (sketch == null)
+    {
+      return null;
+    }
     final SummaryStatistics[] stats = new SummaryStatistics[sketch.getNumValues()];
     Arrays.setAll(stats, i -> new SummaryStatistics());
     final ArrayOfDoublesSketchIterator it = sketch.iterator();

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/sql/ArrayOfDoublesSketchMetricsSumEstimateOperatorConversion.java

@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.aggregation.datasketches.tuple.sql;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.aggregation.PostAggregator;
+import org.apache.druid.query.aggregation.datasketches.tuple.ArrayOfDoublesSketchToMetricsSumEstimatePostAggregator;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.sql.calcite.expression.DruidExpression;
+import org.apache.druid.sql.calcite.expression.OperatorConversions;
+import org.apache.druid.sql.calcite.expression.PostAggregatorVisitor;
+import org.apache.druid.sql.calcite.expression.SqlOperatorConversion;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+
+import javax.annotation.Nullable;
+import java.util.List;
+
+public class ArrayOfDoublesSketchMetricsSumEstimateOperatorConversion implements SqlOperatorConversion
+{
+  private static final String FUNCTION_NAME = "ARRAY_OF_DOUBLES_SKETCH_METRICS_SUM_ESTIMATE";
+  private static final SqlFunction SQL_FUNCTION = OperatorConversions
+      .operatorBuilder(StringUtils.toUpperCase(FUNCTION_NAME))
+      .operandTypes(SqlTypeFamily.ANY)
+      .returnTypeNullableArrayWithNullableElements(SqlTypeName.DOUBLE)
+      .build();
+
+
+  @Override
+  public SqlOperator calciteOperator()
+  {
+    return SQL_FUNCTION;
+  }
+
+  @Override
+  public DruidExpression toDruidExpression(
+      PlannerContext plannerContext,
+      RowSignature rowSignature,
+      RexNode rexNode
+  )
+  {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public PostAggregator toPostAggregator(
+      PlannerContext plannerContext,
+      RowSignature rowSignature,
+      RexNode rexNode,
+      PostAggregatorVisitor postAggregatorVisitor
+  )
+  {
+    final List<RexNode> operands = ((RexCall) rexNode).getOperands();
+    final PostAggregator firstOperand = OperatorConversions.toPostAggregator(
+        plannerContext,
+        rowSignature,
+        operands.get(0),
+        postAggregatorVisitor,
+        true
+    );
+
+    if (firstOperand == null) {
+      return null;
+    }
+
+    return new ArrayOfDoublesSketchToMetricsSumEstimatePostAggregator(
+        postAggregatorVisitor.getOutputNamePrefix() + postAggregatorVisitor.getAndIncrementCounter(),
+        firstOperand
+    );
+  }
+
+}

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/sql/ArrayOfDoublesSketchSetBaseOperatorConversion.java

@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.aggregation.datasketches.tuple.sql;
+
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.SqlOperandCountRanges;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.aggregation.PostAggregator;
+import org.apache.druid.query.aggregation.datasketches.tuple.ArrayOfDoublesSketchSetOpPostAggregator;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.sql.calcite.expression.DruidExpression;
+import org.apache.druid.sql.calcite.expression.OperatorConversions;
+import org.apache.druid.sql.calcite.expression.PostAggregatorVisitor;
+import org.apache.druid.sql.calcite.expression.SqlOperatorConversion;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+
+import javax.annotation.Nullable;
+import java.util.ArrayList;
+import java.util.List;
+
+public abstract class ArrayOfDoublesSketchSetBaseOperatorConversion implements SqlOperatorConversion
+{
+  public ArrayOfDoublesSketchSetBaseOperatorConversion()
+  {
+  }
+
+  @Override
+  public SqlOperator calciteOperator()
+  {
+    return makeSqlFunction();
+  }
+
+  @Nullable
+  @Override
+  public DruidExpression toDruidExpression(
+      PlannerContext plannerContext,
+      RowSignature rowSignature,
+      RexNode rexNode
+  )
+  {
+    plannerContext.setPlanningError("%s can only be used on aggregates. " +
+        "It cannot be used directly on a column or on a scalar expression.", getFunctionName());
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public PostAggregator toPostAggregator(
+      PlannerContext plannerContext,
+      RowSignature rowSignature,
+      RexNode rexNode,
+      PostAggregatorVisitor postAggregatorVisitor
+  )
+  {
+    final List<RexNode> operands = ((RexCall) rexNode).getOperands();
+    final List<PostAggregator> inputPostAggs = new ArrayList<>();
+    Integer nominalEntries = null;
+    Integer numberOfvalues = null;
+
+    for (int i = 0; i < operands.size(); i++) {
+      RexNode operand = operands.get(i);
+      if (i == 0 && operand.isA(SqlKind.LITERAL) && SqlTypeFamily.INTEGER.contains(operand.getType())) {
+        nominalEntries = RexLiteral.intValue(operand);
+      } else if (i == 1 && operand.isA(SqlKind.LITERAL) && SqlTypeFamily.INTEGER.contains(operand.getType())) {
+        numberOfvalues = RexLiteral.intValue(operand);
+      } else {
+        final PostAggregator convertedPostAgg = OperatorConversions.toPostAggregator(
+            plannerContext,
+            rowSignature,
+            operand,
+            postAggregatorVisitor,
+            true
+        );
+
+        if (convertedPostAgg == null) {
+          return null;
+        } else {
+          inputPostAggs.add(convertedPostAgg);
+        }
+      }
+    }
+
+    return new ArrayOfDoublesSketchSetOpPostAggregator(
+        postAggregatorVisitor.getOutputNamePrefix() + postAggregatorVisitor.getAndIncrementCounter(),
+        getSetOperationName(),
+        nominalEntries,
+        numberOfvalues,
+        inputPostAggs
+    );
+  }
+
+  private SqlFunction makeSqlFunction()
+  {
+    return new SqlFunction(
+        getFunctionName(),
+        SqlKind.OTHER_FUNCTION,
+        ArrayOfDoublesSketchSqlOperators.RETURN_TYPE_INFERENCE,
+        null,
+        OperandTypes.variadic(SqlOperandCountRanges.from(2)),
+        SqlFunctionCategory.USER_DEFINED_FUNCTION
+    );
+  }
+
+  public abstract String getSetOperationName();
+
+  public String getFunctionName()
+  {
+    return StringUtils.format("ARRAY_OF_DOUBLES_SKETCH_%s", getSetOperationName());
+  }
+
+}