dtenedor commented on code in PR #51298: URL: https://github.com/apache/spark/pull/51298#discussion_r2302261956
########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ThetaSketchUtils.scala: ########## @@ -0,0 +1,43 @@ +/* + * 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.sql.errors.QueryExecutionErrors + +object ThetaSketchUtils { + // Bounds copied from DataSketches' ThetaUtil Review Comment: please expand to mention what their actual meanings are in this codebase? ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala: ########## @@ -0,0 +1,652 @@ +/* + * 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 + */ +// 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, + mutableAggBufferOffset: Int = 0, + inputAggBufferOffset: Int = 0) + 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) + if (v != null) { + 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 + */ +// 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, + mutableAggBufferOffset: Int = 0, Review Comment: Let's remove the default values for these parameters to protect ourselves against someone forgetting to assign them properly in the future. -- 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
