[SPARK-42999][Connect] Dataset#foreach, foreachPartition #40628

zhenlineo · 2023-04-01T00:08:32Z

What changes were proposed in this pull request?

Implements missing methods in Dataset: foreach and foreachPartition.
PR based on top of #40581
The impl of foreachPartition is based on top of mapPartitions + count.

Why are the changes needed?

Add missing methods in Dataset.

Does this PR introduce any user-facing change?

No

How was this patch tested?

E2E tests.

HyukjinKwon · 2023-04-03T00:23:14Z

cc @xinrong-meng

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/UdfUtils.scala

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala

hvanhovell

Two small comments. Otherwise GTG.

hvanhovell

LGTM

@zhenlineo can you update this branch?

xinrong-meng · 2023-04-06T23:08:06Z

Thanks for working on that! We will enable foreach and foreachPartition in Python Client's DataFrame based on that.

hvanhovell · 2023-04-07T01:28:19Z

Merging, thanks!

LuciferYang · 2023-04-10T03:06:36Z

...connect/client/jvm/src/test/scala/org/apache/spark/sql/UserDefinedFunctionE2ETestSuite.scala

@@ -128,4 +130,72 @@ class UserDefinedFunctionE2ETestSuite extends RemoteSparkSession {
      .collect()
    assert(result.sorted.toSeq === Seq(23, 25, 25, 27))
  }
