viirya commented on code in PR #58850:
URL: https://github.com/apache/spark/pull/58850#discussion_r4029576496
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala:
##########
@@ -1260,6 +1260,92 @@ class ArrowCachedBatchSerializerSuite extends QueryTest
with SharedSparkSession
spark.sparkContext.parallelize(values.map(v => Row(v)), 1),
StructType(Seq(StructField("v", dt, nullable = true))))
+ // 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, Spark type, row count, null count) per
cached batch of a
+ // single-column relation, read back through the serializer's own columnar
path. The Spark type
+ // is the one ArrowColumnVector recovers from the read-side Arrow field via
+ // ArrowUtils.fromArrowField, so a type whose fractional-second precision
lives in that field's
+ // metadata is reported as it comes back rather than as it went in.
+ private def cachedVectors(relation: InMemoryRelation): Array[(String,
DataType, 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, column.dataType(),
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, so the
precision rides in
+ // the Arrow field metadata rather than in the value. The cache read path
rebuilds its Arrow
+ // schema from the cache schema and wraps each vector in an
ArrowColumnVector, whose type is
+ // recovered from that field, so the column's type on the way out pins the
tag-and-recover
+ // path: dropping the precision key on either side falls back to
TimeType(6). 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(_)))
+ val vectors = cachedVectors(relation)
+ assert(vectors.map(_._1).toSet === Set("TimeNanoVector"),
s"TIME($precision), $pattern")
+ assert(vectors.map(_._2).toSet === Set[DataType](TimeType(precision)),
+ s"TIME($precision), $pattern: the precision must survive the read-side
Arrow field")
+ assert(vectors.map(_._3).sum === rows, s"TIME($precision), $pattern")
+ assert(vectors.map(_._4).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)
Review Comment:
Could we include values with nonzero sub-millisecond components here?
`7001000000L` nanoseconds is exactly 7001 milliseconds, so every generated
interval is millisecond-aligned. A regression that truncates the stored
microseconds with `micros / 1000 * 1000` would leave all these values
unchanged, and the type/null-count assertions would still pass.
Using a multiplier such as `7001001000L` would exercise microsecond
precision while retaining both signs. Explicit `+/-1` and `+/-999` microsecond
cases would also help cover truncation around zero.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala:
##########
@@ -1260,6 +1260,92 @@ class ArrowCachedBatchSerializerSuite extends QueryTest
with SharedSparkSession
spark.sparkContext.parallelize(values.map(v => Row(v)), 1),
StructType(Seq(StructField("v", dt, nullable = true))))
+ // 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, Spark type, row count, null count) per
cached batch of a
+ // single-column relation, read back through the serializer's own columnar
path. The Spark type
+ // is the one ArrowColumnVector recovers from the read-side Arrow field via
+ // ArrowUtils.fromArrowField, so a type whose fractional-second precision
lives in that field's
+ // metadata is reported as it comes back rather than as it went in.
+ private def cachedVectors(relation: InMemoryRelation): Array[(String,
DataType, 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, column.dataType(),
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, so the
precision rides in
+ // the Arrow field metadata rather than in the value. The cache read path
rebuilds its Arrow
+ // schema from the cache schema and wraps each vector in an
ArrowColumnVector, whose type is
+ // recovered from that field, so the column's type on the way out pins the
tag-and-recover
+ // path: dropping the precision key on either side falls back to
TimeType(6). 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) {
Review Comment:
Small coverage/naming point: `TimeType` supports every integer precision
from 0 through 9, so this leaves 1, 2, 4, 5, 7, and 8 untested despite “every
precision” in the PR title and test name.
Could we iterate over `TimeType.MIN_PRECISION to TimeType.MAX_PRECISION`?
Alternatively, if these four representative precisions are intentional to limit
test runtime, please adjust the wording accordingly.
--
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]
