andygrove commented on code in PR #6383:
URL: https://github.com/apache/datafusion-comet/pull/6383#discussion_r4184296344
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala:
##########
@@ -326,41 +337,73 @@ object CometWriteFilesExec extends Logging {
} {
execIterator.close()
}
-
- rowsWritten
}
/**
- * Report `count` rows to each stats tracker.
+ * Create a task's callback for reporting rows while each batch's vectors
are still valid.
*
* `WriteTaskStatsTracker.newRow` is a per-row callback, but the only
implementation Spark
- * ships, [[BasicWriteTaskStatsTracker]], ignores the row and just counts.
Comet has columnar
- * batches rather than `InternalRow`s here, so it passes an empty row
instead of materializing
- * every row just to hand it straight back. A tracker that actually inspects
row contents would
- * therefore see empty rows, so warn rather than silently report wrong
statistics.
+ * ships, [[BasicWriteTaskStatsTracker]], ignores the row and just counts.
Preserve that fast
+ * path without accessing row contents. Subclasses may override `newRow` to
inspect values, so
+ * only the exact basic class qualifies. All other trackers receive
materialized rows:
+ * `ColumnarBatchRow` does not implement every `InternalRow` read method,
such as `anyNull`.
*
- * The loop is per-tracker on the outside so the hot inner loop has a single
receiver and no
- * per-row closure; the trackers are independent per-file counters, so their
relative
- * interleaving carries no meaning.
+ * The projection is created once per task and reuses its row, just as
Spark's write path may
+ * reuse rows. Trackers that retain a row beyond `newRow` must copy it. This
callback does not
+ * take ownership of the batch.
*
* Visible for testing: nothing Spark ships lets a caller install a
third-party tracker on a V1
- * write, so this is the only way to exercise the warning.
+ * write.
*/
def recordRows(
statsTrackers: Seq[WriteTaskStatsTracker],
filePath: String,
- count: Long): Unit = {
- statsTrackers.foreach { tracker =>
- if (!tracker.isInstanceOf[BasicWriteTaskStatsTracker]) {
- logWarning(
- s"${tracker.getClass.getName} receives row counts but not row
contents from Comet's " +
- "native Parquet writer. Set
spark.comet.parquet.write.enabled=false if this tracker " +
- "needs to inspect written rows.")
+ schema: StructType): ColumnarBatch => Unit = {
+ val (countTrackers, rowTrackers) =
+ statsTrackers.partition(_.getClass ==
classOf[BasicWriteTaskStatsTracker])
+ val rowProjection = if (rowTrackers.nonEmpty) {
+ // Arrow erases logical parameters such as interval endpoints. Compare
against the
+ // reconstructed schema too, while preserving the declared schema for
row materialization.
+ val arrowSchema = CometUtils.fromArrowSchema(
+ CometUtils.toArrowSchema(schema, CometArrowStream.NATIVE_TIMEZONE))
+ Some((UnsafeProjection.create(schema), arrowSchema))
+ } else {
+ None
+ }
+ batch => {
+ countTrackers.foreach { tracker =>
+ var i = 0
+ while (i < batch.numRows()) {
+ tracker.newRow(filePath, InternalRow.empty)
+ i += 1
+ }
}
- var i = 0L
- while (i < count) {
- tracker.newRow(filePath, InternalRow.empty)
- i += 1
+ rowProjection.foreach { case (project, arrowSchema) =>
+ // ScanExec can cast differing physical types after import, but a JVM
row projection
+ // would read those buffers with the wrong getters before that cast.
Fail explicitly
+ // rather than report incorrect statistics. Names, nullability, and
logical parameters
+ // erased by the Arrow bridge do not affect the getters used by the
row projection.
+ val matchingTypes = batch.numCols() == schema.length &&
schema.fields.indices.forall {
+ i =>
+ DataType.equalsStructurally(
+ batch.column(i).dataType(),
+ schema(i).dataType,
+ ignoreNullability = true) || DataType.equalsStructurally(
+ batch.column(i).dataType(),
+ arrowSchema(i).dataType,
+ ignoreNullability = true)
+ }
+ if (!matchingTypes) {
+ throw new UnsupportedOperationException(
+ "Comet's native Parquet writer cannot report row statistics for a
batch whose " +
+ "types differ from the write schema. Set
spark.comet.parquet.write.enabled=false " +
Review Comment:
This needs a merge with main, and the conflict lands in `recordRows`. #6494
renamed `spark.comet.parquet.write.enabled` to
`spark.comet.write.parquet.enabled` outright with no alias, so after the merge
this message would point users at a key that no longer does anything. The
conflicting hunk is main's version of the old warning, which this PR deletes,
so taking your side resolves it cleanly and leaves the stale key here without
any visible conflict. Could you build the message from
`CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key` instead, so the next rename
can't strand it?
Otherwise this looks good to me. I ran the writer suite with the PR merged
into current main on Spark 4.1. I also ran the same recording tracker through
Spark's own `WriteFilesExec` and through `CometWriteFilesExec` on 4.1 and 4.2,
and the trackers saw identical rows across about 20 column types, including
NaN, `-0.0`, timestamps before 1970, binary, intervals and nested types.
--
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]
