srielau commented on code in PR #58549:
URL: https://github.com/apache/spark/pull/58549#discussion_r4051149249
##########
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:
Only the UDT-backed case is incompatible with Connect; the nested
`ArrayType(CharType(2))` case passed the expected rejection in the previous
Connect CI run. Moving the whole method out of the shared mixin removes Connect
coverage for ordinary recursive CHAR/VARCHAR schemas. Please keep that case in
`BasePythonDataSourceTestsMixin` and move only the UDT-storage case to the
Classic suite.
##########
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:
This drops the pass-through columns that exposed the
`IndexOutOfBoundsException` in the previous CI run. It also changes plan
topology: as the suite's earlier comment explains, pruning these columns can
insert a non-columnar `ProjectExec` beneath `ArrowEvalPythonExec`, so these
cases no longer reliably exercise the intended Arrow-backed evaluator. Please
restore a pass-through-column regression and fix the underlying failure rather
than changing the query around it.
--
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]
