Re: [PR] [SPARK-53848][SQL] Add ability to support Alpha family in Theta Aggregates [spark]

["cboumalh (via GitHub)" <gi...@apache.org>]

cboumalh commented on code in PR #55300:
URL: https://github.com/apache/spark/pull/55300#discussion_r3074103671


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala:
##########
@@ -71,49 +73,73 @@ case class FinalizedSketch(sketch: CompactSketch) extends ThetaSketchState {
 // scalastyle:off line.size.limit
 @ExpressionDescription(
   usage = """
-    _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch compact binary representation.
+    _FUNC_(expr, lgNomEntries, family) - Returns the ThetaSketch compact binary representation.
       `lgNomEntries` (optional) is the log-base-2 of nominal entries, with nominal entries deciding
-      the number buckets or slots for the ThetaSketch. """,
+      the number buckets or slots for the ThetaSketch.
+      `family` (optional) is the sketch family, either 'QUICKSELECT' or 'ALPHA' (defaults to
+      'QUICKSELECT').""",
   examples = """
     Examples:
+      > SELECT theta_sketch_estimate(_FUNC_(col)) FROM VALUES (1), (1), (2), (2), (3) tab(col);
+       3
       > SELECT theta_sketch_estimate(_FUNC_(col, 12)) FROM VALUES (1), (1), (2), (2), (3) tab(col);
        3
+      > SELECT theta_sketch_estimate(_FUNC_(col, 15, 'ALPHA')) 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)
+       first: Expression,
+       second: Expression,
+       third: Expression,
+       override val mutableAggBufferOffset: Int,
+       override val inputAggBufferOffset: Int)
     extends TypedImperativeAggregate[ThetaSketchState]
-    with BinaryLike[Expression]
+    with TernaryLike[Expression]
     with ExpectsInputTypes {
 
   // ThetaSketch config - mark as lazy so that they're not evaluated during tree transformation.
 
-  lazy val lgNomEntries: Int = {
-    if (!right.foldable) {
+
+  private lazy val lgNomEntries: Int = {
+    if (!second.foldable) {
       throw QueryExecutionErrors.thetaLgNomEntriesMustBeConstantError(prettyName)
     }
-    val lgNomEntriesInput = right.eval().asInstanceOf[Int]
+    val lgNomEntriesInput = second.eval().asInstanceOf[Int]
     ThetaSketchUtils.checkLgNomLongs(lgNomEntriesInput, prettyName)
     lgNomEntriesInput
   }
 
-  // Constructors
+  private lazy val family: Family = {
+    if (!third.foldable) {
+      throw QueryExecutionErrors.thetaFamilyMustBeConstantError(prettyName)

Review Comment:
   Ideally we can do this check at the analysis step in checkInputDataTypes to fail early, instead of at runtime. Same applies to the lgNomEntries. I can create a follow up, or we can make the modifications here other reviewers see it fit.



-- 
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-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org