feat: Support Broadcast HashJoin #211
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
|
|
||
| package org.apache.comet; | ||
|
|
||
| import java.io.IOException; | ||
| import java.nio.channels.WritableByteChannel; | ||
|
|
||
| import org.apache.arrow.vector.VectorSchemaRoot; | ||
| import org.apache.arrow.vector.VectorUnloader; | ||
| import org.apache.arrow.vector.compression.NoCompressionCodec; | ||
| import org.apache.arrow.vector.dictionary.DictionaryProvider; | ||
| import org.apache.arrow.vector.ipc.ArrowStreamWriter; | ||
| import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; | ||
|
|
||
| /** | ||
| * A custom `ArrowStreamWriter` that allows writing batches from different root to the same stream. | ||
| * Arrow `ArrowStreamWriter` cannot change the root after initialization. | ||
| */ | ||
| public class CometArrowStreamWriter extends ArrowStreamWriter { | ||
| public CometArrowStreamWriter( | ||
| VectorSchemaRoot root, DictionaryProvider provider, WritableByteChannel out) { | ||
| super(root, provider, out); | ||
| } | ||
|
|
||
| public void writeMoreBatch(VectorSchemaRoot root) throws IOException { | ||
| VectorUnloader unloader = | ||
| new VectorUnloader( | ||
| root, /*includeNullCount*/ true, NoCompressionCodec.INSTANCE, /*alignBuffers*/ true); | ||
|
|
||
| try (ArrowRecordBatch batch = unloader.getRecordBatch()) { | ||
| writeRecordBatch(batch); | ||
| } | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,7 +23,7 @@ import java.nio.channels.ReadableByteChannel | |
|
|
||
| import org.apache.spark.sql.vectorized.ColumnarBatch | ||
|
|
||
| import org.apache.comet.vector._ | ||
|
|
||
| class ArrowReaderIterator(channel: ReadableByteChannel) extends Iterator[ColumnarBatch] { | ||
|
|
||
|
|
@@ -36,6 +36,13 @@ class ArrowReaderIterator(channel: ReadableByteChannel) extends Iterator[Columna | |
| return true | ||
| } | ||
|
|
||
| // 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() | ||
| } | ||
|
|
||
|
Comment on lines
+39
to
+45
There was a problem hiding this comment. We need to release the batch before loading next batch. Because There was a problem hiding this comment. Is this related to the memory leak we saw? There was a problem hiding this comment.
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. There was a problem hiding this comment.
It's not, although I suspected it before too. For shuffle, a channel only contains one batch, so There was a problem hiding this comment.
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. There was a problem hiding this comment. 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. |
||
| batch = nextBatch() | ||
| if (batch.isEmpty) { | ||
| return false | ||
|
|
@@ -50,13 +57,6 @@ class ArrowReaderIterator(channel: ReadableByteChannel) extends Iterator[Columna | |
|
|
||
| val nextBatch = batch.get | ||
|
|
||
| currentBatch = nextBatch | ||
| batch = None | ||
| currentBatch | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One related question what if incoming batches have different dictionary provider?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the batch has its provider, it should be returned in
batchProviderOpt?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm. It seems
getBatchFieldVectorsonly 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.