comphead commented on code in PR #7405:
URL: https://github.com/apache/datafusion-comet/pull/7405#discussion_r4460104197
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -326,13 +443,36 @@ class NativeUtil extends AutoCloseable {
}
override def close(): Unit = {
- // closing the dictionary provider also closes the dictionary arrays
- dictionaryProvider.close()
+ val schemas = planOutputSchemas
+ planOutputSchemas = null
+ try {
+ if (schemas != null) releaseArrowStructs(Array.empty, schemas)
+ } finally {
+ // closing the dictionary provider also closes the dictionary arrays
+ dictionaryProvider.close()
+ }
}
}
object NativeUtil {
+ /**
+ * Set in what `Native.executePlan` returns when native filled the schema
structs as well as the
+ * arrays. The row count is in the low 32 bits. Matches
`OUTPUT_SCHEMA_EXPORTED` in
+ * `jni_api.rs`.
+ */
+ val OUTPUT_SCHEMA_EXPORTED: Long = 1L << 32
+
+ /**
+ * How many times [[NativeUtil.getNextPlanBatch]] has imported an output
schema in this JVM, so
+ * that tests and benchmarks can check that a plan sends it only when it
changes.
+ */
+ private[apache] val planOutputSchemaImports = new LongAdder()
+
+ /** Whether `field` or a field nested in it is dictionary-encoded. */
+ private def hasDictionary(field: Field): Boolean =
+ field.getDictionary != null ||
field.getChildren.asScala.exists(hasDictionary)
Review Comment:
`CometArrowPythonRunnerBase` (under `spark-4.x`) has the same private
`hasDictionary(field: Field)`. Could this one be `private[apache]`, or move to
`CometVectorUtils`, so the Python runner reuses it instead of keeping a second
copy?
##########
spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala:
##########
@@ -312,6 +316,163 @@ class NativeUtilSuite extends CometTestBase {
}
}
+ test("getNextPlanBatch imports batches sent without a schema with the fields
sent last") {
+ withIsolatedStructAllocator { (nativeUtil, allocator, _) =>
+ val schemaAddrs = mutable.ArrayBuffer.empty[Seq[Long]]
+ def nextBatch(values: Seq[Int], withSchema: Boolean): ColumnarBatch =
+ Using.resource(new IntVector("value", allocator)) { vector =>
+ vector.allocateNew(values.length)
+ values.indices.foreach(i => vector.setSafe(i, values(i)))
+ vector.setValueCount(values.length)
+ nativeUtil
+ .getNextPlanBatch(
+ 1,
+ (arrays, schemas) => {
+ schemaAddrs += schemas.toSeq
+ val schema = if (withSchema) ArrowSchema.wrap(schemas(0)) else
null
+ Data.exportVector(allocator, vector, null,
ArrowArray.wrap(arrays(0)), schema)
+ val rows = values.length.toLong
+ if (withSchema) rows | NativeUtil.OUTPUT_SCHEMA_EXPORTED else
rows
+ })
+ .get
+ }
+
+ val imports = NativeUtil.planOutputSchemaImports.sum()
+ val batches = Seq(
+ nextBatch(Seq(1, 2), withSchema = true),
+ nextBatch(Seq(3), withSchema = false),
+ nextBatch(Seq(4, 5, 6), withSchema = false))
Review Comment:
Could this sequence also send a batch of another type with its schema (a
`BigIntVector`, for example) and then one more without a schema, and check that
the last one reads back as longs? With the current tests, a `getNextPlanBatch`
that kept only the first schema's fields would still pass: the schema arrives
once here and in the dictionary test, and with the cache off the fields it
keeps are never used. The native side of that sequence is covered by
`output_schema_cache_writes_the_schema_only_when_it_changes`, but nothing on
the JVM side reads a batch without a schema after a second schema.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
