You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2024/12/16 11:36:57 UTC
(spark) branch master updated: [SPARK-50581][SQL] fix support for UDAF in Dataset.observe()
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 44ab3492455e [SPARK-50581][SQL] fix support for UDAF in Dataset.observe()
44ab3492455e is described below
commit 44ab3492455e71ebc46d545ce2323be41c03633f
Author: tomsisso <to...@definity.ai>
AuthorDate: Mon Dec 16 20:36:45 2024 +0900
[SPARK-50581][SQL] fix support for UDAF in Dataset.observe()
### What changes were proposed in this pull request?
Mark inputProjection field as transient in org.apache.spark.sql.execution.aggregate.ScalaAggregator.
### Why are the changes needed?
To support UDAF in Dataset.observe() which currently fails due to serialization exception.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Added test.
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #49190 from tomsisso/support-udaf-in-observe.
Authored-by: tomsisso <to...@definity.ai>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
.../spark/sql/execution/aggregate/udaf.scala | 16 ++++----
.../spark/sql/util/DataFrameCallbackSuite.scala | 48 +++++++++++++++++++++-
2 files changed, 55 insertions(+), 9 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
index 09d9915022a6..1197a16a35e9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
@@ -501,17 +501,17 @@ case class ScalaAggregator[IN, BUF, OUT](
with Logging {
// input and buffer encoders are resolved by ResolveEncodersInScalaAgg
- private[this] lazy val inputDeserializer = inputEncoder.createDeserializer()
- private[this] lazy val bufferSerializer = bufferEncoder.createSerializer()
- private[this] lazy val bufferDeserializer = bufferEncoder.createDeserializer()
- private[this] lazy val outputEncoder = encoderFor(agg.outputEncoder)
- private[this] lazy val outputSerializer = outputEncoder.createSerializer()
+ @transient private[this] lazy val inputDeserializer = inputEncoder.createDeserializer()
+ @transient private[this] lazy val bufferSerializer = bufferEncoder.createSerializer()
+ @transient private[this] lazy val bufferDeserializer = bufferEncoder.createDeserializer()
+ @transient private[this] lazy val outputEncoder = encoderFor(agg.outputEncoder)
+ @transient private[this] lazy val outputSerializer = outputEncoder.createSerializer()
def dataType: DataType = outputEncoder.objSerializer.dataType
def inputTypes: Seq[DataType] = inputEncoder.schema.map(_.dataType)
- override lazy val deterministic: Boolean = isDeterministic
+ @transient override lazy val deterministic: Boolean = isDeterministic
def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ScalaAggregator[IN, BUF, OUT] =
copy(mutableAggBufferOffset = newMutableAggBufferOffset)
@@ -519,7 +519,7 @@ case class ScalaAggregator[IN, BUF, OUT](
def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ScalaAggregator[IN, BUF, OUT] =
copy(inputAggBufferOffset = newInputAggBufferOffset)
- private[this] lazy val inputProjection = UnsafeProjection.create(children)
+ @transient private[this] lazy val inputProjection = UnsafeProjection.create(children)
def createAggregationBuffer(): BUF = agg.zero
@@ -533,7 +533,7 @@ case class ScalaAggregator[IN, BUF, OUT](
if (outputEncoder.isSerializedAsStruct) row else row.get(0, dataType)
}
- private[this] lazy val bufferRow = new UnsafeRow(bufferEncoder.namedExpressions.length)
+ @transient private[this] lazy val bufferRow = new UnsafeRow(bufferEncoder.namedExpressions.length)
def serialize(agg: BUF): Array[Byte] =
bufferSerializer(agg).asInstanceOf[UnsafeRow].getBytes()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
index be91f5e789e2..7e6f10bcc46f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/util/DataFrameCallbackSuite.scala
@@ -22,7 +22,7 @@ import java.lang.{Long => JLong}
import scala.collection.mutable.ArrayBuffer
import org.apache.spark._
-import org.apache.spark.sql.{functions, Dataset, QueryTest, Row, SparkSession}
+import org.apache.spark.sql.{functions, Dataset, Encoder, Encoders, QueryTest, Row, SparkSession}
import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project}
import org.apache.spark.sql.execution.{QueryExecution, WholeStageCodegenExec}
@@ -30,6 +30,7 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.command.{CreateDataSourceTableAsSelectCommand, LeafRunnableCommand}
import org.apache.spark.sql.execution.datasources.InsertIntoHadoopFsRelationCommand
import org.apache.spark.sql.execution.datasources.json.JsonFileFormat
+import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.StringType
@@ -339,6 +340,51 @@ class DataFrameCallbackSuite extends QueryTest
}
}
+ test("SPARK-50581: support observe with udaf") {
+ withUserDefinedFunction(("someUdaf", true)) {
+ spark.udf.register("someUdaf", functions.udaf(new Aggregator[JLong, JLong, JLong] {
+ def zero: JLong = 0L
+ def reduce(b: JLong, a: JLong): JLong = a + b
+ def merge(b1: JLong, b2: JLong): JLong = b1 + b2
+ def finish(r: JLong): JLong = r
+ def bufferEncoder: Encoder[JLong] = Encoders.LONG
+ def outputEncoder: Encoder[JLong] = Encoders.LONG
+ }))
+
+ val df = spark.range(100)
+
+ val metricMaps = ArrayBuffer.empty[Map[String, Row]]
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = {
+ if (qe.observedMetrics.nonEmpty) {
+ metricMaps += qe.observedMetrics
+ }
+ }
+
+ override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = {
+ // No-op
+ }
+ }
+ try {
+ spark.listenerManager.register(listener)
+
+ // udaf usage in observe is not working (serialization exception)
+ df.observe(
+ name = "my_metrics",
+ expr("someUdaf(id)").as("agg")
+ )
+ .collect()
+
+ sparkContext.listenerBus.waitUntilEmpty()
+ assert(metricMaps.size === 1)
+ assert(metricMaps.head("my_metrics") === Row(4950L))
+
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+ }
+ }
+
private def validateObservedMetrics(df: Dataset[JLong]): Unit = {
val metricMaps = ArrayBuffer.empty[Map[String, Row]]
val listener = new QueryExecutionListener {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org