HyukjinKwon commented on code in PR #59474:
URL: https://github.com/apache/spark/pull/59474#discussion_r4278133992
##########
sql/api/src/main/scala/org/apache/spark/sql/catalyst/trees/origin.scala:
##########
@@ -152,8 +152,13 @@ object CurrentOrigin {
"(org\\.apache\\.spark\\.sql\\." +
"(?:(classic|connect)\\.)?" +
"(?:functions|Column|ColumnName|SQLImplicits|Dataset|DataFrameStatFunctions|DatasetHolder"
+
- "|SparkSession|ColumnNodeToProtoConverter)" +
+ "|SparkSession|ColumnNodeToProtoConverter" +
+ // Typed APIs build their UDF column inside these helpers (e.g. the
Connect client encodes
Review Comment:
**Nit (P3):** This comment sits before the
`KeyValueGroupedDataset|KeyValueGroupedDatasetImpl` alternatives but cites
`groupByKey/filter/reduce`. In the Connect client, `filter` and `reduce` call
`udf.apply(col("*"))` directly inside `Dataset` (already matched) and never go
through these classes. Only `groupByKey` builds its grouping UDF in
`KeyValueGroupedDatasetImpl`. What fixes `filter`/`reduce` and ordinary UDF
calls is the new `org.apache.spark.sql.expressions.UserDefinedFunction`
alternative below, which has no comment. A future reader could easily take that
entry for redundant. Could you split the explanation so each alternative is
justified where it appears? UserDefinedFunction would cover every UDF column,
since all are built via `UserDefinedFunction.apply`.
KeyValueGroupedDataset/Impl would cover the Connect groupByKey factory and
columns built in KVGD methods.
**Verification:**
- **Inspection:** The reworded comments match the actual call paths (connect
Dataset filter/reduce call udf.apply directly; groupByKey goes through
KeyValueGroupedDatasetImpl), contain only ASCII, and keep lines within 100
characters.
- **Compile:** The sql/api module still compiles, and the concatenated
pattern string is byte-identical to before.
##########
sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -51,6 +53,33 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest
with RemoteSparkSessi
assert(values == Arrays.asList[Int](5, 5))
}
+ test("SPARK-60247: typed groupByKey origin points to the user call site") {
+ // The Connect client encodes `groupByKey`'s grouping key as a Scala UDF
applied to
+ // `col("*")`, built inside `KeyValueGroupedDatasetImpl`. The DataFrame
query context
+ // captured for that column must skip the `UserDefinedFunction.apply` and
+ // `KeyValueGroupedDatasetImpl` frames so the call site lands on the user
code that called
+ // `groupByKey` (this suite), not on those internal helpers. Assert on the
serialized origin
+ // of the grouping expression to guard the frame filter in `CurrentOrigin`.
+ val ds = spark
+ .range(10)
+ .groupByKey(v => v % 2)
+ .count()
+
+ val groupingUdf = ds.plan.getRoot.getAggregate.getGroupingExpressions(0)
Review Comment:
**Nit (P3):** This test only checks `getGroupingExpressions(0)`, the
grouping UDF built in `KeyValueGroupedDatasetImpl$.apply`. Its call chain
(`Column` -> `UserDefinedFunction.apply` -> `KeyValueGroupedDatasetImpl$.apply`
-> `Dataset.groupByKey` -> test) never passes through a plain
`KeyValueGroupedDataset` frame. So removing `|KeyValueGroupedDataset` from
`sparkCodePattern` (keeping `|KeyValueGroupedDatasetImpl`) would not fail this
test or the `DatasetSuite` one. That token is what moves the call site for
columns built inside `KeyValueGroupedDataset` methods. For example, `.count()`
delegates to `sql.KeyValueGroupedDataset.count`, which builds
`agg(cnt(lit(1)).as(...))`. Could you also assert on the origin of the
aggregate expression in this same plan
(`getAggregate.getAggregateExpressions(0)`) so that behavior is guarded too?
**Verification:**
- **Regression:** An assertion on the origin of a column built inside a
KeyValueGroupedDataset method (e.g. the count aggregate) fails when the
KeyValueGroupedDataset alternative is dropped, and passes with it present.
- **Behavior:** The existing grouping-UDF call-site assertion still passes.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
