HIVE-23031 rewrite distinct

common/src/java/org/apache/hadoop/hive/conf/HiveConf.java

@@ -2465,6 +2465,19 @@ public static enum ConfVars {
         "If the number of references to a CTE clause exceeds this threshold, Hive will materialize it\n" +
         "before executing the main query block. -1 will disable this feature."),
 
+    HIVE_OPTIMIZE_BI_ENABLED("hive.optimize.bi.enabled", false,
+        "Enables query rewrites based on approximate functions(sketches)."),
+
+    HIVE_OPTIMIZE_BI_REWRITE_COUNTDISTINCT_ENABLED("hive.optimize.bi.rewrite.countdistinct.enabled",
+        true,
+        "Enables to rewrite COUNT(DISTINCT(X)) queries to be rewritten to use sketch functions."),
+
+    HIVE_OPTIMIZE_BI_REWRITE_COUNT_DISTINCT_SKETCH(
+        "hive.optimize.bi.rewrite.countdistinct.sketch", "hll",
+        new StringSet("hll", "cpc", "theta"),
+        "Defines which sketch type to use when rewriting COUNT(DISTINCT(X)) expressions. "
+            + "Distinct counting can be done with: hll,cpc or theta"),
+
     // Statistics
     HIVE_STATS_ESTIMATE_STATS("hive.stats.estimate", true,
         "Estimate statistics in absence of statistics."),

itests/src/test/resources/testconfiguration.properties

@@ -820,7 +820,9 @@ minillaplocal.query.files=\
   schq_ingest.q,\
   sketches_hll.q,\
   sketches_theta.q,\
+  sketches_rewrite.q,\
   sketches_materialized_view_rollup.q,\
+  sketches_materialized_view_rollup2.q,\
   table_access_keys_stats.q,\
   temp_table_llap_partitioned.q,\
   tez_bmj_schema_evolution.q,\

ql/src/java/org/apache/hadoop/hive/ql/exec/DataSketchesFunctions.java

@@ -18,21 +18,28 @@
 
 package org.apache.hadoop.hive.ql.exec;
 
+import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+
+import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeImpl;
 import org.apache.calcite.rel.type.RelProtoDataType;
 import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.type.InferTypes;
 import org.apache.calcite.sql.type.OperandTypes;
 import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.hadoop.hive.ql.optimizer.calcite.HiveTypeSystemImpl;
 import org.apache.hadoop.hive.ql.optimizer.calcite.functions.HiveMergeableAggregate;
+import org.apache.hadoop.hive.ql.optimizer.calcite.reloperators.HiveSqlFunction;
 import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFResolver2;
 import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
 import org.apache.hive.plugin.api.HiveUDFPlugin;
