dtenedor commented on code in PR #51298:
URL: https://github.com/apache/spark/pull/51298#discussion_r2308190303
##########
python/pyspark/sql/functions/builtin.py:
##########
@@ -25704,6 +25704,378 @@ def hll_union(
return _invoke_function("hll_union", _to_java_column(col1),
_to_java_column(col2))
+@_try_remote_functions
+def theta_sketch_agg(
Review Comment:
This PySpark integration part and Scala DataFrame support LGTM.
##########
sql/api/src/main/scala/org/apache/spark/sql/functions.scala:
##########
@@ -3552,6 +3715,154 @@ object functions {
hll_union(Column(columnName1), Column(columnName2),
allowDifferentLgConfigK)
}
+ /**
+ * Subtracts two binary representations of Datasketches ThetaSketch objects,
using a
+ * Datasketches AnotB object. Uses default log nominal entries.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_difference(c1: Column, c2: Column): Column =
+ Column.fn("theta_difference", c1, c2)
+
+ /**
+ * Subtracts two binary representations of Datasketches ThetaSketch objects,
using a
+ * Datasketches AnotB object. Uses default log nominal entries.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_difference(columnName1: String, columnName2: String): Column = {
+ theta_difference(Column(columnName1), Column(columnName2))
+ }
+
+ /**
+ * Subtracts two binary representations of Datasketches ThetaSketch objects,
using a
+ * Datasketches AnotB object. Allows setting of log nominal entries for the
difference buffer.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_difference(c1: Column, c2: Column, lgNomEntries: Int): Column =
+ Column.fn("theta_difference", c1, c2, lit(lgNomEntries))
+
+ /**
+ * Subtracts two binary representations of Datasketches ThetaSketch objects,
using a
+ * Datasketches AnotB object. Allows setting of log nominal entries for the
difference buffer.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_difference(columnName1: String, columnName2: String, lgNomEntries:
Int): Column = {
+ theta_difference(Column(columnName1), Column(columnName2), lgNomEntries)
+ }
+
+ /**
+ * Intersects two binary representations of Datasketches ThetaSketch
objects, using a
+ * Datasketches Intersection object. Uses default log nominal entries.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_intersection(c1: Column, c2: Column): Column =
+ Column.fn("theta_intersection", c1, c2)
+
+ /**
+ * Intersects two binary representations of Datasketches ThetaSketch
objects, using a
+ * Datasketches Intersection object. Uses default log nominal entries.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_intersection(columnName1: String, columnName2: String): Column = {
+ theta_intersection(Column(columnName1), Column(columnName2))
+ }
+
+ /**
+ * Intersects two binary representations of Datasketches ThetaSketch
objects, using a
+ * Datasketches Intersection object. Allows setting of log nominal entries
for the intersection
+ * buffer.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_intersection(c1: Column, c2: Column, lgNomEntries: Int): Column =
+ Column.fn("theta_intersection", c1, c2, lit(lgNomEntries))
+
+ /**
+ * Intersects two binary representations of Datasketches ThetaSketch
objects, using a
+ * Datasketches Intersection object. Allows setting of log nominal entries
for the intersection
+ * buffer.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_intersection(columnName1: String, columnName2: String,
lgNomEntries: Int): Column = {
+ theta_intersection(Column(columnName1), Column(columnName2), lgNomEntries)
+ }
+
+ /**
+ * Returns the estimated number of unique values given the binary
representation of a
+ * Datasketches ThetaSketch.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_sketch_estimate(c: Column): Column =
Column.fn("theta_sketch_estimate", c)
+
+ /**
+ * Returns the estimated number of unique values given the binary
representation of a
+ * Datasketches ThetaSketch.
+ *
+ * @group misc_funcs
+ * @since 4.1.0
+ */
+ def theta_sketch_estimate(columnName: String): Column = {
+ theta_sketch_estimate(Column(columnName))
+ }
+
+ /**
+ * Merges two binary representations of Datasketches ThetaSketch objects,
using a Datasketches
+ * Union object. Uses default log nominal entries.
Review Comment:
What do these mentions of log nominal entries mean?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala:
##########
@@ -0,0 +1,321 @@
+/*
+ * 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.common.SketchesArgumentException
+import org.apache.datasketches.memory.Memory
+import org.apache.datasketches.theta.{CompactSketch, SetOperation}
+
+import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes,
Expression, ExpressionDescription, Literal}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.util.ThetaSketchUtils
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.types.{AbstractDataType, BinaryType, DataType,
IntegerType, LongType}
+
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr) - Returns the estimated number of unique values given the
binary representation
+ of a Datasketches ThetaSketch. """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(theta_sketch_agg(col)) FROM VALUES (1), (1), (2), (2),
(3) tab(col);
+ 3
+ """,
+ group = "misc_funcs",
+ since = "4.1.0")
+case class ThetaSketchEstimate(child: Expression)
+ extends UnaryExpression
+ with CodegenFallback
+ with ExpectsInputTypes {
+ override def nullIntolerant: Boolean = true
+
+ override protected def withNewChildInternal(newChild: Expression):
ThetaSketchEstimate =
+ copy(child = newChild)
+
+ override def prettyName: String = "theta_sketch_estimate"
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType)
+
+ override def dataType: DataType = LongType
+
+ override def nullSafeEval(input: Any): Any = {
+ val buffer = input.asInstanceOf[Array[Byte]]
+
+ val memory = try {
+ Memory.wrap(buffer)
+ } catch {
+ case _: IllegalArgumentException | _: IndexOutOfBoundsException =>
+ throw QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val sketch = try {
+ CompactSketch.wrap(memory)
+ } catch {
+ case _: SketchesArgumentException =>
+ throw QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ Math.round(sketch.getEstimate)
+ }
+}
+
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(first, second, lgNomEntries) - Merges two binary representations of
+ Datasketches ThetaSketch objects, using a Datasketches Union object. Set
+ lgNomEntries to a value between 4 and 26 to find the unions of sketches
with different
+ union buffer sizes values (defaults to 12). """,
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(theta_sketch_agg(col1),
theta_sketch_agg(col2))) FROM VALUES (1, 4), (1, 4), (2, 5), (2, 5), (3, 6)
tab(col1, col2);
+ 6
+ """,
+ group = "misc_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaUnion(first: Expression, second: Expression, third: Expression)
+ extends TernaryExpression
+ with CodegenFallback
+ with ExpectsInputTypes {
+ override def nullIntolerant: Boolean = true
+
+ def this(first: Expression, second: Expression) = {
+ this(first, second, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS))
+ }
+
+ def this(first: Expression, second: Expression, third: Int) = {
+ this(first, second, Literal(third))
+ }
+
+ override protected def withNewChildrenInternal(
+ newFirst: Expression,
+ newSecond: Expression,
+ newThird: Expression): ThetaUnion =
+ copy(first = newFirst, second = newSecond, third = newThird)
+
+ override def prettyName: String = "theta_union"
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType, BinaryType,
IntegerType)
+
+ override def dataType: DataType = BinaryType
+
+ override def nullSafeEval(value1: Any, value2: Any, value3: Any): Any = {
+ val logNominalEntries = value3.asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(logNominalEntries)
+
+ val memory1 = try {
+ Memory.wrap(value1.asInstanceOf[Array[Byte]])
+ } catch {
+ case _: IllegalArgumentException | _: IndexOutOfBoundsException =>
+ throw QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val sketch1 = try {
+ CompactSketch.wrap(memory1)
+ } catch {
+ case _: SketchesArgumentException =>
+ throw QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val memory2 = try {
Review Comment:
This code from here until L148 seems like a dup of L122-L134 above. Should
we dedup it into a shared helper?
Same for L218 and L300 below. Could we make just one helper that we reuse 2
times for each relevant function in this file?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala:
##########
@@ -0,0 +1,682 @@
+/*
+ * 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.aggregate
+
+import org.apache.datasketches.common.SketchesArgumentException
+import org.apache.datasketches.memory.Memory
+import org.apache.datasketches.theta.{CompactSketch, Intersection,
SetOperation, Sketch, Union, UpdateSketch, UpdateSketchBuilder}
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes,
Expression, ExpressionDescription, Literal}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.catalyst.trees.BinaryLike
+import org.apache.spark.sql.catalyst.util.{ArrayData, CollationFactory,
ThetaSketchUtils}
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.internal.types.StringTypeWithCollation
+import org.apache.spark.sql.types.{AbstractDataType, ArrayType, BinaryType,
DataType, DoubleType, FloatType, IntegerType, LongType, StringType,
TypeCollection}
+import org.apache.spark.unsafe.types.UTF8String
+
+sealed trait ThetaSketchState {
+ def serialize(): Array[Byte]
+ def eval(): Array[Byte]
+}
+case class UpdatableSketchBuffer(sketch: UpdateSketch) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+ override def eval(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+}
+case class UnionAggregationBuffer(union: Union) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = union.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] = union.getResult.toByteArrayCompressed
+}
+case class IntersectionAggregationBuffer(intersection: Intersection) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+}
+case class FinalizedSketch(sketch: CompactSketch) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = sketch.toByteArrayCompressed
+ override def eval(): Array[Byte] = sketch.toByteArrayCompressed
+}
+
+/**
+ * The ThetaSketchAgg function utilizes a Datasketches ThetaSketch instance to
count a
+ * probabilistic approximation of the number of unique values in a given
column, and outputs the
+ * binary representation of the ThetaSketch.
+ *
+ * See [[https://datasketches.apache.org/docs/Theta/ThetaSketches.html]] for
more information.
+ *
+ * @param left
+ * child expression against which unique counting will occur
+ * @param right
+ * the log-base-2 of nomEntries decides the number of buckets for the sketch
+ * @param mutableAggBufferOffset
+ * offset for mutable aggregation buffer
+ * @param inputAggBufferOffset
+ * offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's compact binary
representation.
+ `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with
Nominal Entries deciding
+ the number buckets or slots for the ThetaSketch. """,
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(col, 12)) FROM VALUES (1), (1),
(2), (2), (3) tab(col);
+ 3
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaSketchAgg(
+ left: Expression,
+ right: Expression,
+ override val mutableAggBufferOffset: Int,
+ override val inputAggBufferOffset: Int)
+ extends TypedImperativeAggregate[ThetaSketchState]
+ with BinaryLike[Expression]
+ with ExpectsInputTypes {
+
+ // ThetaSketch config - mark as lazy so that they're not evaluated during
tree transformation.
+
+ lazy val lgNomEntries: Int = {
+ val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput)
+ lgNomEntriesInput
+ }
+
+ // Constructors
+
+ def this(child: Expression) = {
+ this(child, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS), 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Expression) = {
+ this(child, lgNomEntries, 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Int) = {
+ this(child, Literal(lgNomEntries), 0, 0)
+ }
+
+ // Copy constructors required by ImperativeAggregate
+
+ override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+ override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): ThetaSketchAgg =
+ copy(left = newLeft, right = newRight)
+
+ // Overrides for TypedImperativeAggregate
+
+ override def prettyName: String = "theta_sketch_agg"
+
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ TypeCollection(
+ IntegerType,
+ LongType,
+ FloatType,
+ DoubleType,
+ StringTypeWithCollation(supportsTrimCollation = true),
+ BinaryType,
+ ArrayType(IntegerType),
+ ArrayType(LongType)),
+ IntegerType)
+
+ override def dataType: DataType = BinaryType
+
+ override def nullable: Boolean = false
+
+ /**
+ * Instantiate an UpdateSketch instance using the lgNomEntries param.
+ *
+ * @return
+ * an UpdateSketch instance wrapped with UpdatableSketchBuffer
+ */
+ override def createAggregationBuffer(): ThetaSketchState = {
+ val builder = new UpdateSketchBuilder
+ builder.setLogNominalEntries(lgNomEntries)
+ UpdatableSketchBuffer(builder.build)
+ }
+
+ /**
+ * Evaluate the input row and update the UpdateSketch instance with the
row's value. The update
+ * function only supports a subset of Spark SQL types, and an exception will
be thrown for
+ * unsupported types.
+ *
+ * @param updateBuffer
+ * A previously initialized UpdateSketch instance
+ * @param input
+ * An input row
+ */
+ override def update(updateBuffer: ThetaSketchState, input: InternalRow):
ThetaSketchState =
+ updateBuffer match {
+ case UpdatableSketchBuffer(sketch) =>
+ val v = left.eval(input)
+ v match {
+ case null => UpdatableSketchBuffer(sketch)
+ case _ =>
+ left.dataType match {
+ case IntegerType =>
+ sketch.update(v.asInstanceOf[Int].toLong) // Promote to long
+ case LongType =>
+ sketch.update(v.asInstanceOf[Long])
+ case DoubleType =>
+ sketch.update(v.asInstanceOf[Double])
+ case FloatType =>
Review Comment:
please make sure that we have golden-file tests in the .sql time exercising
all of these supported input types, for each of the different functions.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala:
##########
@@ -0,0 +1,682 @@
+/*
+ * 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.aggregate
+
+import org.apache.datasketches.common.SketchesArgumentException
+import org.apache.datasketches.memory.Memory
+import org.apache.datasketches.theta.{CompactSketch, Intersection,
SetOperation, Sketch, Union, UpdateSketch, UpdateSketchBuilder}
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes,
Expression, ExpressionDescription, Literal}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.catalyst.trees.BinaryLike
+import org.apache.spark.sql.catalyst.util.{ArrayData, CollationFactory,
ThetaSketchUtils}
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.internal.types.StringTypeWithCollation
+import org.apache.spark.sql.types.{AbstractDataType, ArrayType, BinaryType,
DataType, DoubleType, FloatType, IntegerType, LongType, StringType,
TypeCollection}
+import org.apache.spark.unsafe.types.UTF8String
+
+sealed trait ThetaSketchState {
+ def serialize(): Array[Byte]
+ def eval(): Array[Byte]
+}
+case class UpdatableSketchBuffer(sketch: UpdateSketch) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+ override def eval(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+}
+case class UnionAggregationBuffer(union: Union) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = union.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] = union.getResult.toByteArrayCompressed
+}
+case class IntersectionAggregationBuffer(intersection: Intersection) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+}
+case class FinalizedSketch(sketch: CompactSketch) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = sketch.toByteArrayCompressed
+ override def eval(): Array[Byte] = sketch.toByteArrayCompressed
+}
+
+/**
+ * The ThetaSketchAgg function utilizes a Datasketches ThetaSketch instance to
count a
+ * probabilistic approximation of the number of unique values in a given
column, and outputs the
+ * binary representation of the ThetaSketch.
+ *
+ * See [[https://datasketches.apache.org/docs/Theta/ThetaSketches.html]] for
more information.
+ *
+ * @param left
+ * child expression against which unique counting will occur
+ * @param right
+ * the log-base-2 of nomEntries decides the number of buckets for the sketch
+ * @param mutableAggBufferOffset
+ * offset for mutable aggregation buffer
+ * @param inputAggBufferOffset
+ * offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's compact binary
representation.
+ `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with
Nominal Entries deciding
+ the number buckets or slots for the ThetaSketch. """,
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(col, 12)) FROM VALUES (1), (1),
(2), (2), (3) tab(col);
+ 3
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaSketchAgg(
+ left: Expression,
+ right: Expression,
+ override val mutableAggBufferOffset: Int,
+ override val inputAggBufferOffset: Int)
+ extends TypedImperativeAggregate[ThetaSketchState]
+ with BinaryLike[Expression]
+ with ExpectsInputTypes {
+
+ // ThetaSketch config - mark as lazy so that they're not evaluated during
tree transformation.
+
+ lazy val lgNomEntries: Int = {
+ val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput)
+ lgNomEntriesInput
+ }
+
+ // Constructors
+
+ def this(child: Expression) = {
+ this(child, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS), 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Expression) = {
+ this(child, lgNomEntries, 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Int) = {
+ this(child, Literal(lgNomEntries), 0, 0)
+ }
+
+ // Copy constructors required by ImperativeAggregate
+
+ override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+ override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): ThetaSketchAgg =
+ copy(left = newLeft, right = newRight)
+
+ // Overrides for TypedImperativeAggregate
+
+ override def prettyName: String = "theta_sketch_agg"
+
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ TypeCollection(
+ IntegerType,
+ LongType,
+ FloatType,
+ DoubleType,
+ StringTypeWithCollation(supportsTrimCollation = true),
+ BinaryType,
+ ArrayType(IntegerType),
+ ArrayType(LongType)),
+ IntegerType)
+
+ override def dataType: DataType = BinaryType
+
+ override def nullable: Boolean = false
+
+ /**
+ * Instantiate an UpdateSketch instance using the lgNomEntries param.
+ *
+ * @return
+ * an UpdateSketch instance wrapped with UpdatableSketchBuffer
+ */
+ override def createAggregationBuffer(): ThetaSketchState = {
+ val builder = new UpdateSketchBuilder
+ builder.setLogNominalEntries(lgNomEntries)
+ UpdatableSketchBuffer(builder.build)
+ }
+
+ /**
+ * Evaluate the input row and update the UpdateSketch instance with the
row's value. The update
+ * function only supports a subset of Spark SQL types, and an exception will
be thrown for
+ * unsupported types.
+ *
+ * @param updateBuffer
+ * A previously initialized UpdateSketch instance
+ * @param input
+ * An input row
+ */
+ override def update(updateBuffer: ThetaSketchState, input: InternalRow):
ThetaSketchState =
+ updateBuffer match {
+ case UpdatableSketchBuffer(sketch) =>
+ val v = left.eval(input)
+ v match {
+ case null => UpdatableSketchBuffer(sketch)
+ case _ =>
+ left.dataType match {
+ case IntegerType =>
+ sketch.update(v.asInstanceOf[Int].toLong) // Promote to long
+ case LongType =>
+ sketch.update(v.asInstanceOf[Long])
+ case DoubleType =>
+ sketch.update(v.asInstanceOf[Double])
+ case FloatType =>
+ sketch.update(v.asInstanceOf[Float].toDouble) // Promote to
double
+ case st: StringType =>
+ val cKey =
+ CollationFactory.getCollationKey(v.asInstanceOf[UTF8String],
st.collationId)
+ sketch.update(cKey.toString)
+ case BinaryType =>
+ val bytes = v.asInstanceOf[Array[Byte]]
+ if (bytes.nonEmpty) sketch.update(bytes)
+ case ArrayType(IntegerType, _) =>
+ val arr = v.asInstanceOf[ArrayData].toIntArray()
+ if (arr.nonEmpty) sketch.update(arr)
+ case ArrayType(LongType, _) =>
+ val arr = v.asInstanceOf[ArrayData].toLongArray()
+ if (arr.nonEmpty) sketch.update(arr)
+ case _ =>
+ throw new SparkUnsupportedOperationException(
+ errorClass = "_LEGACY_ERROR_TEMP_3121",
+ messageParameters = Map("dataType" ->
left.dataType.toString))
+ }
+ UpdatableSketchBuffer(sketch) // Return updated sketch wrapped
again
+ }
+ case _ =>
+ updateBuffer
+ }
+
+ /**
+ * Merges an input Compact sketch into the UpdateSketch which is acting as
the aggregation
+ * buffer.
+ *
+ * @param updateBuffer
+ * the UpdateSketch or Union instance used to store the aggregation result
+ * @param input
+ * An input UpdateSketch, Union, or Compact sketch instance
+ */
+ override def merge(
+ updateBuffer: ThetaSketchState,
+ input: ThetaSketchState): ThetaSketchState = {
+ // Helper function to create union only when needed
+ def createUnionWith(sketch1: Sketch, sketch2: Sketch):
UnionAggregationBuffer = {
+ val union =
SetOperation.builder.setLogNominalEntries(lgNomEntries).buildUnion
+ union.union(sketch1)
+ union.union(sketch2)
+ UnionAggregationBuffer(union)
+ }
+
+ (updateBuffer, input) match {
+ // REUSE existing union - this is the most efficient path
+ case (UnionAggregationBuffer(existingUnion),
UpdatableSketchBuffer(sketch)) =>
+ existingUnion.union(sketch.compact)
+ UnionAggregationBuffer(existingUnion)
+ case (UnionAggregationBuffer(existingUnion), FinalizedSketch(sketch)) =>
+ existingUnion.union(sketch)
+ UnionAggregationBuffer(existingUnion)
+ case (UnionAggregationBuffer(union1), UnionAggregationBuffer(union2)) =>
+ union1.union(union2.getResult)
+ UnionAggregationBuffer(union1)
+ // CREATE new union only when necessary
+ case (UpdatableSketchBuffer(sketch1), UpdatableSketchBuffer(sketch2)) =>
+ createUnionWith(sketch1.compact, sketch2.compact)
+ case (UpdatableSketchBuffer(sketch1), FinalizedSketch(sketch2)) =>
+ createUnionWith(sketch1.compact, sketch2)
+ // Should never make it here, but added cases for defensive programming
+ case (FinalizedSketch(sketch1), UpdatableSketchBuffer(sketch2)) =>
+ createUnionWith(sketch1, sketch2.compact)
+ case (FinalizedSketch(sketch1), FinalizedSketch(sketch2)) =>
+ createUnionWith(sketch1, sketch2)
+ case _ => throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+ }
+
+ /**
+ * Returns a Compact sketch derived from the input column or expression
+ *
+ * @param sketchState
+ * Union instance used as an aggregation buffer
+ * @return
+ * A Compact binary sketch
+ */
+ override def eval(sketchState: ThetaSketchState): Any = {
+ sketchState.eval()
+ }
+
+ /** Convert the underlying UpdateSketch/Union into an Compact byte array */
+ override def serialize(sketchState: ThetaSketchState): Array[Byte] = {
+ sketchState.serialize()
+ }
+
+ /** Wrap the byte array into a Compact sketch instance */
+ override def deserialize(buffer: Array[Byte]): ThetaSketchState = {
+ if (buffer.nonEmpty) {
+ FinalizedSketch(CompactSketch.heapify(Memory.wrap(buffer)))
+ } else {
+ this.createAggregationBuffer()
+ }
+ }
+}
+
+/**
+ * The ThetaUnionAgg function ingests and merges Datasketches ThetaSketch
instances previously
+ * produced by the ThetaSketchAgg function, and outputs the merged ThetaSketch.
+ *
+ * See [[https://datasketches.apache.org/docs/Theta/ThetaSketches.html]] for
more information.
+ *
+ * @param left
+ * Child expression against which unique counting will occur
+ * @param right
+ * the log-base-2 of nomEntries decides the number of buckets for the sketch
+ * @param mutableAggBufferOffset
+ * offset for mutable aggregation buffer
+ * @param inputAggBufferOffset
+ * offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's Compact binary
representation.
+ `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with
Nominal Entries deciding
+ the number buckets or slots for the ThetaSketch.""",
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(sketch)) FROM (SELECT
theta_sketch_agg(col) as sketch FROM VALUES (1) tab(col) UNION ALL SELECT
theta_sketch_agg(col, 20) as sketch FROM VALUES (1) tab(col));
+ 1
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaUnionAgg(
+ left: Expression,
+ right: Expression,
+ override val mutableAggBufferOffset: Int,
+ override val inputAggBufferOffset: Int)
+ extends TypedImperativeAggregate[ThetaSketchState]
+ with BinaryLike[Expression]
+ with ExpectsInputTypes {
+
+ // ThetaSketch config - mark as lazy so that they're not evaluated during
tree transformation.
+
+ lazy val lgNomEntries: Int = {
+ val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput)
+ lgNomEntriesInput
+ }
+
+ // Constructors
+
+ def this(child: Expression) = {
+ this(child, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS), 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Expression) = {
+ this(child, lgNomEntries, 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Int) = {
+ this(child, Literal(lgNomEntries), 0, 0)
+ }
+
+ // Copy constructors required by ImperativeAggregate
+
+ override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ThetaUnionAgg =
+ copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+ override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int):
ThetaUnionAgg =
+ copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): ThetaUnionAgg =
+ copy(left = newLeft, right = newRight)
+
+ // Overrides for TypedImperativeAggregate
+
+ override def prettyName: String = "theta_union_agg"
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType, IntegerType)
+
+ override def dataType: DataType = BinaryType
+
+ override def nullable: Boolean = false
+
+ /**
+ * Instantiate a Union instance using the lgNomEntries param.
+ *
+ * @return
+ * a Union instance wrapped with UnionAggregationBuffer
+ */
+ override def createAggregationBuffer(): ThetaSketchState = {
+ UnionAggregationBuffer(
+ SetOperation.builder
+ .setLogNominalEntries(lgNomEntries)
+ .buildUnion)
+ }
+
+ /**
+ * Update the Union instance with the Compact sketch byte array obtained
from the row.
+ *
+ * @param unionBuffer
+ * A previously initialized Union instance
+ * @param input
+ * An input row
+ */
+ override def update(unionBuffer: ThetaSketchState, input: InternalRow):
ThetaSketchState = {
+ val v = left.eval(input)
+ v match {
+ case null => unionBuffer // Input is null, return buffer unchanged
+ case _ =>
+ left.dataType match {
+ case BinaryType =>
+ val memory = try {
+ Memory.wrap(v.asInstanceOf[Array[Byte]])
+ } catch {
+ case _: IllegalArgumentException | _: IndexOutOfBoundsException
=>
+ throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val inputSketch = try {
+ CompactSketch.wrap(memory)
+ } catch {
+ case _: SketchesArgumentException =>
+ throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val union = unionBuffer match {
+ case UnionAggregationBuffer(existingUnionBuffer) =>
existingUnionBuffer
+ case _ => throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+ union.union(inputSketch)
+ UnionAggregationBuffer(union)
+ case _ =>
+ throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+ }
+ }
+
+ /**
+ * Merges an input Compact sketch into the Union which is acting as the
aggregation buffer.
+ *
+ * @param unionBuffer
+ * The Union instance used to store the aggregation result
+ * @param input
+ * An input Union or Compact sketch instance
+ */
+ override def merge(unionBuffer: ThetaSketchState, input: ThetaSketchState):
ThetaSketchState = {
+ (unionBuffer, input) match {
+ // Both are unions, merge them directly
+ case (UnionAggregationBuffer(unionLeft),
UnionAggregationBuffer(unionRight)) =>
+ unionLeft.union(unionRight.getResult)
+ UnionAggregationBuffer(unionLeft)
+ // The input was serialized then deserialized
+ case (UnionAggregationBuffer(union), FinalizedSketch(sketch)) =>
+ union.union(sketch)
+ UnionAggregationBuffer(union)
+ // Should never make it here, but added cases for defensive programming
+ case (FinalizedSketch(sketch1), FinalizedSketch(sketch2)) =>
+ val union =
SetOperation.builder.setLogNominalEntries(lgNomEntries).buildUnion
+ union.union(sketch1)
+ union.union(sketch2)
+ UnionAggregationBuffer(union)
+ case (FinalizedSketch(sketch), UnionAggregationBuffer(union)) =>
+ union.union(sketch)
+ UnionAggregationBuffer(union)
+ case _ => throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+ }
+
+ /**
+ * Returns a Compact sketch derived from the merged sketches
+ *
+ * @param sketchState
+ * Union instance used as an aggregation buffer
+ * @return
+ * A Compact binary sketch
+ */
+ override def eval(sketchState: ThetaSketchState): Any = {
+ sketchState.eval()
+ }
+
+ /** Convert the underlying Union into an Compact byte array */
+ override def serialize(sketchState: ThetaSketchState): Array[Byte] = {
+ sketchState.serialize()
+ }
+
+ /** Wrap the byte array into a Compact sketch instance */
+ override def deserialize(buffer: Array[Byte]): ThetaSketchState = {
+ if (buffer.nonEmpty) {
+ FinalizedSketch(CompactSketch.heapify(Memory.wrap(buffer)))
+ } else {
+ this.createAggregationBuffer()
+ }
+ }
+}
+
+/**
+ * The ThetaIntersectionAgg function ingests and intersects Datasketches
ThetaSketch instances
+ * previously produced by the ThetaSketchAgg function, and outputs the
intersected ThetaSketch.
+ *
+ * See [[https://datasketches.apache.org/docs/Theta/ThetaSketches.html]] for
more information.
+ *
+ * @param left
+ * Child expression against which unique counting will occur
+ * @param right
+ * the log-base-2 of nomEntries decides the number of buckets for the sketch
+ * @param mutableAggBufferOffset
+ * offset for mutable aggregation buffer
+ * @param inputAggBufferOffset
+ * offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's Compact binary
representation.
+ `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with
Nominal Entries deciding
+ the number buckets or slots for the ThetaSketch.""",
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(sketch)) FROM (SELECT
theta_sketch_agg(col) as sketch FROM VALUES (1) tab(col) UNION ALL SELECT
theta_sketch_agg(col, 20) as sketch FROM VALUES (1) tab(col));
+ 1
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaIntersectionAgg(
+ left: Expression,
+ right: Expression,
+ override val mutableAggBufferOffset: Int,
+ override val inputAggBufferOffset: Int)
+ extends TypedImperativeAggregate[ThetaSketchState]
+ with BinaryLike[Expression]
+ with ExpectsInputTypes {
+
+ // ThetaSketch config - mark as lazy so that they're not evaluated during
tree transformation.
+
+ lazy val lgNomEntries: Int = {
+ val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput)
+ lgNomEntriesInput
+ }
+
+ // Constructors
+
+ def this(child: Expression) = {
+ this(child, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS), 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Expression) = {
+ this(child, lgNomEntries, 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Int) = {
+ this(child, Literal(lgNomEntries), 0, 0)
+ }
+
+ // Copy constructors required by ImperativeAggregate
+
+ override def withNewMutableAggBufferOffset(
+ newMutableAggBufferOffset: Int): ThetaIntersectionAgg =
+ copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+ override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int):
ThetaIntersectionAgg =
+ copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): ThetaIntersectionAgg =
+ copy(left = newLeft, right = newRight)
+
+ // Overrides for TypedImperativeAggregate
+
+ override def prettyName: String = "theta_intersection_agg"
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType, IntegerType)
+
+ override def dataType: DataType = BinaryType
+
+ override def nullable: Boolean = false
+
+ /**
+ * Instantiate an Intersection instance using the lgNomEntries param.
+ *
+ * @return
+ * an Intersection instance wrapped with IntersectionAggregationBuffer
+ */
+ override def createAggregationBuffer(): ThetaSketchState = {
+ IntersectionAggregationBuffer(
+ SetOperation.builder
+ .setLogNominalEntries(lgNomEntries)
+ .buildIntersection)
+ }
+
+ /**
+ * Update the Intersection instance with the Compact sketch byte array
obtained from the row.
+ *
+ * @param intersectionBuffer
+ * A previously initialized Intersection instance
+ * @param input
+ * An input row
+ */
+ override def update(
+ intersectionBuffer: ThetaSketchState,
+ input: InternalRow): ThetaSketchState = {
+ val v = left.eval(input)
+ v match {
+ case null => intersectionBuffer // Input is null, return buffer unchanged
+ case _ =>
+ left.dataType match {
+ case BinaryType =>
+ val memory = try {
+ Memory.wrap(v.asInstanceOf[Array[Byte]])
+ } catch {
+ case _: IllegalArgumentException | _: IndexOutOfBoundsException
=>
+ throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val inputSketch = try {
+ CompactSketch.wrap(memory)
+ } catch {
+ case _: SketchesArgumentException =>
+ throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+
+ val intersection = intersectionBuffer match {
+ case IntersectionAggregationBuffer(existingIntersection) =>
existingIntersection
+ case _ => throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+ intersection.intersect(inputSketch)
+ IntersectionAggregationBuffer(intersection)
+ case _ =>
+ throw
QueryExecutionErrors.thetaInvalidInputSketchBuffer(prettyName)
+ }
+ }
+ }
+
+ /**
+ * Merges an input Compact sketch into the Intersection which is acting as
the aggregation
+ * buffer.
+ *
+ * @param intersectionBuffer
+ * The Intersection instance used to store the aggregation result
+ * @param input
+ * An input Intersection or Compact sketch instance
+ */
+ override def merge(
+ intersectionBuffer: ThetaSketchState,
+ input: ThetaSketchState): ThetaSketchState = {
+ (intersectionBuffer, input) match {
+ // Both are intersections, merge them directly
+ case (
+ IntersectionAggregationBuffer(intersectLeft),
+ IntersectionAggregationBuffer(intersectRight)) =>
+ intersectLeft.intersect(intersectRight.getResult)
+ IntersectionAggregationBuffer(intersectLeft)
+ // The input was serialized then deserialized
+ case (IntersectionAggregationBuffer(intersection),
FinalizedSketch(sketch)) =>
+ intersection.intersect(sketch)
+ IntersectionAggregationBuffer(intersection)
+ // Should never make it here, but added cases for defensive programming
Review Comment:
super nit, could we please make every comment in the PR comprise complete
sentences, each starting with a capital letter and ending with punctuation?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala:
##########
@@ -0,0 +1,682 @@
+/*
+ * 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.aggregate
+
+import org.apache.datasketches.common.SketchesArgumentException
+import org.apache.datasketches.memory.Memory
+import org.apache.datasketches.theta.{CompactSketch, Intersection,
SetOperation, Sketch, Union, UpdateSketch, UpdateSketchBuilder}
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes,
Expression, ExpressionDescription, Literal}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.catalyst.trees.BinaryLike
+import org.apache.spark.sql.catalyst.util.{ArrayData, CollationFactory,
ThetaSketchUtils}
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.internal.types.StringTypeWithCollation
+import org.apache.spark.sql.types.{AbstractDataType, ArrayType, BinaryType,
DataType, DoubleType, FloatType, IntegerType, LongType, StringType,
TypeCollection}
+import org.apache.spark.unsafe.types.UTF8String
+
+sealed trait ThetaSketchState {
+ def serialize(): Array[Byte]
+ def eval(): Array[Byte]
+}
+case class UpdatableSketchBuffer(sketch: UpdateSketch) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+ override def eval(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+}
+case class UnionAggregationBuffer(union: Union) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = union.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] = union.getResult.toByteArrayCompressed
+}
+case class IntersectionAggregationBuffer(intersection: Intersection) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+}
+case class FinalizedSketch(sketch: CompactSketch) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = sketch.toByteArrayCompressed
+ override def eval(): Array[Byte] = sketch.toByteArrayCompressed
+}
+
+/**
+ * The ThetaSketchAgg function utilizes a Datasketches ThetaSketch instance to
count a
+ * probabilistic approximation of the number of unique values in a given
column, and outputs the
+ * binary representation of the ThetaSketch.
+ *
+ * See [[https://datasketches.apache.org/docs/Theta/ThetaSketches.html]] for
more information.
+ *
+ * @param left
+ * child expression against which unique counting will occur
+ * @param right
+ * the log-base-2 of nomEntries decides the number of buckets for the sketch
+ * @param mutableAggBufferOffset
+ * offset for mutable aggregation buffer
+ * @param inputAggBufferOffset
+ * offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's compact binary
representation.
+ `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with
Nominal Entries deciding
+ the number buckets or slots for the ThetaSketch. """,
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(col, 12)) FROM VALUES (1), (1),
(2), (2), (3) tab(col);
+ 3
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaSketchAgg(
+ left: Expression,
+ right: Expression,
+ override val mutableAggBufferOffset: Int,
+ override val inputAggBufferOffset: Int)
+ extends TypedImperativeAggregate[ThetaSketchState]
+ with BinaryLike[Expression]
+ with ExpectsInputTypes {
+
+ // ThetaSketch config - mark as lazy so that they're not evaluated during
tree transformation.
+
+ lazy val lgNomEntries: Int = {
+ val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput)
+ lgNomEntriesInput
+ }
+
+ // Constructors
+
+ def this(child: Expression) = {
+ this(child, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS), 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Expression) = {
+ this(child, lgNomEntries, 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Int) = {
+ this(child, Literal(lgNomEntries), 0, 0)
+ }
+
+ // Copy constructors required by ImperativeAggregate
+
+ override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+ override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): ThetaSketchAgg =
+ copy(left = newLeft, right = newRight)
+
+ // Overrides for TypedImperativeAggregate
+
+ override def prettyName: String = "theta_sketch_agg"
+
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ TypeCollection(
+ IntegerType,
+ LongType,
+ FloatType,
+ DoubleType,
Review Comment:
super nit, please sort alphabetically?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/ThetaSketchUtilsSuite.scala:
##########
@@ -0,0 +1,64 @@
+/*
+ * 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.util
+
+import org.apache.spark.{SparkFunSuite, SparkRuntimeException}
+
+class ThetaSketchUtilsSuite extends SparkFunSuite {
+
+ test("checkLgNomLongs: accepts values within valid range") {
+ val validValues =
+ Seq(ThetaSketchUtils.MIN_LG_NOM_LONGS, 10, 20,
ThetaSketchUtils.MAX_LG_NOM_LONGS)
+ validValues.foreach { value =>
+ // Should not throw
+ ThetaSketchUtils.checkLgNomLongs(value)
+ }
+ }
+
+ test("checkLgNomLongs: throws exception for values below minimum") {
+ val invalidValues = Seq(ThetaSketchUtils.MIN_LG_NOM_LONGS - 1, 0, -5)
+ invalidValues.foreach { value =>
+ val e = intercept[SparkRuntimeException] {
+ ThetaSketchUtils.checkLgNomLongs(value)
+ }
+ assert(
+ e.getMessage.contains(
Review Comment:
Can you please switch these unit tests to use `checkError` instead of
asserting on substrings of the error message? Example: [1].
[1]
https://github.com/apache/spark/blob/596d03fc541cd9433fdfd511e131aa1cbc603a85/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala#L101-L107
##########
python/pyspark/sql/connect/functions/builtin.py:
##########
@@ -4229,6 +4229,90 @@ def hll_union(
hll_union.__doc__ = pysparkfuncs.hll_union.__doc__
+def theta_sketch_agg(
+ col: "ColumnOrName",
+ lgNomEntries: Optional[Union[int, Column]] = None,
+) -> Column:
+ if lgNomEntries is None:
+ return _invoke_function_over_columns("theta_sketch_agg", col)
+ else:
+ return _invoke_function_over_columns("theta_sketch_agg", col,
lit(lgNomEntries))
Review Comment:
super nit, this string is repeated from L4237 above, can we dedup it? Same
for other new code below, and also L4190 and L4203 above.
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -5560,6 +5560,18 @@
],
"sqlState" : "428EK"
},
+ "THETA_INVALID_INPUT_SKETCH_BUFFER" : {
+ "message" : [
+ "Invalid call to <function>; only valid Theta sketch buffers are
supported as inputs (such as those produced by the `theta_sketch_agg`
function)."
+ ],
+ "sqlState" : "22546"
+ },
+ "THETA_INVALID_LG_NOM_ENTRIES" : {
+ "message" : [
+ "Invalid Theta sketch call; the `lgNomEntries` value must be between
<min> and <max>, inclusive: <value>."
Review Comment:
Can we add "Invalid call to <function>" to this error message as well? It
could help make the error message easier to understand, like the previous one.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala:
##########
@@ -0,0 +1,682 @@
+/*
+ * 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.aggregate
+
+import org.apache.datasketches.common.SketchesArgumentException
+import org.apache.datasketches.memory.Memory
+import org.apache.datasketches.theta.{CompactSketch, Intersection,
SetOperation, Sketch, Union, UpdateSketch, UpdateSketchBuilder}
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{ExpectsInputTypes,
Expression, ExpressionDescription, Literal}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.TypedImperativeAggregate
+import org.apache.spark.sql.catalyst.trees.BinaryLike
+import org.apache.spark.sql.catalyst.util.{ArrayData, CollationFactory,
ThetaSketchUtils}
+import org.apache.spark.sql.errors.QueryExecutionErrors
+import org.apache.spark.sql.internal.types.StringTypeWithCollation
+import org.apache.spark.sql.types.{AbstractDataType, ArrayType, BinaryType,
DataType, DoubleType, FloatType, IntegerType, LongType, StringType,
TypeCollection}
+import org.apache.spark.unsafe.types.UTF8String
+
+sealed trait ThetaSketchState {
+ def serialize(): Array[Byte]
+ def eval(): Array[Byte]
+}
+case class UpdatableSketchBuffer(sketch: UpdateSketch) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+ override def eval(): Array[Byte] =
sketch.rebuild.compact.toByteArrayCompressed
+}
+case class UnionAggregationBuffer(union: Union) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = union.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] = union.getResult.toByteArrayCompressed
+}
+case class IntersectionAggregationBuffer(intersection: Intersection) extends
ThetaSketchState {
+ override def serialize(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+ override def eval(): Array[Byte] =
intersection.getResult.toByteArrayCompressed
+}
+case class FinalizedSketch(sketch: CompactSketch) extends ThetaSketchState {
+ override def serialize(): Array[Byte] = sketch.toByteArrayCompressed
+ override def eval(): Array[Byte] = sketch.toByteArrayCompressed
+}
+
+/**
+ * The ThetaSketchAgg function utilizes a Datasketches ThetaSketch instance to
count a
+ * probabilistic approximation of the number of unique values in a given
column, and outputs the
+ * binary representation of the ThetaSketch.
+ *
+ * See [[https://datasketches.apache.org/docs/Theta/ThetaSketches.html]] for
more information.
+ *
+ * @param left
+ * child expression against which unique counting will occur
+ * @param right
+ * the log-base-2 of nomEntries decides the number of buckets for the sketch
+ * @param mutableAggBufferOffset
+ * offset for mutable aggregation buffer
+ * @param inputAggBufferOffset
+ * offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's compact binary
representation.
+ `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with
Nominal Entries deciding
+ the number buckets or slots for the ThetaSketch. """,
+ examples = """
+ Examples:
+ > SELECT theta_sketch_estimate(_FUNC_(col, 12)) FROM VALUES (1), (1),
(2), (2), (3) tab(col);
+ 3
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ThetaSketchAgg(
+ left: Expression,
+ right: Expression,
+ override val mutableAggBufferOffset: Int,
+ override val inputAggBufferOffset: Int)
+ extends TypedImperativeAggregate[ThetaSketchState]
+ with BinaryLike[Expression]
+ with ExpectsInputTypes {
+
+ // ThetaSketch config - mark as lazy so that they're not evaluated during
tree transformation.
+
+ lazy val lgNomEntries: Int = {
+ val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+ ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput)
+ lgNomEntriesInput
+ }
+
+ // Constructors
+
+ def this(child: Expression) = {
+ this(child, Literal(ThetaSketchUtils.DEFAULT_LG_NOM_LONGS), 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Expression) = {
+ this(child, lgNomEntries, 0, 0)
+ }
+
+ def this(child: Expression, lgNomEntries: Int) = {
+ this(child, Literal(lgNomEntries), 0, 0)
+ }
+
+ // Copy constructors required by ImperativeAggregate
+
+ override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+ override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int):
ThetaSketchAgg =
+ copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): ThetaSketchAgg =
+ copy(left = newLeft, right = newRight)
+
+ // Overrides for TypedImperativeAggregate
+
+ override def prettyName: String = "theta_sketch_agg"
+
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ TypeCollection(
+ IntegerType,
+ LongType,
+ FloatType,
+ DoubleType,
+ StringTypeWithCollation(supportsTrimCollation = true),
+ BinaryType,
+ ArrayType(IntegerType),
+ ArrayType(LongType)),
+ IntegerType)
+
+ override def dataType: DataType = BinaryType
+
+ override def nullable: Boolean = false
+
+ /**
+ * Instantiate an UpdateSketch instance using the lgNomEntries param.
+ *
+ * @return
+ * an UpdateSketch instance wrapped with UpdatableSketchBuffer
+ */
+ override def createAggregationBuffer(): ThetaSketchState = {
+ val builder = new UpdateSketchBuilder
+ builder.setLogNominalEntries(lgNomEntries)
+ UpdatableSketchBuffer(builder.build)
+ }
+
+ /**
+ * Evaluate the input row and update the UpdateSketch instance with the
row's value. The update
+ * function only supports a subset of Spark SQL types, and an exception will
be thrown for
+ * unsupported types.
+ *
+ * @param updateBuffer
+ * A previously initialized UpdateSketch instance
+ * @param input
+ * An input row
+ */
+ override def update(updateBuffer: ThetaSketchState, input: InternalRow):
ThetaSketchState =
+ updateBuffer match {
+ case UpdatableSketchBuffer(sketch) =>
+ val v = left.eval(input)
+ v match {
+ case null => UpdatableSketchBuffer(sketch)
+ case _ =>
+ left.dataType match {
Review Comment:
There are still many levels of indentation here. Perhaps we could improve it
with a helper `var`s, something like
```
// <some descriptive comment here>
var makeUpdatableSketchBuffer = false
// <some descriptive comment here>
var inputDataType = Option.empty[DataType]
updateBuffer match {
case UpdatableSketchBuffer(sketch) =>
val v = left.eval(input)
v match {
case null =>
makeUpdatableSketchBuffer = true
case _ =>
makeUpdatableSketchBuffer = true
inputDataType = Some(left.dataType)
case _ =>
}
inputDataType.foreach {
_ match {
case IntegerType =>
sketch.update(v.asInstanceOf[Int].toLong) // Promote to long
...
}
}
if (makeUpdatableSketchBuffer) {
UpdatableSketchBuffer(sketch)
} else {
updateBuffer
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
