[SPARK-42669][CONNECT] Short circuit local relation RPCs #40782

beliefer · 2023-04-13T13:37:38Z

What changes were proposed in this pull request?

Operations on LocalRelation can mostly be done locally (without sending RPCs).
We should leverage this.

Why are the changes needed?

Avoid sending RPCs for LocalRelation.

Does this PR introduce any user-facing change?

'No'.
New feature.

How was this patch tested?

Exists test cases.

beliefer · 2023-04-14T01:18:18Z

The failure GA is unrelated to this PR.
ping @hvanhovell @zhengruifeng cc @dongjoon-hyun @amaliujia

zhengruifeng · 2023-04-14T01:26:06Z

sorry, test_parity_torch_distributor became flaky, I am going to disable related cases for now

amaliujia · 2023-04-14T03:31:43Z

So this happens when

val df = createDataFrame(...)
df.collect()

?

amaliujia · 2023-04-14T03:34:04Z

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

+      // Short circuit local relation RPCs
+      val localRelation = plan.getRoot.getLocalRelation
+      new SparkResult(
+        Seq.empty[proto.ExecutePlanResponse].iterator.asJava,


I am wondering if you can populate the data from the LocalRelation into an instance of proto.ExecutePlanResponse thus we do not need to branch the code in SparkResult?

beliefer · 2023-04-14T06:34:46Z

So this happens when

It happens when the root plan is LocalRelation.

amaliujia

LGTM

The code changes are reduced much less now.

hvanhovell · 2023-04-18T01:00:33Z

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

+      reader.loadNextBatch()
+      val batch = proto.ExecutePlanResponse.ArrowBatch
+        .newBuilder()
+        .setRowCount(reader.getVectorSchemaRoot.getRowCount)


Do we need to set row count? You can just modify SparkResult to make the assert a bit more lenient.

hvanhovell · 2023-04-18T01:01:05Z

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

+      val localRelation = plan.getRoot.getLocalRelation
+      val response = proto.ExecutePlanResponse.newBuilder()
+      val reader = new ArrowStreamReader(localRelation.getData.newInput(), allocator)
+      reader.loadNextBatch()


I think we need to close the reader, otherwise the vector root is never cleaned up (memory leak).

Sorry. I forgot it.

beliefer · 2023-04-18T12:41:35Z

@hvanhovell The failed GA is unrelated to this PR.

hvanhovell · 2023-04-21T02:31:47Z

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

+        .newBuilder()
+        .setData(localRelation.getData)
+        .build()
+      response.setArrowBatch(batch).setIsLocalBuilt(true)


Why do we need this? I don't think we should change the protocol for this.

hvanhovell · 2023-04-21T02:32:31Z

...ctor/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala

-            assert(root.getRowCount == response.getArrowBatch.getRowCount) // HUH!
+            assert(
+              response.getIsLocalBuilt ||
+                root.getRowCount == response.getArrowBatch.getRowCount


I'd just make the row count optional. It is kinda weird it is in the protocol anyway.

You means change the protocol that make the row count optional ?

Error: Field "1" on message "ArrowBatch" moved from outside to inside a oneof. [8](https://github.com/beliefer/spark/actions/runs/4761585011/jobs/8462987474#step:5:9) Error: buf found 1 breaking changes.

It seems 1 breaking change. @hvanhovell

I guess adding the row count as optional to LocalRelation message is more useful.
The server can check if the expected number of rows are provided, and we can reuse it here to set the row count.

@hvanhovell

@hvanhovell What do you think about my idea if making the row count optional causes the breaking change?

@hvanhovell Could you take a close look ?

ueshin

~~I was trying to implement this in Python client, but I'm not sure whether this is useful or not.~~
~~We need at least one roundtrip to retrieve some configs to provide a right result, e.g., spark.sql.session.timeZone.~~

~~The test below won't pass without retrieving the config:~~
https://github.com/apache/spark/blob/d7712776bf28d7c9c0d523624bec70a4bc2d97a0/python/pyspark/sql/tests/test_arrow.py#L335-L358

~~I'm not sure it's the same in Scala client, though.~~
~~Just FYI.~~

After some more investigation, I realized we need to retrieve some configs in some cases anyway.
I think we should go ahead with this. @beliefer

ueshin · 2023-04-21T20:57:59Z

...ctor/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala

-            assert(root.getRowCount == response.getArrowBatch.getRowCount) // HUH!
+            assert(
+              response.getIsLocalBuilt ||
+                root.getRowCount == response.getArrowBatch.getRowCount


I guess adding the row count as optional to LocalRelation message is more useful.
The server can check if the expected number of rows are provided, and we can reuse it here to set the row count.

beliefer · 2023-05-18T07:15:07Z

Because there are difference suggestion from @hvanhovell and @ueshin, I don't know how to continue this job.

beliefer · 2023-05-19T10:04:56Z

@ueshin @hvanhovell Recently, #41064 added the rowCount statistics to LocalRelation. In this PR, @ueshin also suggested to add the row count as optional to LocalRelation message.
So I think this is a chance to add optional row count to LocalRelation message.

beliefer · 2023-05-24T03:46:24Z

cc @HyukjinKwon

beliefer · 2023-06-14T13:15:34Z

ping @hvanhovell cc @ueshin

github-actions · 2023-09-23T00:16:45Z

We're closing this PR because it hasn't been updated in a while. This isn't a judgement on the merit of the PR in any way. It's just a way of keeping the PR queue manageable.
If you'd like to revive this PR, please reopen it and ask a committer to remove the Stale tag!

beliefer · 2023-10-14T05:36:09Z

@hvanhovell Do we need this change?

github-actions · 2024-01-23T00:19:43Z

We're closing this PR because it hasn't been updated in a while. This isn't a judgement on the merit of the PR in any way. It's just a way of keeping the PR queue manageable.
If you'd like to revive this PR, please reopen it and ask a committer to remove the Stale tag!

github-actions · 2024-05-03T00:19:01Z

We're closing this PR because it hasn't been updated in a while. This isn't a judgement on the merit of the PR in any way. It's just a way of keeping the PR queue manageable.
If you'd like to revive this PR, please reopen it and ask a committer to remove the Stale tag!

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

amaliujia reviewed Apr 14, 2023

View reviewed changes

hvanhovell reviewed Apr 18, 2023

View reviewed changes

github-actions bot added CORE PYTHON labels Apr 18, 2023

hvanhovell reviewed Apr 21, 2023

View reviewed changes

ueshin reviewed Apr 21, 2023

View reviewed changes

beliefer force-pushed the SPARK-42669 branch from d771277 to c8f4a91 Compare May 4, 2023 01:58

beliefer force-pushed the SPARK-42669 branch from c8f4a91 to badf69d Compare May 19, 2023 11:53

beliefer added 9 commits June 1, 2023 17:47

[SPARK-42669][CONNECT] Short circuit local relation RPC

d81e62c

scala reformat

557c36b

Update code

e446aaf

Update code

e887b5b

Update code

9d633dc

Add optional row count to LocalRelation message.

702115c

Update code

a6336fa

Update code

a83b3a9

Fix conflict

b81e406

beliefer force-pushed the SPARK-42669 branch from 1565235 to b81e406 Compare June 1, 2023 10:32

github-actions bot added the Stale label Sep 23, 2023

github-actions bot closed this Sep 24, 2023

beliefer removed the Stale label Oct 14, 2023

beliefer reopened this Oct 14, 2023

github-actions bot removed the CORE label Oct 14, 2023

github-actions bot added the Stale label Jan 23, 2024

beliefer removed the Stale label Jan 23, 2024

github-actions bot added the Stale label May 3, 2024

github-actions bot closed this May 4, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-42669][CONNECT] Short circuit local relation RPCs #40782

[SPARK-42669][CONNECT] Short circuit local relation RPCs #40782

beliefer commented Apr 13, 2023

beliefer commented Apr 14, 2023 •

edited

zhengruifeng commented Apr 14, 2023

amaliujia commented Apr 14, 2023

amaliujia Apr 14, 2023

beliefer Apr 14, 2023 •

edited

beliefer commented Apr 14, 2023

amaliujia left a comment

hvanhovell Apr 18, 2023

hvanhovell Apr 18, 2023

beliefer Apr 18, 2023

beliefer commented Apr 18, 2023

hvanhovell Apr 21, 2023

hvanhovell Apr 21, 2023

beliefer Apr 21, 2023

beliefer Apr 21, 2023 •

edited

ueshin Apr 21, 2023

beliefer Apr 23, 2023

ueshin May 3, 2023

beliefer May 10, 2023

ueshin left a comment •

edited

ueshin Apr 21, 2023

beliefer commented May 18, 2023

beliefer commented May 19, 2023

beliefer commented May 24, 2023

beliefer commented Jun 14, 2023

github-actions bot commented Sep 23, 2023

beliefer commented Oct 14, 2023

github-actions bot commented Jan 23, 2024

github-actions bot commented May 3, 2024

[SPARK-42669][CONNECT] Short circuit local relation RPCs #40782

[SPARK-42669][CONNECT] Short circuit local relation RPCs #40782

Conversation

beliefer commented Apr 13, 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?

beliefer commented Apr 14, 2023 • edited

zhengruifeng commented Apr 14, 2023

amaliujia commented Apr 14, 2023

Choose a reason for hiding this comment

beliefer Apr 14, 2023 • edited

Choose a reason for hiding this comment

beliefer commented Apr 14, 2023

amaliujia left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

beliefer commented Apr 18, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

beliefer Apr 21, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ueshin left a comment • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

beliefer commented May 18, 2023

beliefer commented May 19, 2023

beliefer commented May 24, 2023

beliefer commented Jun 14, 2023

github-actions bot commented Sep 23, 2023

beliefer commented Oct 14, 2023

github-actions bot commented Jan 23, 2024

github-actions bot commented May 3, 2024

beliefer commented Apr 14, 2023 •

edited

beliefer Apr 14, 2023 •

edited

beliefer Apr 21, 2023 •

edited

ueshin left a comment •

edited