Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]

["Yicong-Huang (via GitHub)" <gi...@apache.org>]

Yicong-Huang commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3454230793


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -4777,6 +4777,25 @@ object SQLConf {
         "must be positive.")
       .createWithDefault(100)
 
+  val PYTHON_UDF_MAX_BYTES_PER_BATCH =
+    buildConf("spark.sql.execution.python.udf.maxBytesPerBatch")
+      .internal()
+      .doc("Byte-size cap for a batch of rows sent to a worker for regular (pickle-serialized, " +
+        "non-Arrow) Python UDF evaluation. Without it, BatchEvalPythonExec batches purely by " +
+        "row count (spark.sql.execution.python.udf.maxRecordsPerBatch), so one oversized batch " +
+        "can OOM the executor; a finite value caps the batch at the min of the row-count and " +
+        "byte limits. The size is a best-effort per-row estimate of the pickled size of the " +
+        "converted row (its accuracy is observable via the pythonEstimatedInputBytes vs " +
+        "pythonDataSent metrics); a row larger than the cap still yields a one-row batch. " +
+        "-1 (the default) means no limit.")
+      .version("5.0.0")

Review Comment:
   should this be 4.3.0?



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -4777,6 +4777,25 @@ object SQLConf {
         "must be positive.")
       .createWithDefault(100)
 
+  val PYTHON_UDF_MAX_BYTES_PER_BATCH =
+    buildConf("spark.sql.execution.python.udf.maxBytesPerBatch")
+      .internal()
+      .doc("Byte-size cap for a batch of rows sent to a worker for regular (pickle-serialized, " +
+        "non-Arrow) Python UDF evaluation. Without it, BatchEvalPythonExec batches purely by " +
+        "row count (spark.sql.execution.python.udf.maxRecordsPerBatch), so one oversized batch " +
+        "can OOM the executor; a finite value caps the batch at the min of the row-count and " +
+        "byte limits. The size is a best-effort per-row estimate of the pickled size of the " +
+        "converted row (its accuracy is observable via the pythonEstimatedInputBytes vs " +
+        "pythonDataSent metrics); a row larger than the cap still yields a one-row batch. " +
+        "-1 (the default) means no limit.")
+      .version("5.0.0")
+      .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+      .bytesConf(ByteUnit.BYTE)
+      .checkValue(x => x == -1 || (x > 0 && x <= Int.MaxValue),

Review Comment:
   Int max is about 2GB, just want to make sure this is the idea range. as a reference, arrow (I know we are talking about non-arrow path) also has limit around 2GB, so maybe the lower limit is enough. but do we want to shrink the batch size cap for our batches?



##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -52,6 +54,183 @@ class BatchEvalPythonExecSuite extends SharedSparkSession
     }
   }
 
