feat: Support Broadcast HashJoin #211

viirya · 2024-03-16T18:01:54Z

Which issue does this PR close?

Closes #202.

Rationale for this change

What changes are included in this PR?

How are these changes tested?

viirya

This is based on #194.

viirya · 2024-03-16T18:57:22Z

common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/ArrowReaderIterator.scala

+    // Release the previous batch.
+    // If it is not released, when closing the reader, arrow library will complain about
+    // memory leak.
+    if (currentBatch != null) {
+      currentBatch.close()
+    }
+


We need to release the batch before loading next batch. Because ArrowStreamReader loads data into same vectors of root internally. After loading next batch, close will release the just loaded batch instead of previous batch.

Is this related to the memory leak we saw?

Because ArrowStreamReader loads data into same vectors of root internally. After loading next batch, close will release the just loaded batch instead of previous batch.

This sounds like a data corruption problem. If the just loaded batch is closed/released, the just loaded ColumnarBatch would be corrupted? But it seems like that the CI passes without any issue previously.

When working on #206, I also found out it might be inconvenient to use Arrow Java's memory API. It requires extra caution to allocate and release ArrowBuf correctly.

Is this related to the memory leak we saw?

It's not, although I suspected it before too. For shuffle, a channel only contains one batch, so ArrowReaderIterator doesn't hit this issue.

This sounds like a data corruption problem. If the just loaded batch is closed/released, the just loaded ColumnarBatch would be corrupted? But it seems like that the CI passes without any issue previously.

When working on #206, I also found out it might be inconvenient to use Arrow Java's memory API. It requires extra caution to allocate and release ArrowBuf correctly.

Due to #211 (comment), this issue is not exposed before.

I feel that Arrow Java API is hard to use and somehow counter-intuitive, especially compared with arrow-rs.

Yes I feel the same pain when using Java Arrow. I think in the long term we'd better to switch away from it. It should be relatively easy except the Java Arrow Flight feature.

codecov-commenter · 2024-03-16T19:39:03Z

Codecov Report

Attention: Patch coverage is 49.33333% with 38 lines in your changes are missing coverage. Please review.

Project coverage is 33.38%. Comparing base (ce63ff8) to head (de73821).

Files	Patch %	Lines
...n/scala/org/apache/spark/sql/comet/operators.scala	60.71%	2 Missing and 9 partials ⚠️
...org/apache/comet/CometSparkSessionExtensions.scala	61.53%	4 Missing and 6 partials ⚠️
.../java/org/apache/comet/CometArrowStreamWriter.java	0.00%	6 Missing ⚠️
...ain/scala/org/apache/comet/vector/NativeUtil.scala	0.00%	5 Missing ⚠️
.../scala/org/apache/comet/serde/QueryPlanSerde.scala	42.85%	1 Missing and 3 partials ⚠️
.../comet/execution/shuffle/ArrowReaderIterator.scala	0.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main     #211      +/-   ##
============================================
+ Coverage     33.32%   33.38%   +0.06%     
- Complexity      769      776       +7     
============================================
  Files           107      108       +1     
  Lines         37037    37099      +62     
  Branches       8106     8129      +23     
============================================
+ Hits          12342    12386      +44     
- Misses        22098    22099       +1     
- Partials       2597     2614      +17

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala

viirya · 2024-03-21T21:09:26Z

cc @sunchao

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala

sunchao · 2024-03-22T07:01:52Z

common/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/ArrowReaderIterator.scala

+    // Release the previous batch.
+    // If it is not released, when closing the reader, arrow library will complain about
+    // memory leak.
+    if (currentBatch != null) {
+      currentBatch.close()
+    }
+


Is this related to the memory leak we saw?

advancedxy · 2024-03-22T07:57:06Z

