This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 80e4a2e9 chore: Close dictionary provider when iterator is closed
(#904)
80e4a2e9 is described below
commit 80e4a2e98faeacc2fb0c041bd75158014921b9d6
Author: Andy Grove <[email protected]>
AuthorDate: Tue Sep 3 15:01:44 2024 -0600
chore: Close dictionary provider when iterator is closed (#904)
* Close dictionary provider when iterator is closed
* add more docs
* format
---
.../java/org/apache/arrow/c/ArrowImporter.java | 5 +----
.../scala/org/apache/comet/vector/NativeUtil.scala | 26 +++++++++++++++++++++-
.../scala/org/apache/comet/CometExecIterator.scala | 3 ++-
3 files changed, 28 insertions(+), 6 deletions(-)
diff --git a/common/src/main/java/org/apache/arrow/c/ArrowImporter.java
b/common/src/main/java/org/apache/arrow/c/ArrowImporter.java
index 1f0cbd41..d0d1754e 100644
--- a/common/src/main/java/org/apache/arrow/c/ArrowImporter.java
+++ b/common/src/main/java/org/apache/arrow/c/ArrowImporter.java
@@ -40,15 +40,12 @@ public class ArrowImporter {
}
Field importField(ArrowSchema schema, CDataDictionaryProvider provider) {
- Field var4;
try {
- var4 = importer.importField(schema, provider);
+ return importer.importField(schema, provider);
} finally {
schema.release();
schema.close();
}
-
- return var4;
}
public FieldVector importVector(
diff --git a/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala
b/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala
index 5149c734..751aa6bb 100644
--- a/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala
+++ b/common/src/main/scala/org/apache/comet/vector/NativeUtil.scala
@@ -30,13 +30,32 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.comet.CometArrowAllocator
+/**
+ * Provides functionality for importing Arrow vectors from native code and
wrapping them as
+ * CometVectors.
+ *
+ * Also provides functionality for exporting Comet columnar batches to native
code.
+ *
+ * Each instance of NativeUtil creates an instance of CDataDictionaryProvider
(a
+ * DictionaryProvider that is used in C Data Interface for imports).
+ *
+ * NativeUtil must be closed after use to release resources in the dictionary
provider.
+ */
class NativeUtil {
import Utils._
+ /** Use the global allocator */
private val allocator = CometArrowAllocator
- private val dictionaryProvider: CDataDictionaryProvider = new
CDataDictionaryProvider
+
+ /** ArrowImporter does not hold any state and does not need to be closed */
private val importer = new ArrowImporter(allocator)
+ /**
+ * Dictionary provider to use for the lifetime of this instance of
NativeUtil. The dictionary
+ * provider is closed when NativeUtil is closed.
+ */
+ private val dictionaryProvider: CDataDictionaryProvider = new
CDataDictionaryProvider
+
/**
* Exports a Comet `ColumnarBatch` into a list of memory addresses that can
be consumed by the
* native execution.
@@ -133,6 +152,11 @@ class NativeUtil {
new ColumnarBatch(arrayVectors.toArray, maxNumRows)
}
+
+ def close(): Unit = {
+ // closing the dictionary provider also closes the dictionary arrays
+ dictionaryProvider.close()
+ }
}
object NativeUtil {
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index 29eb2f0c..dcdc8ae9 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -48,7 +48,7 @@ class CometExecIterator(
extends Iterator[ColumnarBatch] {
private val nativeLib = new Native()
- private val nativeUtil = new NativeUtil
+ private val nativeUtil = new NativeUtil()
private val cometBatchIterators = inputs.map { iterator =>
new CometBatchIterator(iterator, nativeUtil)
}.toArray
@@ -157,6 +157,7 @@ class CometExecIterator(
currentBatch.close()
currentBatch = null
}
+ nativeUtil.close()
nativeLib.releasePlan(plan)
// The allocator thoughts the exported ArrowArray and ArrowSchema
structs are not released,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
