MaxGekk commented on code in PR #58850:
URL: https://github.com/apache/spark/pull/58850#discussion_r4029330936
##########
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:
Both points are right, and I traced each against master rather than taking
them on trust. Fixed in 96d1ab9.
On the guard: agreed, the three assertions did not establish it.
`relation.output` is `child.output`, fixed at cache time by
`InMemoryRelation.apply`, so that assertion only re-read the input schema; the
read path rebuilds its Arrow schema from `cacheSchema` via
`ArrowUtils.toArrowSchema(..., losslessInternalTypes = true)` and the payload
is the record batch alone, so nothing about the precision travels in the cached
bytes (and `isCompatibleWithDeclaredField` ignores field metadata on the write
side too); and all four precisions land in `TimeNanoVector` with the value
untouched while the test pre-truncates its inputs, so neither `checkAnswer` nor
the vector-class check could separate them.
There is a metadata round trip in the cache read path, though, one level
below where the test was looking: `new ArrowColumnVector(vector)` is
`this(ArrowUtils.fromArrowField(vector.getField()))`, and the vector's field is
the one `toArrowSchema(cacheSchema, ...)` produced, tagged with
`SPARK::time::precision` by `toPrecisionTaggedArrowField`. So the column's type
on the way out is the Spark type recovered from the round-tripped Arrow field,
which is the guard you asked for. `cachedVectors` now reports that type and
each case asserts it is `TimeType(precision)`; the tautological assertion is
gone.
I checked that it bites rather than assuming it: with `fromArrowField`
mutated to read a key that is not there, the TIME case fails with
`Set(TimeType(6)) did not equal Set(TimeType(0))` - the canonical-microsecond
fallback - at precisions 0, 3 and 9, while `checkAnswer` still passes, which is
exactly your point that the values alone cannot carry this.
The "Why" is reworded to match: the precision is reconstructed from the
cache schema on read rather than persisted in the cached bytes, and what the
new assertion pins is the tag-and-recover path that reconstruction goes
through. Thank you for reading it that closely.
##########
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:
Correct, and that one is straightforwardly my mistake: the block went in
between the `cachedStats` doc comment and `cachedStats` itself, so the comment
ended up describing `cachedRelation`, which returns an `InMemoryRelation` and
no stats at all. `cachedRelation`, `cachedVectors`, `nullPatterns` and the two
tests now sit after `singlePartDf`, which puts `cachedStats` back under its own
comment and keeps the new helpers next to the helper they build on.
--
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]
