Re: [PR] fix: report row contents to native write statistics trackers [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org