common/src/main/scala/org/apache/comet/vector/NativeUtil.scala

    var rowCount = 0

    batches.foreach { batch =>
      val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
-      val root = schemaRoot.getOrElse(new VectorSchemaRoot(fieldVectors.asJava))
+      val root = new VectorSchemaRoot(fieldVectors.asJava)
      val provider = batchProviderOpt.getOrElse(dictionaryProvider)


One related question what if incoming batches have different dictionary provider?

If the batch has its provider, it should be returned in batchProviderOpt ?

But the writer is reused. Once the writer is created, new dictionary provider(if different from previous one) from new batches is never used/ written?

Oh I see. I suppose that the dictionary provider is same across batches. This seems to be the reason why there is dictionary provider, i.e. to store dictionary values for arrays/batches.

Hmm. It seems getBatchFieldVectors only checks same dictionary provider across arrays but not batches. Maybe we should add that too? Anyway, it's kind of out of this PR's scope. Maybe in a separate issue to track that.

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

advancedxy

LGTM

advancedxy · 2024-03-23T15:54:56Z

common/src/main/scala/org/apache/comet/vector/NativeUtil.scala

    var rowCount = 0

    batches.foreach { batch =>
      val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
-      val root = schemaRoot.getOrElse(new VectorSchemaRoot(fieldVectors.asJava))
+      val root = new VectorSchemaRoot(fieldVectors.asJava)
      val provider = batchProviderOpt.getOrElse(dictionaryProvider)


Hmm. It seems getBatchFieldVectors only checks same dictionary provider across arrays but not batches. Maybe we should add that too? Anyway, it's kind of out of this PR's scope. Maybe in a separate issue to track that.

viirya · 2024-03-26T04:59:38Z

@sunchao any more comments?

sunchao

LGTM. Can we merge #210 first and trigger another CI run?

viirya · 2024-03-26T06:01:43Z

Let me take a look at #210 tomorrow (today is late, I might not be able to finish it).

I'll hold this until #210 is merged.

viirya · 2024-03-26T18:18:02Z

Merged. Thanks.

* feat: Support HashJoin * Add comment * Clean up test * Fix join filter * Fix clippy * Use consistent function with sort merge join * Add note about left semi and left anti joins * feat: Support BroadcastHashJoin * Move tests * Remove unused import * Add function to parse join parameters * Remove duplicate code * For review

viirya added 7 commits March 13, 2024 13:12

feat: Support HashJoin

8078466

Add comment

fefcf01

Clean up test

cc13619

Fix join filter

334d7d9

Fix clippy

cbd87cf

Use consistent function with sort merge join

c95659c

Add note about left semi and left anti joins

895160b

viirya commented Mar 16, 2024

View reviewed changes

viirya force-pushed the broadcast_hash_join branch 3 times, most recently from c3ed3ee to 1b50512 Compare March 16, 2024 18:54

viirya commented Mar 16, 2024

View reviewed changes

feat: Support BroadcastHashJoin

772588b

viirya force-pushed the broadcast_hash_join branch from 1b50512 to 772588b Compare March 16, 2024 21:56

viirya added 6 commits March 18, 2024 20:01

Merge remote-tracking branch 'upstream/main' into broadcast_hash_join

4351b14

Move tests

d329d3f

Remove unused import

85afec8

Add function to parse join parameters

64240fa

Merge remote-tracking branch 'upstream/main' into broadcast_hash_join

ba4bada

Remove duplicate code

187ba36

viirya force-pushed the broadcast_hash_join branch from 50895b1 to 187ba36 Compare March 21, 2024 21:03

viirya commented Mar 21, 2024

View reviewed changes

spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala Show resolved Hide resolved

sunchao reviewed Mar 22, 2024

View reviewed changes

advancedxy reviewed Mar 22, 2024

View reviewed changes

For review

ce9e1ef

advancedxy approved these changes Mar 23, 2024

View reviewed changes

sunchao approved these changes Mar 26, 2024

View reviewed changes

Merge remote-tracking branch 'upstream/main' into broadcast_hash_join

de73821

viirya merged commit 0826772 into apache:main Mar 26, 2024
28 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Support Broadcast HashJoin #211

feat: Support Broadcast HashJoin #211

viirya commented Mar 16, 2024

viirya left a comment

viirya Mar 16, 2024

sunchao Mar 22, 2024

advancedxy Mar 22, 2024

viirya Mar 22, 2024

viirya Mar 22, 2024

sunchao Mar 23, 2024

codecov-commenter commented Mar 16, 2024 •

edited

Loading

viirya commented Mar 21, 2024

sunchao Mar 22, 2024

advancedxy Mar 22, 2024 •

edited

Loading

viirya Mar 23, 2024

advancedxy Mar 23, 2024

viirya Mar 23, 2024

advancedxy Mar 23, 2024

advancedxy left a comment

advancedxy Mar 23, 2024

viirya commented Mar 26, 2024

sunchao left a comment

viirya commented Mar 26, 2024

viirya commented Mar 26, 2024

feat: Support Broadcast HashJoin #211

feat: Support Broadcast HashJoin #211

Conversation

viirya commented Mar 16, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

viirya 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov-commenter commented Mar 16, 2024 • edited Loading

Codecov Report

viirya commented Mar 21, 2024

Choose a reason for hiding this comment

advancedxy Mar 22, 2024 • edited Loading

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

advancedxy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

viirya commented Mar 26, 2024

sunchao left a comment

Choose a reason for hiding this comment

viirya commented Mar 26, 2024

viirya commented Mar 26, 2024

codecov-commenter commented Mar 16, 2024 •

edited

Loading

advancedxy Mar 22, 2024 •

edited

Loading