@@ -48,9 +55,9 @@ public final class DataSketchesFunctions implements HiveUDFPlugin {
 
   private static final String DATASKETCHES_PREFIX = "ds";
 
-  private static final String DATA_TO_SKETCH = "sketch";
+  public static final String DATA_TO_SKETCH = "sketch";
+  public static final String SKETCH_TO_ESTIMATE = "estimate";
   private static final String SKETCH_TO_ESTIMATE_WITH_ERROR_BOUNDS = "estimate_bounds";
-  private static final String SKETCH_TO_ESTIMATE = "estimate";
   private static final String SKETCH_TO_STRING = "stringify";
   private static final String UNION_SKETCH = "union";
   private static final String UNION_SKETCH1 = "union_f";
@@ -73,12 +80,12 @@ public final class DataSketchesFunctions implements HiveUDFPlugin {
   private static final String SKETCH_TO_VARIANCES = "variances";
   private static final String SKETCH_TO_PERCENTILE = "percentile";
 
-  private final List<SketchDescriptor> sketchClasses;
+  private final Map<String, SketchDescriptor> sketchClasses;
   private final ArrayList<UDFDescriptor> descriptors;
 
   private DataSketchesFunctions() {
-    this.sketchClasses = new ArrayList<SketchDescriptor>();
-    this.descriptors = new ArrayList<HiveUDFPlugin.UDFDescriptor>();
+    this.sketchClasses = new HashMap<>();
+    this.descriptors = new ArrayList<>();
     registerHll();
     registerCpc();
     registerKll();
@@ -96,19 +103,31 @@ public Iterable<UDFDescriptor> getDescriptors() {
     return descriptors;
   }
 
+  public SketchFunctionDescriptor getSketchFunction(String className, String function) {
+    if (!sketchClasses.containsKey(className)) {
+      throw new IllegalArgumentException(String.format("Sketch-class '%s' doesn't exists", className));
+    }
+    SketchDescriptor sc = sketchClasses.get(className);
+    if (!sc.fnMap.containsKey(function)) {
+      throw new IllegalArgumentException(String.format("The Sketch-class '%s' doesn't have a '%s' method", function));
+    }
+    return sketchClasses.get(className).fnMap.get(function);
+  }
+
   private void buildDescritors() {
-    for (SketchDescriptor sketchDescriptor : sketchClasses) {
+    for (SketchDescriptor sketchDescriptor : sketchClasses.values()) {
       descriptors.addAll(sketchDescriptor.fnMap.values());
     }
   }
 
   private void buildCalciteFns() {
-    for (SketchDescriptor sd : sketchClasses) {
+    for (SketchDescriptor sd : sketchClasses.values()) {
       // Mergability is exposed to Calcite; which enables to use it during rollup.
       RelProtoDataType sketchType = RelDataTypeImpl.proto(SqlTypeName.BINARY, true);
 
       SketchFunctionDescriptor sketchSFD = sd.fnMap.get(DATA_TO_SKETCH);
       SketchFunctionDescriptor unionSFD = sd.fnMap.get(UNION_SKETCH);
+      SketchFunctionDescriptor estimateSFD = sd.fnMap.get(SKETCH_TO_ESTIMATE);
 
       if (sketchSFD == null || unionSFD == null) {
         continue;
@@ -128,14 +147,27 @@ private void buildCalciteFns() {
           OperandTypes.family(),
           unionFn);
 
+
       unionSFD.setCalciteFunction(unionFn);
       sketchSFD.setCalciteFunction(sketchFn);
+      if (estimateSFD != null && estimateSFD.getReturnRelDataType().isPresent()) {
+
+        SqlFunction estimateFn = new HiveSqlFunction(estimateSFD.name,
+            SqlKind.OTHER_FUNCTION,
+            ReturnTypes.explicit(estimateSFD.getReturnRelDataType().get().getSqlTypeName()),
+            InferTypes.ANY_NULLABLE,
+            OperandTypes.family(),
+            SqlFunctionCategory.USER_DEFINED_FUNCTION,
+            true,
+            false);
+
+        estimateSFD.setCalciteFunction(estimateFn);
+      }
     }
   }
 
-
   private void registerHiveFunctionsInternal(Registry system) {
-    for (SketchDescriptor sketchDescriptor : sketchClasses) {
+    for (SketchDescriptor sketchDescriptor : sketchClasses.values()) {
       Collection<SketchFunctionDescriptor> functions = sketchDescriptor.fnMap.values();
       for (SketchFunctionDescriptor fn : functions) {
         if (UDF.class.isAssignableFrom(fn.udfClass)) {
@@ -165,6 +197,7 @@ private static class SketchFunctionDescriptor implements HiveUDFPlugin.UDFDescri
     String name;
     Class<?> udfClass;
     private SqlFunction calciteFunction;
+    private Class<?> returnType;
 
     public SketchFunctionDescriptor(String name, Class<?> udfClass) {
       this.name = name;
@@ -181,6 +214,19 @@ public String getFunctionName() {
       return name;
     }
 
+    public Optional<RelDataType> getReturnRelDataType() {
+      if (returnType == null) {
+        return Optional.empty();
+      } else {
+        JavaTypeFactoryImpl typeFactory = new JavaTypeFactoryImpl(new HiveTypeSystemImpl());
+        return Optional.of(typeFactory.createType(returnType));
+      }
+    }
+
+    public void setReturnType(Class<?> returnType) {
+      this.returnType = returnType;
+    }
+
     @Override
     public Optional<SqlFunction> getCalciteFunction() {
       return Optional.ofNullable(calciteFunction);
@@ -189,6 +235,11 @@ public Optional<SqlFunction> getCalciteFunction() {
     public void setCalciteFunction(SqlFunction calciteFunction) {
       this.calciteFunction = calciteFunction;
     }
+
+    @Override
+    public String toString() {
+      return getClass().getCanonicalName() + "[" + name + "]";
+    }
   }
 
   private static class SketchDescriptor {
@@ -201,7 +252,28 @@ public SketchDescriptor(String string) {
     }
 
     private void register(String name, Class<?> clazz) {
-      fnMap.put(name, new SketchFunctionDescriptor(functionPrefix + name, clazz));
+      SketchFunctionDescriptor value = new SketchFunctionDescriptor(functionPrefix + name, clazz);
+      if (UDF.class.isAssignableFrom(clazz)) {
+        Optional<Method> evaluateMethod = getEvaluateMethod(clazz);
+        if (evaluateMethod.isPresent()) {
+          value.setReturnType(evaluateMethod.get().getReturnType());
+        }
+      }
+      fnMap.put(name, value);
+    }
+
+    private Optional<Method> getEvaluateMethod(Class<?> clazz) {
+      List<Method> evaluateMethods = new ArrayList<Method>();
+      for (Method method : clazz.getMethods()) {
+        if ("evaluate".equals(method.getName())) {
+          evaluateMethods.add(method);
+        }
+      }
+      if (evaluateMethods.size() == 1) {
+        return Optional.of(evaluateMethods.get(0));
+      } else {
+        return Optional.empty();
+      }
     }
   }
 
@@ -214,7 +286,7 @@ private void registerHll() {
     sd.register(SKETCH_TO_STRING, org.apache.datasketches.hive.hll.SketchToStringUDF.class);
     sd.register(UNION_SKETCH1, org.apache.datasketches.hive.hll.UnionSketchUDF.class);
     sd.register(UNION_SKETCH, org.apache.datasketches.hive.hll.UnionSketchUDAF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("hll", sd);
   }
 
   private void registerCpc() {
@@ -228,7 +300,7 @@ private void registerCpc() {
     sd.register(SKETCH_TO_STRING, org.apache.datasketches.hive.cpc.SketchToStringUDF.class);
     sd.register(UNION_SKETCH1, org.apache.datasketches.hive.cpc.UnionSketchUDF.class);
     sd.register(UNION_SKETCH, org.apache.datasketches.hive.cpc.UnionSketchUDAF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("cpc", sd);
   }
 
   private void registerKll() {
@@ -244,7 +316,7 @@ private void registerKll() {
     sd.register(GET_QUANTILES, org.apache.datasketches.hive.kll.GetQuantilesUDF.class);
     sd.register(GET_QUANTILE, org.apache.datasketches.hive.kll.GetQuantileUDF.class);
     sd.register(GET_RANK, org.apache.datasketches.hive.kll.GetRankUDF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("kll", sd);
   }
 
   private void registerTheta() {
@@ -258,7 +330,7 @@ private void registerTheta() {
     sd.register(INTERSECT_SKETCH, org.apache.datasketches.hive.theta.IntersectSketchUDAF.class);
     sd.register(SKETCH_TO_ESTIMATE, org.apache.datasketches.hive.theta.EstimateSketchUDF.class);
     sd.register(EXCLUDE_SKETCH, org.apache.datasketches.hive.theta.ExcludeSketchUDF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("theta", sd);
 
   }
 
@@ -284,7 +356,7 @@ private void registerTupleArrayOfDoubles() {
         org.apache.datasketches.hive.tuple.ArrayOfDoublesSketchToQuantilesSketchUDF.class);
     sd.register(SKETCH_TO_VALUES, org.apache.datasketches.hive.tuple.ArrayOfDoublesSketchToValuesUDTF.class);
     sd.register(SKETCH_TO_VARIANCES, org.apache.datasketches.hive.tuple.ArrayOfDoublesSketchToVariancesUDF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("tuple_arrayofdouble", sd);
   }
 
   private void registerTupleDoubleSummary() {
@@ -294,7 +366,7 @@ private void registerTupleDoubleSummary() {
     sd.register(UNION_SKETCH, org.apache.datasketches.hive.tuple.UnionDoubleSummarySketchUDAF.class);
     sd.register(SKETCH_TO_ESTIMATE, org.apache.datasketches.hive.tuple.DoubleSummarySketchToEstimatesUDF.class);
     sd.register(SKETCH_TO_PERCENTILE, org.apache.datasketches.hive.tuple.DoubleSummarySketchToPercentileUDF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("tuple_doublesummary", sd);
   }
 
   private void registerQuantiles() {
@@ -312,7 +384,7 @@ private void registerFrequencies() {
     sd.register(UNION_SKETCH, org.apache.datasketches.hive.frequencies.UnionStringsSketchUDAF.class);
     sd.register(GET_FREQUENT_ITEMS,
         org.apache.datasketches.hive.frequencies.GetFrequentItemsFromStringsSketchUDTF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("freq", sd);
   }
 
   private void registerQuantilesString() {
@@ -327,7 +399,7 @@ private void registerQuantilesString() {
     sd.register(GET_PMF, org.apache.datasketches.hive.quantiles.GetPmfFromStringsSketchUDF.class);
     sd.register(GET_QUANTILE, org.apache.datasketches.hive.quantiles.GetQuantileFromStringsSketchUDF.class);
     sd.register(GET_QUANTILES, org.apache.datasketches.hive.quantiles.GetQuantilesFromStringsSketchUDF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("quantile_strings", sd);
   }
 
   private void registerQuantilesDoubles() {
@@ -342,7 +414,7 @@ private void registerQuantilesDoubles() {
     sd.register(GET_PMF, org.apache.datasketches.hive.quantiles.GetPmfFromDoublesSketchUDF.class);
     sd.register(GET_QUANTILE, org.apache.datasketches.hive.quantiles.GetQuantileFromDoublesSketchUDF.class);
     sd.register(GET_QUANTILES, org.apache.datasketches.hive.quantiles.GetQuantilesFromDoublesSketchUDF.class);
-    sketchClasses.add(sd);
+    sketchClasses.put("quantile_doubles", sd);
   }
 
 }