sunchao commented on code in PR #5560:
URL: https://github.com/apache/datafusion-comet/pull/5560#discussion_r3991507909
##########
spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala:
##########
@@ -336,6 +394,242 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite
with Matchers {
}
}
+ for (failSerialization <- Seq(false, true)) {
+ test(s"dictionary inputs materialize logical values (failure:
$failSerialization)") {
+ val sourceAllocator = new RootAllocator(Long.MaxValue)
+ val writerAllocator = new RootAllocator(Long.MaxValue)
+ val intType = new ArrowType.Int(32, true)
+ val textEncoding = new DictionaryEncoding(11L, false, intType)
+ val binaryEncoding = new DictionaryEncoding(12L, false, intType)
+ val textValues = new VarCharVector("text", sourceAllocator)
+ val binaryValues = new VarBinaryVector("data", sourceAllocator)
+ val textIndices =
+ new IntVector("text", new FieldType(true, intType, textEncoding),
sourceAllocator)
+ val binaryIndices =
+ new IntVector("data", new FieldType(true, intType, binaryEncoding),
sourceAllocator)
+ val dictionaries = Map(
+ textEncoding.getId -> new Dictionary(textValues, textEncoding),
+ binaryEncoding.getId -> new Dictionary(binaryValues, binaryEncoding))
+ val provider = new DictionaryProvider {
+ override def lookup(id: Long): Dictionary = dictionaries(id)
+
+ override def getDictionaryIds: java.util.Set[java.lang.Long] =
+ dictionaries.keys.map(id => java.lang.Long.valueOf(id)).toSet.asJava
+ }
+ val columns = Seq[CometDecodedVector](
+ new CometDictionaryVector(
+ new CometPlainVector(textIndices),
+ new CometDictionary(new CometPlainVector(textValues)),
+ provider),
+ new CometDictionaryVector(
+ new CometPlainVector(binaryIndices),
+ new CometDictionary(new CometPlainVector(binaryValues)),
+ provider))
+ var failWrites = false
+ val output = new ByteArrayOutputStream() {
+ override def write(bytes: Array[Byte], offset: Int, length: Int): Unit
= {
+ if (failWrites) {
+ throw new IOException("injected dictionary IPC write failure")
+ }
+ super.write(bytes, offset, length)
+ }
+ }
+ try {
+ textValues.allocateNew()
+ Seq("same", "", "λ中文").zipWithIndex.foreach { case (value, index) =>
+ textValues.setSafe(index, value.getBytes(StandardCharsets.UTF_8))
+ }
+ textValues.setValueCount(3)
+ binaryValues.allocateNew()
+ Seq(Array[Byte](1, 2), Array.emptyByteArray, Array[Byte](0,
-1)).zipWithIndex.foreach {
+ case (value, index) => binaryValues.setSafe(index, value)
+ }
+ binaryValues.setValueCount(3)
+ textIndices.allocateNew()
+ Seq(0, 1, 0, 2).zipWithIndex.foreach { case (value, index) =>
+ textIndices.setSafe(index, value)
+ }
+ textIndices.setNull(2)
+ textIndices.setValueCount(4)
+ binaryIndices.allocateNew()
+ Seq(2, 0, 0, 1).zipWithIndex.foreach { case (value, index) =>
+ binaryIndices.setSafe(index, value)
+ }
+ binaryIndices.setNull(2)
+ binaryIndices.setValueCount(4)
+
+ val sourceVectors = Seq(textValues, binaryValues, textIndices,
binaryIndices)
+ val sourceBuffers = sourceVectors.flatMap(_.getFieldBuffers.asScala)
+ val sourceRefs = sourceBuffers.map(_.refCnt())
+ val sourceBytes = sourceAllocator.getAllocatedMemory
+
+ def writeDictionaryBatch(): Unit =
+ withMaterializedInputVectors(columns, writerAllocator) { vectors =>
+ vectors.map(_.getField.getDictionary) shouldBe Seq(null, null)
+ vectors.head.getObject(0).toString shouldBe "same"
+ vectors.head.getObject(1).toString shouldBe ""
+ vectors.head.isNull(2) shouldBe true
+ vectors.head.getObject(3).toString shouldBe "λ中文"
+ vectors(1).getObject(0).asInstanceOf[Array[Byte]] shouldBe
Array[Byte](0, -1)
+ vectors(1).isNull(2) shouldBe true
+
+ withWriter(vectors.map(_.getField), writerAllocator,
Channels.newChannel(output)) {
+ channel =>
+ failWrites = failSerialization
+ try {
+ serializeBatch(new WriteChannel(channel), vectors, 4,
writerAllocator)
+ } finally {
+ failWrites = false
+ }
+ }
+ }
+
+ if (failSerialization) {
+ val error = intercept[IOException](writeDictionaryBatch())
+ error.getMessage shouldBe "injected dictionary IPC write failure"
+ } else {
+ writeDictionaryBatch()
+ withReader(output.toByteArray) { reader =>
+ reader.loadNextBatch() shouldBe true
+ val struct =
reader.getVectorSchemaRoot.getVector(0).asInstanceOf[StructVector]
+ val resultText = struct.getChild("text")
+ val resultData = struct.getChild("data")
+ resultText.getField.getType shouldBe ArrowType.Utf8.INSTANCE
+ resultData.getField.getType shouldBe ArrowType.Binary.INSTANCE
+ resultText.getObject(0).toString shouldBe "same"
+ resultText.getObject(1).toString shouldBe ""
+ resultText.isNull(2) shouldBe true
+ resultText.getObject(3).toString shouldBe "λ中文"
+ resultData.getObject(0).asInstanceOf[Array[Byte]] shouldBe
Array[Byte](0, -1)
+ resultData.isNull(2) shouldBe true
+ reader.loadNextBatch() shouldBe false
+ }
+ }
+
+ writerAllocator.getAllocatedMemory shouldBe 0L
+ sourceAllocator.getAllocatedMemory shouldBe sourceBytes
+ sourceBuffers.map(_.refCnt()) shouldBe sourceRefs
+ textValues.getObject(0).toString shouldBe "same"
+ binaryValues.getObject(2).asInstanceOf[Array[Byte]] shouldBe
Array[Byte](0, -1)
+ } finally {
+ columns.foreach(_.close())
+ writerAllocator.close()
+ sourceAllocator.close()
+ }
+ }
+ }
+
+ test("dictionary inputs are sliced before decoding to the Arrow batch
limits") {
Review Comment:
Added both requested regressions: [`dictionary slices keep plain nullable
and nested columns aligned through
IPC`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L768)
verifies all four columns and nulls across `[3, 3, 3, 1]` batches, and
[`randomized dictionary ranges match the soft-limit oracle and cover every
input
row`](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala#L914)
calls `inputBatchRanges` directly over 100 seeded configurations. It checks
exact agreement with the soft-limit oracle, contiguous complete coverage, and
record limits. Both passed locally on Spark 4.1.
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -596,9 +596,10 @@ object CometConf extends ShimCometConf {
.withAlternative("spark.comet.shuffle.preferDictionary.ratio")
.category(CATEGORY_SHUFFLE)
.doc(
- "The ratio of total values to distinct values in a string column to
decide whether to " +
+ "The ratio of total values to distinct values in a string or binary
column to decide " +
+ "whether to " +
Review Comment:
Reflowed [the
description](https://github.com/apache/datafusion-comet/blob/e75103f8b51bab3ec440bd9ee0fcb16b0ec4e069/spark/src/main/scala/org/apache/comet/CometConf.scala#L597)
so “whether to prefer dictionary encoding” stays together. The concatenated
documentation string is unchanged, including “string or binary”; verified the
string equality and formatting checks.
--
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]
