srielau commented on code in PR #58549:
URL: https://github.com/apache/spark/pull/58549#discussion_r4051205217
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -103,6 +115,231 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("Arrow-backed source: CHAR/VARCHAR output checks") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") {
+ val charUDF = TestTypedScalarPandasUDF(
+ name = "arrow_char_udf", returnType = CharType(4))
+ val varcharUDF = TestTypedScalarPandasUDF(
+ name = "arrow_varchar_udf", returnType = VarcharType(3))
+ registerTestUDF(charUDF, spark)
+ registerTestUDF(varcharUDF, spark)
+
+ val df = readArrowSource(numRows = 10)
+ val padded = df.selectExpr(
+ "id", "name", "value", "data",
+ "arrow_char_udf(id) as udf_id")
+ val arrowExec = collectNodes[ArrowEvalPythonExec](
+ padded.queryExecution.executedPlan).head
+ assert(arrowExec.child.supportsColumnar,
+ "ArrowEvalPythonExec should retain its Arrow-backed columnar child")
+ assert(padded.collect().map(_.getString(4)).toSeq ===
+ (0 until 10).map(_.toString.padTo(4, ' ').mkString))
+
+ val exception = intercept[SparkRuntimeException] {
+ df.selectExpr(
+ "id", "name", "value", "data",
+ "arrow_varchar_udf(name) as udf_name").collect()
+ }
+ assert(exception.getMessage.contains("EXCEED_LIMIT_LENGTH"))
+ }
+ }
+
+ test("Arrow-backed source: row queue preserves ordinary UDT sibling") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") {
+ val charUDF = TestTypedScalarPandasUDF(
+ name = "row_queue_char_udf", returnType = CharType(4))
+ val udtUDF = TestTypedScalarPandasUDF(
+ name = "row_queue_udt_udf", returnType = new StringStorageUDT())
+ registerTestUDF(charUDF, spark)
+ registerTestUDF(udtUDF, spark)
+
+ val result = readArrowSource(numRows = 3).selectExpr(
+ "row_queue_char_udf(id) as char_result",
+ "row_queue_udt_udf(id) as udt_result")
+ val arrowExec =
collectNodes[ArrowEvalPythonExec](result.queryExecution.executedPlan).head
+
+ assert(!ColumnarArrowEvalPythonEvaluatorFactory.canUseArrowColumnar(
+ Some(Array(0, 0)), isArrow = true, arrowExec.udfs))
+
assert(result.schema("udt_result").dataType.isInstanceOf[StringStorageUDT])
+ assert(result.collect().map(row => (row.getString(0),
row.getString(1))).toSeq ===
+ Seq(("0 ", "0"), ("1 ", "1"), ("2 ", "2")))
+ }
+ }
+
+ test("Arrow-backed source: default policy exposes checked CHAR/VARCHAR
output as STRING") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true",
+ SQLConf.LEGACY_CHAR_VARCHAR_AS_STRING.key -> "false",
+ SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "false",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false") {
+ val charUDF = TestTypedScalarPandasUDF(
+ name = "default_arrow_char_udf", returnType = CharType(4))
+ registerTestUDF(charUDF, spark)
+
+ val result = readArrowSource(numRows = 1)
+ .selectExpr("default_arrow_char_udf(id) as udf_id")
+
+ assert(result.schema.head.dataType === StringType)
+ assert(result.head().getString(0) === "0 ")
+ }
+ }
+
+ test("Arrow-backed source: higher-order CHAR/VARCHAR output checks") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") {
+ val charUDF = TestTypedScalarPandasUDF(
+ name = "arrow_hof_char_udf", returnType = CharType(4))
+ val varcharUDF = TestTypedScalarPandasUDF(
+ name = "arrow_hof_varchar_udf", returnType = VarcharType(3))
+ registerTestUDF(charUDF, spark)
+ registerTestUDF(varcharUDF, spark)
+
+ val df = readArrowSource(numRows = 10)
+ val padded = df.selectExpr(
+ "id", "name", "value", "data",
+ "transform(array(id), x -> arrow_hof_char_udf(x)) as udf_values")
+ val arrowExec = collectNodes[ArrowEvalPythonExec](
+ padded.queryExecution.executedPlan).head
+ assert(arrowExec.child.supportsColumnar,
+ "ArrowEvalPythonExec should retain its Arrow-backed columnar child")
+ assert(padded.collect().map(_.getSeq[String](4)).toSeq ===
+ (0 until 10).map(index => Seq(index.toString.padTo(4, ' ').mkString)))
+
+ val exception = intercept[SparkRuntimeException] {
+ df.selectExpr(
+ "id", "name", "value", "data",
+ "transform(array(name), x -> arrow_hof_varchar_udf(x)) as
udf_values").collect()
+ }
+ assert(exception.getMessage.contains("EXCEED_LIMIT_LENGTH"))
+ }
+ }
+
+ test("Arrow-backed source: nested checked output keeps unchecked sibling
ordinal") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") {
+ val charUDF = TestTypedScalarPandasUDF(
+ name = "mixed_nested_char_udf", returnType = CharType(4))
+ val varcharUDF = TestTypedScalarPandasUDF(
+ name = "mixed_unchecked_varchar_udf", returnType = VarcharType(3))
+ registerTestUDF(charUDF, spark)
+ registerTestUDF(varcharUDF, spark)
+
+ val checkedNested =
+ transform(array(col("id")), value =>
charUDF(value)).as("checked_nested")
+ val unchecked = withSQLConf(
+ SQLConf.LEGACY_CHAR_VARCHAR_AS_STRING.key -> "true",
+ SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "false",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false") {
+ varcharUDF(col("name")).as("unchecked")
+ }
+ val result = readArrowSource(numRows = 2).select(
+ col("id"),
+ col("name"),
+ col("value"),
+ col("data"),
+ checkedNested,
+ unchecked,
+ varcharUDF(col("id")).as("checked"))
+ val arrowExec =
collectNodes[ArrowEvalPythonExec](result.queryExecution.executedPlan).head
+
+ assert(arrowExec.child.supportsColumnar)
+ assert(result.schema("checked_nested").dataType ===
ArrayType(CharType(4)))
+ assert(result.schema("unchecked").dataType === StringType)
+ assert(result.schema("checked").dataType === VarcharType(3))
+ assert(result.collect()
+ .map(row => (row.getSeq[String](4), row.getString(5),
row.getString(6))).toSeq ===
+ Seq((Seq("0 "), "row_0", "0"), (Seq("1 "), "row_1", "1")))
+
+ val exception = intercept[SparkRuntimeException] {
+ readArrowSource(numRows = 1).select(
+ col("id"),
+ col("name"),
+ col("value"),
+ col("data"),
+ checkedNested,
+ unchecked,
+ varcharUDF(col("name")).as("checked")).collect()
+ }
+ assert(exception.getMessage.contains("EXCEED_LIMIT_LENGTH"))
+ }
+ }
+
+ test("Arrow-backed source: legacy CHAR/VARCHAR output remains unchecked") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED.key -> "true",
+ SQLConf.LEGACY_CHAR_VARCHAR_AS_STRING.key -> "true",
+ SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "false",
+ SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false") {
+ val charUDF = TestTypedScalarPandasUDF(
+ name = "legacy_arrow_char_udf", returnType = CharType(4))
+ val varcharUDF = TestTypedScalarPandasUDF(
+ name = "legacy_arrow_varchar_udf", returnType = VarcharType(3))
+ registerTestUDF(charUDF, spark)
+ registerTestUDF(varcharUDF, spark)
+
+ val result = readArrowSource(numRows = 10).selectExpr(
Review Comment:
Fixed in `105258cd524`. I restored the nullable pass-through columns and
fixed the underlying lifetime bug in the optimized evaluator. The input reader
can close source-owned Arrow vectors once the Python runner consumes its input,
so the pass-through queue now stores independent `splitAndTransfer` views and
closes any queued views on task completion. The original topology and
regression assertions are preserved. The focused Scala suite compiles locally;
pandas-dependent cases require CI on this host.
##########
python/pyspark/sql/tests/test_python_datasource.py:
##########
@@ -2024,7 +2034,43 @@ def test_data_source_read_with_udf_perf_profiler(self):
self.assertEqual(stdout, "")
-class PythonDataSourceTests(BasePythonDataSourceTestsMixin,
ReusedSQLTestCase): ...
+class PythonDataSourceTests(BasePythonDataSourceTestsMixin, ReusedSQLTestCase):
+ def test_data_source_char_varchar_return_type_is_unsupported(self):
+ class CharStorageUDT(UserDefinedType):
+ @classmethod
+ def sqlType(cls):
+ return CharType(2)
+
+ @classmethod
+ def module(cls):
+ return __name__
+
+ @classmethod
+ def scalaUDT(cls):
+ return ""
+
+ def serialize(self, obj):
+ return obj
+
+ def deserialize(self, datum):
+ return datum
+
+ schemas = [
Review Comment:
Fixed in `105258cd524`. The nested `ArrayType(CharType(2))` rejection is
back in `BasePythonDataSourceTestsMixin`, so Connect continues to run it. Only
the UDT-storage case is now a separate Classic-only test. A local ownership
check confirms the shared case is inherited by `PythonDataSourceParityTests`
while the UDT case is not.
--
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]