+
+  test("Dataset foreach") {


How to test the new case with maven ? @zhenlineo

I run the following commands:

build/mvn clean install -pl connector/connect/server -am -DskipTests build/mvn clean install -pl assembly -am -DskipTests build/mvn clean install -pl connector/connect/client

there are 15 test failed in this one:

- Dataset typed filter *** FAILED *** io.grpc.StatusRuntimeException: INTERNAL: org.apache.spark.sql.UserDefinedFunctionE2ETestSuite at io.grpc.Status.asRuntimeException(Status.java:535) at io.grpc.stub.ClientCalls$BlockingResponseStream.hasNext(ClientCalls.java:660) at org.apache.spark.sql.connect.client.SparkResult.org$apache$spark$sql$connect$client$SparkResult$$processResponses(SparkResult.scala:61) at org.apache.spark.sql.connect.client.SparkResult.length(SparkResult.scala:106) at org.apache.spark.sql.connect.client.SparkResult.toArray(SparkResult.scala:123) at org.apache.spark.sql.Dataset.$anonfun$collect$1(Dataset.scala:2687) at org.apache.spark.sql.Dataset.withResult(Dataset.scala:3088) at org.apache.spark.sql.Dataset.collect(Dataset.scala:2686) at org.apache.spark.sql.Dataset.collectAsList(Dataset.scala:2700) at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.$anonfun$new$1(UserDefinedFunctionE2ETestSuite.scala:38) ... *** 15 TESTS FAILED ***

Fixing now. Will "at" you on the followup PR.

FYI, it seems flaky:

[info] - Dataset foreach *** FAILED *** (133 milliseconds) [info] "INTERNAL: RST_STREAM closed stream. HTTP/2 error code: PROTOCOL_ERROR" did not contain "Hello foreach" (UserDefinedFunctionE2ETestSuite.scala:141) [info] org.scalatest.exceptions.TestFailedException: [info] at org.scalatest.Assertions.newAssertionFailedException(Assertions.scala:472) [info] at org.scalatest.Assertions.newAssertionFailedException$(Assertions.scala:471) [info] at org.scalatest.Assertions$.newAssertionFailedException(Assertions.scala:1231) [info] at org.scalatest.Assertions$AssertionsHelper.macroAssert(Assertions.scala:1295) [info] at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.$anonfun$new$19(UserDefinedFunctionE2ETestSuite.scala:141) [info] at org.scalatest.OutcomeOf.outcomeOf(OutcomeOf.scala:85) [info] at org.scalatest.OutcomeOf.outcomeOf$(OutcomeOf.scala:83) [info] at org.scalatest.OutcomeOf$.outcomeOf(OutcomeOf.scala:104) [info] at org.scalatest.Transformer.apply(Transformer.scala:22) [info] at org.scalatest.Transformer.apply(Transformer.scala:20) [info] at org.scalatest.funsuite.AnyFunSuiteLike$$anon$1.apply(AnyFunSuiteLike.scala:226) [info] at org.scalatest.TestSuite.withFixture(TestSuite.scala:196) [info] at org.scalatest.TestSuite.withFixture$(TestSuite.scala:195) [info] at org.scalatest.funsuite.AnyFunSuite.withFixture(AnyFunSuite.scala:1564) [info] at org.scalatest.funsuite.AnyFunSuiteLike.invokeWithFixture$1(AnyFunSuiteLike.scala:224) [info] at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$runTest$1(AnyFunSuiteLike.scala:236) [info] at org.scalatest.SuperEngine.runTestImpl(Engine.scala:306) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTest(AnyFunSuiteLike.scala:236) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTest$(AnyFunSuiteLike.scala:218) [info] at org.scalatest.funsuite.AnyFunSuite.runTest(AnyFunSuite.scala:1564) [info] at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$runTests$1(AnyFunSuiteLike.scala:269) [info] at org.scalatest.SuperEngine.$anonfun$runTestsInBranch$1(Engine.scala:413) [info] at scala.collection.immutable.List.foreach(List.scala:431) [info] at org.scalatest.SuperEngine.traverseSubNodes$1(Engine.scala:401) [info] at org.scalatest.SuperEngine.runTestsInBranch(Engine.scala:396) [info] at org.scalatest.SuperEngine.runTestsImpl(Engine.scala:475) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTests(AnyFunSuiteLike.scala:269) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTests$(AnyFunSuiteLike.scala:268) [info] at org.scalatest.funsuite.AnyFunSuite.runTests(AnyFunSuite.scala:1564) [info] at org.scalatest.Suite.run(Suite.scala:1114) [info] at org.scalatest.Suite.run$(Suite.scala:1096) [info] at org.scalatest.funsuite.AnyFunSuite.org$scalatest$funsuite$AnyFunSuiteLike$$super$run(AnyFunSuite.scala:1564) [info] at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$run$1(AnyFunSuiteLike.scala:273) [info] at org.scalatest.SuperEngine.runImpl(Engine.scala:535) [info] at org.scalatest.funsuite.AnyFunSuiteLike.run(AnyFunSuiteLike.scala:273) [info] at org.scalatest.funsuite.AnyFunSuiteLike.run$(AnyFunSuiteLike.scala:272) [info] at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.org$scalatest$BeforeAndAfterAll$$super$run(UserDefinedFunctionE2ETestSuite.scala:36) [info] at org.scalatest.BeforeAndAfterAll.liftedTree1$1(BeforeAndAfterAll.scala:213) [info] at org.scalatest.BeforeAndAfterAll.run(BeforeAndAfterAll.scala:210) [info] at org.scalatest.BeforeAndAfterAll.run$(BeforeAndAfterAll.scala:208) [info] at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.run(UserDefinedFunctionE2ETestSuite.scala:36) [info] at org.scalatest.tools.Framework.org$scalatest$tools$Framework$$runSuite(Framework.scala:321) [info] at org.scalatest.tools.Framework$ScalaTestTask.execute(Framework.scala:517) [info] at sbt.ForkMain$Run.lambda$runTest$1(ForkMain.java:413) [info] at java.util.concurrent.FutureTask.run(FutureTask.java:266) [info] at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149) [info] at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624) [info] at java.lang.Thread.run(Thread.java:750) [info] - Dataset foreach - java *** FAILED *** (123 milliseconds) [info] "INTERNAL: RST_STREAM closed stream. HTTP/2 error code: PROTOCOL_ERROR" did not contain "Hello foreach" (UserDefinedFunctionE2ETestSuite.scala:154) [info] org.scalatest.exceptions.TestFailedException: [info] at org.scalatest.Assertions.newAssertionFailedException(Assertions.scala:472) [info] at org.scalatest.Assertions.newAssertionFailedException$(Assertions.scala:471) [info] at org.scalatest.Assertions$.newAssertionFailedException(Assertions.scala:1231)

https://github.com/apache/spark/actions/runs/4748841722/jobs/8435481918

This actually seems pretty flaky.

