srielau commented on code in PR #58549:
URL: https://github.com/apache/spark/pull/58549#discussion_r3945717220
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala:
##########
@@ -207,6 +207,18 @@ object EvaluatePython {
case c: Int => c.toLong
}
+ case c: CharType => (obj: Any) => nullSafeConvert(obj) {
+ case _ =>
+ CharVarcharCodegenUtils.charTypeWriteSideCheck(
Review Comment:
Fixed in 5ca8ceae9f9. EvaluatePython now evaluates
shouldApplyWriteSideLengthCheck once and passes that policy through recursive
conversion. Both row and columnar evaluator projections use the same policy,
and legacy-mode scalar/nested UDF regression tests verify that CHAR remains
unpadded and over-length VARCHAR remains accepted.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala:
##########
@@ -557,13 +559,15 @@ private[sql] object ArrowConverters extends Logging {
val rdd = session.sparkContext
.parallelize(batchesInDriver.toImmutableArraySeq,
batchesInDriver.length)
.mapPartitions { batchesInExecutors =>
- ArrowConverters.fromBatchIterator(
+ val rows = ArrowConverters.fromBatchIterator(
batchesInExecutors,
schema,
timeZoneId,
errorOnDuplicatedFieldNames,
largeVarTypes,
TaskContext.get())
+ val projection = UnsafeProjection.create(checkedAttrs, attrs)
+ rows.map(row => projection(row).copy(): InternalRow)
Review Comment:
Fixed in 5ca8ceae9f9. The RDD path now preserves the original direct
fromBatchIterator result unless the schema contains CHAR/VARCHAR and the active
write-side policy requires checks. The projection/copy overhead is therefore
limited to schemas that need assignment enforcement; legacy mode also retains
the direct path.
##########
python/pyspark/sql/pandas/types.py:
##########
@@ -133,7 +135,7 @@ def to_arrow_type(
arrow_type = pa.float64()
elif isinstance(dt, DecimalType):
arrow_type = pa.decimal128(dt.precision, dt.scale)
- elif isinstance(dt, StringType):
+ elif isinstance(dt, (StringType, CharType, VarcharType)):
Review Comment:
Fixed in 5ca8ceae9f9. Return-type validation now recursively detects
CHAR/VARCHAR and keeps it unsupported for non-scalar map, grouped-map,
cogrouped-map, and aggregate pandas/Arrow eval types. Scalar Arrow-optimized,
pandas, and Arrow UDF eval types retain support. Added focused negative
coverage across those non-scalar API categories.
##########
python/pyspark/sql/tests/arrow/test_arrow_python_udf.py:
##########
@@ -270,18 +272,53 @@ def f(v: float):
rounded = df.select(f("v").alias("d")).first().d
self.assertEqual(rounded, Decimal("1.233999999999999986"))
- def test_err_return_type(self):
- with self.assertRaises(PySparkNotImplementedError) as pe:
- udf(lambda x: x, VarcharType(10), useArrow=True)
-
- self.check_error(
- exception=pe.exception,
- errorClass="NOT_IMPLEMENTED",
- messageParameters={
- "feature": "Invalid return type with Arrow-optimized Python
UDF: VarcharType(10)"
- },
+ def test_char_varchar_results(self):
+ schema = StructType(
+ [
+ StructField("c", CharType(4)),
+ StructField("v", VarcharType(3)),
+ StructField("nested", ArrayType(CharType(2))),
+ StructField("m", MapType(CharType(2), VarcharType(3))),
+ ]
)
+ with self.sql_conf(
+ {
+ "spark.sql.charVarchar.standardSemantics.enabled": "true",
+ "spark.sql.execution.arrow.pythonUDF.columnarInput.enabled":
"true",
+ }
+ ):
+ result = self.spark.range(1).select(
+ udf(
+ lambda _: ("ab", "xyz", ["z"], {"k": "xy"}),
+ schema,
+ useArrow=True,
+ )("id").alias("s")
+ )
+ self.assertEqual(
+ result.first().s,
+ Row(c="ab ", v="xyz", nested=["z "], m={"k ": "xy"}),
+ )
+
+ pandas_result = self.spark.range(1).select(
+ pandas_udf(lambda values: values,
CharType(4))(lit("ab")).alias("c")
+ )
+ self.assertEqual(pandas_result.first().c, "ab ")
+
+ invalid = self.spark.range(1).select(
+ udf(lambda _: "abcd", VarcharType(3), useArrow=True)("id")
+ )
+ with self.assertRaisesRegex(Exception, "EXCEED_LIMIT_LENGTH"):
+ invalid.collect()
+
+ with tempfile.TemporaryDirectory() as path:
+ self.spark.range(1).write.parquet(path)
+ columnar_input = self.spark.read.parquet(path)
Review Comment:
Fixed in 5ca8ceae9f9. Removed the Parquet case and added coverage to
ArrowColumnarPythonUDFSuite using readArrowSource. The new typed test UDF has a
real CHAR/VARCHAR PythonUDF output and a simple column input, asserts the
Arrow-backed columnar child is retained, and covers both CHAR padding and
VARCHAR overflow.
--
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]
