hvanhovell commented on code in PR #46245: URL: https://github.com/apache/spark/pull/46245#discussion_r1594544696
########## connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala: ########## @@ -0,0 +1,104 @@ +/* + * 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.expressions + +import org.apache.spark.sql.{Encoder, TypedColumn} + +/** + * A base class for user-defined aggregations, which can be used in `Dataset` operations to take + * all of the elements of a group and reduce them to a single value. + * + * For example, the following aggregator extracts an `int` from a specific class and adds them up: + * {{{ + * case class Data(i: Int) + * + * val customSummer = new Aggregator[Data, Int, Int] { + * def zero: Int = 0 + * def reduce(b: Int, a: Data): Int = b + a.i + * def merge(b1: Int, b2: Int): Int = b1 + b2 + * def finish(r: Int): Int = r + * def bufferEncoder: Encoder[Int] = Encoders.scalaInt + * def outputEncoder: Encoder[Int] = Encoders.scalaInt + * } + * + * spark.udf.register("customSummer", udaf(customSummer)) + * val ds: Dataset[Data] = ... + * val aggregated = ds.selectExpr("customSummer(i)") + * }}} + * + * Based loosely on Aggregator from Algebird: https://github.com/twitter/algebird + * + * @tparam IN The input type for the aggregation. + * @tparam BUF The type of the intermediate value of the reduction. + * @tparam OUT The type of the final output result. + * @since 4.0.0 + */ +@SerialVersionUID(2093413866369130093L) +abstract class Aggregator[-IN, BUF, OUT] extends Serializable { + + /** + * A zero value for this aggregation. Should satisfy the property that any b + zero = b. + * + * @since 4.0.0 + */ + def zero: BUF + + /** + * Combine two values to produce a new value. For performance, the function may modify `b` and + * return it instead of constructing new object for b. + * + * @since 4.0.0 + */ + def reduce(b: BUF, a: IN): BUF + + /** + * Merge two intermediate values. + * + * @since 4.0.0 + */ + def merge(b1: BUF, b2: BUF): BUF + + /** + * Transform the output of the reduction. + * + * @since 4.0.0 + */ + def finish(reduction: BUF): OUT + + /** + * Specifies the `Encoder` for the intermediate value type. + * + * @since 4.0.0 + */ + def bufferEncoder: Encoder[BUF] + + /** + * Specifies the `Encoder` for the final output value type. + * + * @since 4.0.0 + */ + def outputEncoder: Encoder[OUT] + + /** + * Returns this `Aggregator` as a `TypedColumn` that can be used in `Dataset`. + * operations. + */ + def toColumn: TypedColumn[IN, OUT] = { Review Comment: I think it should take a wildcard, we have done similar things before. -- 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