+  test("SPARK-57593: getInputIterator records peak pickled batch size as a size metric") {
+    EvaluatePython.registerPicklers()
+    val schema = StructType(Seq(StructField("s", StringType)))
+    // batchSize = 1 so each row is its own pickled batch; the widest row yields the peak.
+    val batchSize = 1
+    val rows = Seq(
+      InternalRow(UTF8String.fromString("a")),
+      InternalRow(UTF8String.fromString("b" * 1000)),
+      InternalRow(UTF8String.fromString("c")))
+
+    // Baseline: the peak pickled-batch size, computed without instrumentation.
+    val expectedPeak = BatchEvalPythonExec
+      .getInputIterator(rows.iterator, schema, batchSize, binaryAsBytes = false)
+      .map(_.length.toLong)
+      .max
+    assert(expectedPeak > 0)
+
+    val metric = SQLMetrics.createSizeMetric(spark.sparkContext, "peak pickled batch size in bytes")
+    // Drain the iterator so every batch is pickled and measured.
+    BatchEvalPythonExec
+      .getInputIterator(rows.iterator, schema, batchSize, binaryAsBytes = false,
+        pythonMetrics = Map("pythonPeakPickledBatchBytes" -> metric))
+      .foreach(_ => ())
+
+    assert(metric.value === expectedPeak)
+  }
+
+  test("SPARK-57593: ByteBoundedAsArrayIterator oversized-batch and estimated-bytes metrics") {
+    def elems(sizes: Long*): Iterator[(Any, Long)] =
+      sizes.iterator.map(s => (0.asInstanceOf[Any], s))
+    def m(): SQLMetric = SQLMetrics.createMetric(spark.sparkContext, "m")
+
+    // Cut once accumulated bytes >= 25 (after 3 rows of 10). estimatedInputBytes sums the
+    // per-row estimate across the whole partition.
+    val ov1 = m(); val est1 = m()
+    val byBytes = new ByteBoundedAsArrayIterator(
+      elems(10, 10, 10, 10, 10, 10, 10), maxRecordsPerBatch = 100, maxBytesPerBatch = 25L,
+      Some(ov1), Some(est1)).toList.map(_.length)
+    assert(byBytes === List(3, 3, 1))
+    assert(ov1.value === 2)   // two batches were cut at the cap; the trailing 1-row batch was not
+    assert(est1.value === 70) // sum of per-row estimates
+
+    // Record cap dominates when smaller; the byte limit is never reached.
+    val ov2 = m()
+    val byCount = new ByteBoundedAsArrayIterator(
+      elems(10, 10, 10, 10, 10), maxRecordsPerBatch = 2, maxBytesPerBatch = 1000L,
+      Some(ov2), None).toList.map(_.length)
+    assert(byCount === List(2, 2, 1))
+    assert(ov2.value === 0)
+
+    // A single row larger than the cap still forms a one-row batch (never zero rows).
+    val ov3 = m()
+    val giant = new ByteBoundedAsArrayIterator(
+      elems(10000), maxRecordsPerBatch = 100, maxBytesPerBatch = 25L,
+      Some(ov3), None).toList.map(_.length)
+    assert(giant === List(1))
+    assert(ov3.value === 1)
+
+    // Size-0 estimates contribute nothing to the estimate sum or the byte cap; the cut happens
+    // on the first row that pushes the estimate to the cap.
+    val ov5 = m(); val est5 = m()
+    val mixed = new ByteBoundedAsArrayIterator(
+      elems(0, 0, 30, 0), maxRecordsPerBatch = 100, maxBytesPerBatch = 25L,
+      Some(ov5), Some(est5)).toList.map(_.length)
+    assert(mixed === List(3, 1))
+    assert(est5.value === 30) // only the size-30 row contributes
+    assert(ov5.value === 1)   // one batch cut at the cap
+  }
+
+  test("SPARK-57593: ByteBoundedAsArrayIterator honors the Iterator contract") {
+    def elems(sizes: Long*): Iterator[(Any, Long)] =
+      sizes.iterator.map(s => (0.asInstanceOf[Any], s))
+    def iterOf(input: Iterator[(Any, Long)], maxRecords: Int = 100): Iterator[Array[Any]] =
+      new ByteBoundedAsArrayIterator(
+        input, maxRecordsPerBatch = maxRecords, maxBytesPerBatch = Int.MaxValue.toLong,
+        None, None)
+
+    // next() past the end throws NoSuchElementException, matching the row-count batching path.
+    val it = iterOf(elems(10))
+    assert(it.hasNext)
+    assert(it.next().length === 1)
+    assert(!it.hasNext)
+    intercept[NoSuchElementException](it.next())
+
+    // An empty input yields no batches and next() throws immediately.
+    val empty = iterOf(elems())
+    assert(!empty.hasNext)
+    intercept[NoSuchElementException](empty.next())
+
+    // A non-positive record limit is rejected up front.
+    intercept[IllegalArgumentException](iterOf(elems(10), maxRecords = 0))
+
+    // A non-positive byte cap is rejected up front too: only finite positive caps may reach this
+    // class (a cap of 0 would otherwise silently degrade every batch to a single row).
+    intercept[IllegalArgumentException](new ByteBoundedAsArrayIterator(
+      elems(10), maxRecordsPerBatch = 100, maxBytesPerBatch = 0L, None, None))
+  }
+
+  test("SPARK-57593: getInputIterator byte-bounds pickle batches when a cap is configured") {
+    EvaluatePython.registerPicklers()
+    val schema = StructType(Seq(StructField("s", StringType)))
+    // GenericInternalRow on purpose: production feeds this path from a MutableProjection whose
+    // target is a GenericInternalRow (never an UnsafeRow), so the byte cap must work on it.
+    def wideRows(): Iterator[InternalRow] = (0 until 10).iterator.map { _ =>
+      new GenericInternalRow(Array[Any](UTF8String.fromString("x" * 1000)))
+    }
+
+    // Default: no byte cap, so all 10 rows form a single batch (batchSize = 100).
+    val defaultBatches = BatchEvalPythonExec
+      .getInputIterator(wideRows(), schema, batchSize = 100, binaryAsBytes = false).size
+    assert(defaultBatches === 1)
+
+    // Finite cap: the wide rows split into multiple smaller batches, and the estimate metric is
+    // populated (regression test: the estimate must be non-zero for production row classes).
+    val estMetric = SQLMetrics.createSizeMetric(spark.sparkContext, "est")
+    val cappedBatches = BatchEvalPythonExec
+      .getInputIterator(wideRows(), schema, batchSize = 100, binaryAsBytes = false,
+        maxBytesPerBatch = 2048L,
+        pythonMetrics = Map("pythonEstimatedInputBytes" -> estMetric)).size
+    assert(cappedBatches > defaultBatches)
+    assert(estMetric.value >= 10 * 1000L) // 10 rows of a 1000-char string, plus overhead
+  }
+
+  test("SPARK-57593: toJava accumulates the pickled-size estimate during conversion") {
+    val acc = new PickledSizeAccumulator
+    def sized(value: Any, dataType: DataType): Long = {
+      EvaluatePython.toJava(value, dataType, binaryAsBytes = true, Some(acc))
+      acc.getAndReset()
+    }
+    // Payload bytes dominate the estimate; UTF8String.numBytes is the exact UTF-8 byte count.
+    val wide = sized(UTF8String.fromString("x" * 10000), StringType)
+    assert(wide >= 10000L && wide <= 10000L + 64L)
+    val binary = sized(Array.fill[Byte](5000)(1), BinaryType)
+    assert(binary >= 5000L && binary <= 5000L + 64L)
+    // Decimals expand into digit strings when pickled; the estimate tracks precision.
+    val dec = sized(Decimal(BigDecimal("1234567890.123456789")), DecimalType(38, 18))
+    assert(dec >= 19L)
+    // Nulls and unknown leaf types yield small positive estimates -- never zero, so the cap
+    // can never go blind on an unrecognized row shape.
+    assert(sized(null, StringType) > 0L)
+    assert(sized(new Object, CalendarIntervalType) > 0L)
+    // Whole-row conversion (needConversion path): nested values are sized in the same pass,
+    // and getAndReset() closes the per-row cycle.
+    val schema = StructType(Seq(StructField("s", StringType), StructField("n", StringType)))
+    val rowSize = sized(
+      new GenericInternalRow(Array[Any](UTF8String.fromString("y" * 2000), null)), schema)
+    assert(rowSize >= 2000L)
+    assert(acc.getAndReset() === 0L) // reset really resets
+  }
+
+  test("SPARK-57593: toJava size estimate is positive for every data type (drift guard)") {
+    // Enforces the accounting invariant on EvaluatePython.toJava: a (new) conversion case that
+    // forgets to account to sizeAcc yields a zero estimate at value level and makes the byte cap
+    // blind to that type. The type list comes from DataTypeTestUtils + composites, so the sweep
+    // extends as Spark grows types, and values come from the seeded RandomDataGenerator.
+    import org.apache.spark.sql.RandomDataGenerator
+    import org.apache.spark.sql.catalyst.CatalystTypeConverters
+    val acc = new PickledSizeAccumulator
+    val sweep = (DataTypeTestUtils.atomicTypes.toSeq ++ Seq(
+      ArrayType(StringType),
+      MapType(StringType, DoubleType),
+      StructType(Seq(StructField("s", StringType), StructField("d", DoubleType)))))
+    var swept = 0
+    sweep.foreach { dt =>
+      RandomDataGenerator.forType(dt, nullable = false, new scala.util.Random(42)).foreach {
+        gen =>
+          val catalystValue = CatalystTypeConverters.createToCatalystConverter(dt)(gen())
+          EvaluatePython.toJava(catalystValue, dt, binaryAsBytes = true, Some(acc))
+          val estimate = acc.getAndReset()
+          assert(estimate > 0L, s"toJava accounted nothing for $dt -- the byte cap is blind " +
+            "to this type; add accounting to the corresponding case")
+          swept += 1
+      }
+    }
+    assert(swept > 10) // the sweep must be exercising a meaningful slice of the type space
+  }
+

Review Comment:
   I might have missed it, but did we have a test for -1 value (no limit)? 



-- 
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: reviews-unsubscribe@spark.apache.org

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


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org