approx_top_k_combine

common/utils/src/main/resources/error/error-conditions.json

@@ -114,6 +114,18 @@
     ],
     "sqlState" : "22004"
   },
+  "APPROX_TOP_K_SKETCH_TYPE_UNMATCHED" : {
+    "message" : [
+      "Combining approx_top_k sketches of different types is not allowed. Found sketches of type <type1> and <type2>."
+    ],
+    "sqlState": "42846"
+  },
+  "APPROX_TOP_K_SKETCH_SIZE_UNMATCHED" : {
+    "message" : [
+      "Combining approx_top_k sketches of different sizes is not allowed. Found sketches of size <size1> and <size2>."
+    ],
+    "sqlState": "42846"
+  },
   "ARITHMETIC_OVERFLOW" : {
     "message" : [
       "<message>.<alternative> If necessary set <config> to \"false\" to bypass this error."

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -528,6 +528,8 @@ object FunctionRegistry {
     expression[HllSketchAgg]("hll_sketch_agg"),
     expression[HllUnionAgg]("hll_union_agg"),
     expression[ApproxTopK]("approx_top_k"),
+    expression[ApproxTopKAccumulate]("approx_top_k_accumulate"),
+    expression[ApproxTopKCombine]("approx_top_k_combine"),
 
     // string functions
     expression[Ascii]("ascii"),
@@ -786,6 +788,7 @@ object FunctionRegistry {
     expression[EqualNull]("equal_null"),
     expression[HllSketchEstimate]("hll_sketch_estimate"),
     expression[HllUnion]("hll_union"),
+    expression[ApproxTopKEstimate]("approx_top_k_estimate"),
 
     // grouping sets
     expression[Grouping]("grouping"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApproxTopKExpressions.scala

@@ -0,0 +1,114 @@
+/*
+ * 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.spark.sql.catalyst.expressions
+
+import org.apache.datasketches.frequencies.ItemsSketch
+import org.apache.datasketches.memory.Memory
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
+import org.apache.spark.sql.catalyst.expressions.aggregate.ApproxTopK
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.types._
+
+/**
+ * An expression that estimates the top K items from a sketch.
+ *
+ * The input is a sketch state that is generated by the ApproxTopKAccumulation function.
+ * The output is an array of structs, each containing a frequent item and its estimated frequency.
+ * The items are sorted by their estimated frequency in descending order.
+ *
+ * @param state The sketch state, which is a struct containing the serialized sketch data,
+ *              the original data type and the max items tracked of the sketch.
+ * @param k     The number of top items to estimate.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = """
+    _FUNC_(state, k) - Returns top k items with their frequency.
+      `k` An optional INTEGER literal greater than 0. If k is not specified, it defaults to 5.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_(approx_top_k_accumulate(expr)) FROM VALUES (0), (0), (1), (1), (2), (3), (4), (4) AS tab(expr);
+       [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+
+      > SELECT _FUNC_(approx_top_k_accumulate(expr), 2) FROM VALUES 'a', 'b', 'c', 'c', 'c', 'c', 'd', 'd' tab(expr);
+       [{"item":"c","count":4},{"item":"d","count":2}]
+  """,
+  group = "misc_funcs",
+  since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKEstimate(state: Expression, k: Expression)
+  extends BinaryExpression
+  with CodegenFallback
+  with ImplicitCastInputTypes {
+
+  def this(child: Expression, topK: Int) = this(child, Literal(topK))
+
+  def this(child: Expression) = this(child, Literal(ApproxTopK.DEFAULT_K))
+
+  private lazy val itemDataType: DataType = {
+    // itemDataType is the type of the "ItemTypeNull" field of the output of ACCUMULATE or COMBINE
+    state.dataType.asInstanceOf[StructType]("ItemTypeNull").dataType
+  }
+
+  override def left: Expression = state
+
+  override def right: Expression = k
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    val defaultCheck = super.checkInputDataTypes()
+    if (defaultCheck.isFailure) {
+      defaultCheck
+    } else if (!k.foldable) {
+      TypeCheckFailure("K must be a constant literal")
+    } else {
+      TypeCheckSuccess
+    }
+  }
+
+  override def dataType: DataType = ApproxTopK.getResultDataType(itemDataType)
+
+  override def eval(input: InternalRow): Any = {
+    // null check
+    ApproxTopK.checkExpressionNotNull(k, "k")
+    // eval
+    val stateEval = left.eval(input)
+    val kEval = right.eval(input)
+    val dataSketchBytes = stateEval.asInstanceOf[InternalRow].getBinary(0)
+    val maxItemsTrackedVal = stateEval.asInstanceOf[InternalRow].getInt(2)
+    val kVal = kEval.asInstanceOf[Int]
+    ApproxTopK.checkK(kVal)
+    ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal, kVal)
+    val itemsSketch = ItemsSketch.getInstance(
+      Memory.wrap(dataSketchBytes), ApproxTopK.genSketchSerDe(itemDataType))
+    ApproxTopK.genEvalResult(itemsSketch, kVal, itemDataType)
+  }
+
+  override protected def withNewChildrenInternal(newState: Expression, newK: Expression)
+  : Expression = copy(state = newState, k = newK)
+
+  override def nullable: Boolean = false
+
+  override def prettyName: String =
+    getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_estimate")
+}