https://github.com/LuciferYang/spark/actions/runs/4749019479/jobs/8435837245
https://github.com/LuciferYang/spark/actions/runs/4749127823/jobs/8436326385

same issue

[info] - Dataset foreachPartition - java *** FAILED *** (92 milliseconds) [info] "INTERNAL: RST_STREAM closed stream. HTTP/2 error code: PROTOCOL_ERROR" did not contain "45 did not equal -1" (UserDefinedFunctionE2ETestSuite.scala:190) [info] org.scalatest.exceptions.TestFailedException: [info] at org.scalatest.Assertions.newAssertionFailedException(Assertions.scala:472) [info] at org.scalatest.Assertions.newAssertionFailedException$(Assertions.scala:471) [info] at org.scalatest.Assertions$.newAssertionFailedException(Assertions.scala:1231) [info] at org.scalatest.Assertions$AssertionsHelper.macroAssert(Assertions.scala:1295) [info] at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.$anonfun$new$29(UserDefinedFunctionE2ETestSuite.scala:190) [info] at org.scalatest.OutcomeOf.outcomeOf(OutcomeOf.scala:85) [info] at org.scalatest.OutcomeOf.outcomeOf$(OutcomeOf.scala:83) [info] at org.scalatest.OutcomeOf$.outcomeOf(OutcomeOf.scala:104) [info] at org.scalatest.Transformer.apply(Transformer.scala:22) [info] at org.scalatest.Transformer.apply(Transformer.scala:20) [info] at org.scalatest.funsuite.AnyFunSuiteLike$$anon$1.apply(AnyFunSuiteLike.scala:226) [info] at org.scalatest.TestSuite.withFixture(TestSuite.scala:196) [info] at org.scalatest.TestSuite.withFixture$(TestSuite.scala:195) [info] at org.scalatest.funsuite.AnyFunSuite.withFixture(AnyFunSuite.scala:1564) [info] at org.scalatest.funsuite.AnyFunSuiteLike.invokeWithFixture$1(AnyFunSuiteLike.scala:224) [info] at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$runTest$1(AnyFunSuiteLike.scala:236) [info] at org.scalatest.SuperEngine.runTestImpl(Engine.scala:306) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTest(AnyFunSuiteLike.scala:236) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTest$(AnyFunSuiteLike.scala:218) [info] at org.scalatest.funsuite.AnyFunSuite.runTest(AnyFunSuite.scala:1564) [info] at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$runTests$1(AnyFunSuiteLike.scala:269) [info] at org.scalatest.SuperEngine.$anonfun$runTestsInBranch$1(Engine.scala:413) [info] at scala.collection.immutable.List.foreach(List.scala:431) [info] at org.scalatest.SuperEngine.traverseSubNodes$1(Engine.scala:401) [info] at org.scalatest.SuperEngine.runTestsInBranch(Engine.scala:396) [info] at org.scalatest.SuperEngine.runTestsImpl(Engine.scala:475) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTests(AnyFunSuiteLike.scala:269) [info] at org.scalatest.funsuite.AnyFunSuiteLike.runTests$(AnyFunSuiteLike.scala:268) [info] at org.scalatest.funsuite.AnyFunSuite.runTests(AnyFunSuite.scala:1564) [info] at org.scalatest.Suite.run(Suite.scala:1114) [info] at org.scalatest.Suite.run$(Suite.scala:1096) [info] at org.scalatest.funsuite.AnyFunSuite.org$scalatest$funsuite$AnyFunSuiteLike$$super$run(AnyFunSuite.scala:1564) [info] at org.scalatest.funsuite.AnyFunSuiteLike.$anonfun$run$1(AnyFunSuiteLike.scala:273) [info] at org.scalatest.SuperEngine.runImpl(Engine.scala:535) [info] at org.scalatest.funsuite.AnyFunSuiteLike.run(AnyFunSuiteLike.scala:273) [info] at org.scalatest.funsuite.AnyFunSuiteLike.run$(AnyFunSuiteLike.scala:272) [info] at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.org$scalatest$BeforeAndAfterAll$$super$run(UserDefinedFunctionE2ETestSuite.scala:36) [info] at org.scalatest.BeforeAndAfterAll.liftedTree1$1(BeforeAndAfterAll.scala:213) [info] at org.scalatest.BeforeAndAfterAll.run(BeforeAndAfterAll.scala:210) [info] at org.scalatest.BeforeAndAfterAll.run$(BeforeAndAfterAll.scala:208) [info] at org.apache.spark.sql.UserDefinedFunctionE2ETestSuite.run(UserDefinedFunctionE2ETestSuite.scala:36) [info] at org.scalatest.tools.Framework.org$scalatest$tools$Framework$$runSuite(Framework.scala:321) [info] at org.scalatest.tools.Framework$ScalaTestTask.execute(Framework.scala:517) [info] at sbt.ForkMain$Run.lambda$runTest$1(ForkMain.java:413) [info] at java.util.concurrent.FutureTask.run(FutureTask.java:266) [info] at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149) [info] at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624) [info] at java.lang.Thread.run(Thread.java:750)

