Implement key-based sampling in the planner

presto-docs/src/main/sphinx/functions/string.rst

@@ -238,3 +238,9 @@ Unicode Functions
     are replaced with `replace`. The replacement string `replace` must either
     be a single character or empty (in which case invalid characters are
     removed).
+
+.. function:: key_sampling_percent(varchar) -> double
+
+    Generates a double value between 0.0 and 1.0 based on the hash of the given ``varchar``.
+    This function is useful for deterministic sampling of data.
+

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -212,6 +212,9 @@ public final class SystemSessionProperties
 
     //TODO: Prestissimo related session properties that are temporarily put here. They will be relocated in the future
     public static final String PRESTISSIMO_SIMPLIFIED_EXPRESSION_EVALUATION_ENABLED = "simplified_expression_evaluation_enabled";
+    public static final String KEY_BASED_SAMPLING_ENABLED = "key_based_sampling_enabled";
+    public static final String KEY_BASED_SAMPLING_PERCENTAGE = "key_based_sampling_percentage";
+    public static final String KEY_BASED_SAMPLING_FUNCTION = "key_based_sampling_function";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -1142,7 +1145,22 @@ public SystemSessionProperties(
                         EXCEEDED_MEMORY_LIMIT_HEAP_DUMP_FILE_DIRECTORY,
                         "Directory to which heap snapshot will be dumped, if heap_dump_on_exceeded_memory_limit_enabled",
                         System.getProperty("java.io.tmpdir"),   // This is intended to be used for debugging purposes only and thus we does not need an associated config property
-                        true));
+                        true),
+                booleanProperty(
+                        KEY_BASED_SAMPLING_ENABLED,
+                        "Key based sampling of tables enabled",
+                        false,
+                        false),
+                doubleProperty(
+                        KEY_BASED_SAMPLING_PERCENTAGE,
+                        "Percentage of keys to be sampled",
+                        0.01,
+                        false),
+                stringProperty(
+                        KEY_BASED_SAMPLING_FUNCTION,
+                        "Sampling function for key based sampling",
+                        "key_sampling_percent",
+                        false));
     }
 
     public static boolean isEmptyJoinOptimization(Session session)
@@ -1165,6 +1183,21 @@ public static boolean isAllowWindowOrderByLiterals(Session session)
         return session.getSystemProperty(ALLOW_WINDOW_ORDER_BY_LITERALS, Boolean.class);
     }
 
+    public static boolean isKeyBasedSamplingEnabled(Session session)
+    {
+        return session.getSystemProperty(KEY_BASED_SAMPLING_ENABLED, Boolean.class);
+    }
+
+    public static double getKeyBasedSamplingPercentage(Session session)
+    {
+        return session.getSystemProperty(KEY_BASED_SAMPLING_PERCENTAGE, Double.class);
+    }
+
+    public static String getKeyBasedSamplingFunction(Session session)
+    {
+        return session.getSystemProperty(KEY_BASED_SAMPLING_FUNCTION, String.class);
+    }
+
     public List<PropertyMetadata<?>> getSessionProperties()
     {
         return sessionProperties;

presto-main/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java

@@ -166,6 +166,7 @@
 import com.facebook.presto.operator.scalar.WordStemFunction;
 import com.facebook.presto.operator.scalar.sql.ArraySqlFunctions;
 import com.facebook.presto.operator.scalar.sql.MapNormalizeFunction;
+import com.facebook.presto.operator.scalar.sql.SimpleSamplingPercent;
 import com.facebook.presto.operator.window.CumulativeDistributionFunction;
 import com.facebook.presto.operator.window.DenseRankFunction;
 import com.facebook.presto.operator.window.FirstValueFunction;
@@ -842,6 +843,7 @@ private List<? extends SqlFunction> getBuildInFunctions(FeaturesConfig featuresC
                 .sqlInvokedScalar(MapNormalizeFunction.class)
                 .sqlInvokedScalars(ArraySqlFunctions.class)
                 .sqlInvokedScalars(ArrayIntersectFunction.class)
+                .sqlInvokedScalars(SimpleSamplingPercent.class)
                 .scalar(DynamicFilterPlaceholderFunction.class)
                 .scalars(EnumCasts.class)
                 .scalars(LongEnumOperators.class)

presto-main/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java

@@ -0,0 +1,33 @@
+/*
+ * Licensed 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 com.facebook.presto.operator.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlType;
+
+public class SimpleSamplingPercent
+{
+    private SimpleSamplingPercent() {}
+
+    @SqlInvokedScalarFunction(value = "key_sampling_percent", deterministic = true, calledOnNullInput = false)
+    @Description("Returns a value between 0.0 and 1.0 using the hash of the given input string")
+    @SqlParameter(name = "input", type = "varchar")
+    @SqlType("double")
+    public static String keySamplingPercent()
+    {
+        return "return (abs(from_ieee754_64(xxhash64(cast(input as varbinary)))) % 100) / 100. ";
+    }
+}

presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java

@@ -121,6 +121,7 @@
 import com.facebook.presto.sql.planner.optimizations.HashGenerationOptimizer;
 import com.facebook.presto.sql.planner.optimizations.ImplementIntersectAndExceptAsUnion;
 import com.facebook.presto.sql.planner.optimizations.IndexJoinOptimizer;
+import com.facebook.presto.sql.planner.optimizations.KeyBasedSampler;
 import com.facebook.presto.sql.planner.optimizations.LimitPushDown;
 import com.facebook.presto.sql.planner.optimizations.MetadataDeleteOptimizer;
 import com.facebook.presto.sql.planner.optimizations.MetadataQueryOptimizer;
@@ -406,6 +407,7 @@ public PlanOptimizers(
                         estimatedExchangesCostCalculator,
                         ImmutableSet.of(new RewriteAggregationIfToFilter(metadata.getFunctionAndTypeManager()))),
                 predicatePushDown,
+                new KeyBasedSampler(metadata, sqlParser),
                 new IterativeOptimizer(
                         ruleStats,
                         statsCalculator,