cboumalh commented on code in PR #51298: URL: https://github.com/apache/spark/pull/51298#discussion_r2302265273
########## 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) { Review Comment: yep makes sense, will fix those -- 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