With the latest master, I've been running the test locally for 100 times, and cannot repro this error. Do we have some way to see how flaky is this test? We can certainly mute these tests if it continue to be flaky. My local run indicate it is not flaky.

@zhenlineo Sorry for forgetting to notify you, after revert 09a4353 everything is ok now

github-actions bot added CONNECT SQL labels Apr 1, 2023

zhenlineo force-pushed the foreach-typed branch from b9f0f93 to 18c41c3 Compare April 1, 2023 00:18

hvanhovell reviewed Apr 3, 2023

View reviewed changes

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/UdfUtils.scala Outdated Show resolved Hide resolved

hvanhovell reviewed Apr 3, 2023

View reviewed changes

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala Outdated Show resolved Hide resolved

hvanhovell requested changes Apr 3, 2023

View reviewed changes

zhenlineo force-pushed the foreach-typed branch 2 times, most recently from 3d6bdfa to 78c5684 Compare April 5, 2023 18:40

hvanhovell approved these changes Apr 6, 2023

View reviewed changes

zhenlineo added 3 commits April 6, 2023 13:37

Foreach and foreachPartition

df40baa

Address comments

9d90ddb

Add col(*) for map too

cc0e9f4

zhenlineo force-pushed the foreach-typed branch from 78c5684 to cc0e9f4 Compare April 6, 2023 20:38

xinrong-meng approved these changes Apr 6, 2023

View reviewed changes

hvanhovell closed this in 4d726a9 Apr 7, 2023

LuciferYang reviewed Apr 10, 2023

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-42999][Connect] Dataset#foreach, foreachPartition #40628

[SPARK-42999][Connect] Dataset#foreach, foreachPartition #40628

zhenlineo commented Apr 1, 2023

HyukjinKwon commented Apr 3, 2023

hvanhovell left a comment

hvanhovell left a comment

xinrong-meng commented Apr 6, 2023

hvanhovell commented Apr 7, 2023

LuciferYang Apr 10, 2023

zhenlineo Apr 10, 2023

HyukjinKwon Apr 20, 2023

HyukjinKwon Apr 20, 2023

LuciferYang Apr 20, 2023

zhenlineo Apr 21, 2023

LuciferYang Apr 22, 2023

[SPARK-42999][Connect] Dataset#foreach, foreachPartition #40628

[SPARK-42999][Connect] Dataset#foreach, foreachPartition #40628

Conversation

zhenlineo commented Apr 1, 2023

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

HyukjinKwon commented Apr 3, 2023

hvanhovell left a comment

Choose a reason for hiding this comment

hvanhovell left a comment

Choose a reason for hiding this comment

xinrong-meng commented Apr 6, 2023

hvanhovell commented Apr 7, 2023

LuciferYang Apr 10, 2023

Choose a reason for hiding this comment

zhenlineo Apr 10, 2023

Choose a reason for hiding this comment

HyukjinKwon Apr 20, 2023

Choose a reason for hiding this comment

HyukjinKwon Apr 20, 2023

Choose a reason for hiding this comment

LuciferYang Apr 20, 2023

Choose a reason for hiding this comment

zhenlineo Apr 21, 2023

Choose a reason for hiding this comment

LuciferYang Apr 22, 2023

Choose a reason for hiding this comment