uros-b commented on code in PR #58850:
URL: https://github.com/apache/spark/pull/58850#discussion_r4025422086
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala:
##########
@@ -1243,6 +1243,85 @@ class ArrowCachedBatchSerializerSuite extends QueryTest
with SharedSparkSession
// Helper: cache a single-column DataFrame (row path) and return its
ArrowCachedBatch stats.
// Stats layout per column: [lowerBound(0), upperBound(1), nullCount(2),
rowCount(3), size(4)].
+ // Helper: the InMemoryRelation behind a cached DataFrame, after populating
the cache. The
+ // DataFrame's queryExecution is memoized, so this must run before anything
else forces it.
+ private def cachedRelation(df: org.apache.spark.sql.DataFrame):
InMemoryRelation = {
+ df.cache()
+ df.count()
+ df.queryExecution.executedPlan.collectFirst {
+ case scan: InMemoryTableScanExec => scan.relation
+ }.get
+ }
+
+ // Helper: (Arrow vector class, row count, null count) per cached batch of a
single-column
+ // relation, read back through the serializer's own columnar path.
+ private def cachedVectors(relation: InMemoryRelation): Array[(String, Int,
Int)] = {
+ val attrs = relation.output
+ relation.cacheBuilder.serializer
+ .convertCachedBatchToColumnarBatch(
+ relation.cacheBuilder.cachedColumnBuffers, attrs, attrs,
spark.sessionState.conf)
+ .mapPartitions { batches =>
+ batches.map { batch =>
+ val column = batch.column(0).asInstanceOf[ArrowColumnVector]
+ (column.getValueVector.getClass.getSimpleName, batch.numRows(),
column.numNulls())
+ }
+ }
+ .collect()
+ }
+
+ private val nullPatterns: Seq[(String, Int => Boolean)] = Seq(
+ ("every 31st row null", i => i % 31 == 0),
+ ("no nulls", _ => false),
+ ("all nulls", _ => true))
+
+ test("SPARK-59571: TIME keeps its precision and its nulls through the cache,
at every " +
+ "precision") {
+ // Every precision of TIME is written to the same TimeNanoVector; the
precision travels in
+ // the Arrow field metadata and comes back through the cached relation's
schema, so the
+ // cached type is checked as well as the values. The values are truncated
to the declared
+ // precision so a precision loss could not hide behind a value the type
would round anyway.
+ val rows = 1000
+ val nanosPerDay = 86400000000000L
+ for (precision <- Seq(0, 3, 6, 9); (pattern, isNull) <- nullPatterns) {
+ val unit = math.pow(10, 9 - precision).toLong
+ def timeAt(i: Int): LocalTime = {
+ val nanos = ((i.toLong * nanosPerDay) / rows + i.toLong * 1234567L) %
nanosPerDay
+ LocalTime.ofNanoOfDay(nanos / unit * unit)
+ }
+ val values = (0 until rows).map(i => if (isNull(i)) null else timeAt(i))
+ val df = singlePartDf(values, TimeType(precision))
+ val relation = cachedRelation(df)
+ checkAnswer(df, values.map(Row(_)))
+ assert(relation.output.head.dataType === TimeType(precision),
+ s"TIME($precision), $pattern: the cached relation must keep the
precision")
+ val vectors = cachedVectors(relation)
+ assert(vectors.map(_._1).toSet === Set("TimeNanoVector"),
s"TIME($precision), $pattern")
+ assert(vectors.map(_._2).sum === rows, s"TIME($precision), $pattern")
+ assert(vectors.map(_._3).sum === (0 until rows).count(isNull),
+ s"TIME($precision), $pattern: null count")
+ df.unpersist()
+ }
+ }
+
+ test("SPARK-59571: day-time intervals of both signs keep their microseconds
through the " +
+ "cache") {
+ val rows = 1000
+ for ((pattern, isNull) <- nullPatterns) {
+ // Whole microseconds, negative for the first half of the rows and
positive for the rest.
+ def intervalAt(i: Int): Duration = Duration.ofNanos((i.toLong - rows /
2) * 7001000000L)
+ val values = (0 until rows).map(i => if (isNull(i)) null else
intervalAt(i))
+ val df = singlePartDf(values, DayTimeIntervalType())
+ val relation = cachedRelation(df)
+ checkAnswer(df, values.map(Row(_)))
+ assert(relation.output.head.dataType === DayTimeIntervalType(), pattern)
+ val vectors = cachedVectors(relation)
+ assert(vectors.map(_._1).toSet === Set("DurationVector"), pattern)
+ assert(vectors.map(_._2).sum === rows, pattern)
+ assert(vectors.map(_._3).sum === (0 until rows).count(isNull),
s"$pattern: null count")
+ df.unpersist()
+ }
+ }
+
Review Comment:
ArrowCachedBatchSerializerSuite.scala (both new tests) - The PR's motivating
guarantee is not actually established by the assertions. The "Why" states that
TIME precision "travels only in the Arrow field metadata, so a regression that
dropped it on the way back through the cache would leave the values intact and
fail no existing test" — implying these tests catch that. Traced against
master, they do not: (1) relation.output.head.dataType === TimeType(precision)
re-checks the input schema — InMemoryRelation.output is cachedPlan.output
(CacheManager / InMemoryRelation.apply), set at cache time and unaffected by
any serialize/deserialize, so the assertion is tautological w.r.t. the
round-trip; (2) the cache read path rebuilds its schema from cacheSchema via
ArrowUtils.toArrowSchema(cacheSchema, ...) (serializer line 1276) and never
calls fromArrowField, and the write side "serializes only the record batch; the
read path reconstructs the schema from cacheSchema" (serializer line 1
152) — so the SPARK::time::precision field metadata is neither persisted in
the cached bytes nor consulted on read; (3) every precision uses the same
physical TimeNanoVector with the stored Long unchanged, and the input values
are pre-truncated to the precision, so both checkAnswer and the vector-class
assertion are precision-independent. A regression in the metadata round-trip
would pass all three checks. The tests still add genuine breadth (all
precisions run, three null patterns incl. the all-null no-validity-buffer
branch, both interval signs, columnar null-count/vector-class checks), but not
the specific precision-round-trip guard claimed. Either strengthen a case to
recover the Spark type from the serialized/round-tripped Arrow field (e.g.
ArrowUtils.fromArrowField on the read-side schema) so the metadata path is
actually exercised, or reword the justification to match what is verified.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala:
##########
@@ -1243,6 +1243,85 @@ class ArrowCachedBatchSerializerSuite extends QueryTest
with SharedSparkSession
// Helper: cache a single-column DataFrame (row path) and return its
ArrowCachedBatch stats.
// Stats layout per column: [lowerBound(0), upperBound(1), nullCount(2),
rowCount(3), size(4)].
+ // Helper: the InMemoryRelation behind a cached DataFrame, after populating
the cache. The
+ // DataFrame's queryExecution is memoized, so this must run before anything
else forces it.
+ private def cachedRelation(df: org.apache.spark.sql.DataFrame):
InMemoryRelation = {
+ df.cache()
+ df.count()
+ df.queryExecution.executedPlan.collectFirst {
+ case scan: InMemoryTableScanExec => scan.relation
+ }.get
+ }
+
+ // Helper: (Arrow vector class, row count, null count) per cached batch of a
single-column
+ // relation, read back through the serializer's own columnar path.
+ private def cachedVectors(relation: InMemoryRelation): Array[(String, Int,
Int)] = {
+ val attrs = relation.output
+ relation.cacheBuilder.serializer
+ .convertCachedBatchToColumnarBatch(
+ relation.cacheBuilder.cachedColumnBuffers, attrs, attrs,
spark.sessionState.conf)
+ .mapPartitions { batches =>
+ batches.map { batch =>
+ val column = batch.column(0).asInstanceOf[ArrowColumnVector]
+ (column.getValueVector.getClass.getSimpleName, batch.numRows(),
column.numNulls())
+ }
+ }
+ .collect()
+ }
+
+ private val nullPatterns: Seq[(String, Int => Boolean)] = Seq(
+ ("every 31st row null", i => i % 31 == 0),
+ ("no nulls", _ => false),
+ ("all nulls", _ => true))
+
+ test("SPARK-59571: TIME keeps its precision and its nulls through the cache,
at every " +
+ "precision") {
+ // Every precision of TIME is written to the same TimeNanoVector; the
precision travels in
+ // the Arrow field metadata and comes back through the cached relation's
schema, so the
+ // cached type is checked as well as the values. The values are truncated
to the declared
+ // precision so a precision loss could not hide behind a value the type
would round anyway.
+ val rows = 1000
+ val nanosPerDay = 86400000000000L
+ for (precision <- Seq(0, 3, 6, 9); (pattern, isNull) <- nullPatterns) {
+ val unit = math.pow(10, 9 - precision).toLong
+ def timeAt(i: Int): LocalTime = {
+ val nanos = ((i.toLong * nanosPerDay) / rows + i.toLong * 1234567L) %
nanosPerDay
+ LocalTime.ofNanoOfDay(nanos / unit * unit)
+ }
+ val values = (0 until rows).map(i => if (isNull(i)) null else timeAt(i))
+ val df = singlePartDf(values, TimeType(precision))
+ val relation = cachedRelation(df)
+ checkAnswer(df, values.map(Row(_)))
+ assert(relation.output.head.dataType === TimeType(precision),
+ s"TIME($precision), $pattern: the cached relation must keep the
precision")
+ val vectors = cachedVectors(relation)
+ assert(vectors.map(_._1).toSet === Set("TimeNanoVector"),
s"TIME($precision), $pattern")
+ assert(vectors.map(_._2).sum === rows, s"TIME($precision), $pattern")
+ assert(vectors.map(_._3).sum === (0 until rows).count(isNull),
+ s"TIME($precision), $pattern: null count")
+ df.unpersist()
+ }
+ }
+
+ test("SPARK-59571: day-time intervals of both signs keep their microseconds
through the " +
+ "cache") {
+ val rows = 1000
+ for ((pattern, isNull) <- nullPatterns) {
+ // Whole microseconds, negative for the first half of the rows and
positive for the rest.
+ def intervalAt(i: Int): Duration = Duration.ofNanos((i.toLong - rows /
2) * 7001000000L)
+ val values = (0 until rows).map(i => if (isNull(i)) null else
intervalAt(i))
+ val df = singlePartDf(values, DayTimeIntervalType())
+ val relation = cachedRelation(df)
+ checkAnswer(df, values.map(Row(_)))
+ assert(relation.output.head.dataType === DayTimeIntervalType(), pattern)
+ val vectors = cachedVectors(relation)
+ assert(vectors.map(_._1).toSet === Set("DurationVector"), pattern)
+ assert(vectors.map(_._2).sum === rows, pattern)
+ assert(vectors.map(_._3).sum === (0 until rows).count(isNull),
s"$pattern: null count")
+ df.unpersist()
+ }
+ }
+
Review Comment:
Also, the new helpers and tests are inserted between the cachedStats
doc-comment ("...return its ArrowCachedBatch stats" / "Stats layout per column:
[lowerBound(0)...]") and the cachedStats method it documents (adjacent in
master), orphaning the comment so it now sits directly above cachedRelation
(which returns an InMemoryRelation, not stats) while cachedStats is left
uncommented below the tests. This is the "split a section of unrelated code"
failure mode called out in CLAUDE.md and makes the comment actively misleading.
Move the new block below cachedStats/singlePartDf, or above the comment.
--
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]
