[SPARK-51739][PYTHON] Validate Arrow schema from mapInArrow & mapInPandas & DataSource

[wengh]

What changes were proposed in this pull request?

Check the actual Arrow batch schema against the declared schema in MapInBatchEvaluator, throwing error if they don't match.

Also fix Pandas to Arrow conversion in ArrowStreamPandasUDFSerializer to respect nullability of output schema fields.

Why are the changes needed?

To improve error message and reject suspicious usage.

Does this PR introduce any user-facing change?

Yes.

Behaviour change

1. Some previously suspicious but accepted schema mismatches are now no longer valid.

   This includes:

   • extraneous fields (previously ignored)
   • wrong order of fields of the same type (previously accepted but in wrong order)
   • expected non-nullable field is actually nullable (previously ignored)

   Example:

   from pyspark.sql.datasource import DataSource, DataSourceReader
   from pyspark.sql.pandas.types import to_arrow_schema
   import pyarrow as pa
   
   expected = StructType.fromDDL("a int, b int")
   actual = StructType.fromDDL("b int, a int")  # wrong order of fields
   
   class TestDataSource(DataSource):
       def schema(self):
           return expected
       def reader(self, schema):
           return TestReader()
   
   class TestReader(DataSourceReader):
       def read(self, partition):
           schema = to_arrow_schema(actual)
           yield pa.record_batch([[1], [2]], schema=schema)
   
   spark.dataSource.register(TestDataSource)
   spark.read.format("TestDataSource").load().show()

   Before:

   +---+---+
   |  a|  b|
   +---+---+
   |  1|  2|
   +---+---+
   

   Now:

   org.apache.spark.SparkException: [ARROW_TYPE_MISMATCH] Invalid schema from pandas_udf(): expected StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)), got StructType(StructField(b,LongType,true),StructField(a,LongType,true)). SQLSTATE: 42K0G
   

2. For other schema mismatches, the error changed from internal error to a clearer ARROW_TYPE_MISMATCH error.

   This includes

   • wrong field types
   • less than expected number of fields

   Example:

   from pyspark.sql.pandas.types import to_arrow_schema
   from pyspark.sql.types import StructType, StructField, IntegerType
   import pyarrow as pa
   
   expected = StructType([StructField("a", IntegerType()), StructField("b", IntegerType())])
   actual = StructType([StructField("a", IntegerType())])  # missing a column
   
   def fun(iterator):
       for batch in iterator:
           schema = to_arrow_schema(actual)
           yield pa.record_batch([[1]], schema=schema)
   
   spark.range(2).mapInArrow(fun, expected).show()

   Before:

   java.lang.ArrayIndexOutOfBoundsException: Index 1 out of bounds for length 1
       at org.apache.spark.sql.vectorized.ArrowColumnVector.getChild(ArrowColumnVector.java:134)
       at org.apache.spark.sql.execution.python.MapInBatchEvaluatorFactory$MapInBatchEvaluator.$anonfun$eval$3(MapInBatchEvaluatorFactory.scala:82)
       ...
   

   Now:

   org.apache.spark.SparkException: [ARROW_TYPE_MISMATCH] Invalid schema from pandas_udf(): expected StructType(StructField(a,IntegerType,true),StructField(b,IntegerType,true)), got StructType(StructField(a,IntegerType,true)). SQLSTATE: 42K0G
   

How was this patch tested?

End-to-end tests in python/pyspark/sql/tests/arrow/test_arrow_map.py

Was this patch authored or co-authored using generative AI tooling?

No

[ueshin]

@HyukjinKwon Should ArrowEvalPythonExec also ignore nullability, and other similar Exec if any? Or this also should NOT ignore it?

[ueshin]

ah, but if this doesn't ignore nullability, it could introduce a breaking change?

[HyukjinKwon]

Hmmm .. yeah .. maybe let's just ignore it for now ..

[wengh]

Yeah, not ignoring nullability breaks many of existing tests (e.g. where expected field is not nullable but actual field is nullable but there's no null values, so the actual schema is technically wrong but not causing any problems)

[ueshin]

Otherwise, LGTM.

[HyukjinKwon]

test failure seems unrelated but mind triggering again to make sure?

[zhengruifeng]

qq: do we have corresponding validation in the python side?

[wengh]

do we have corresponding validation in the python side?

@zhengruifeng
No. DataSource only checks that top level columns have matching names. mapInArrow & mapInPandas don't check at all.

[wengh]

correctly handle non nullable fields required by the arrow_struct_type schema

[HyukjinKwon]

Merged to master.

[allisonwang-db]

Thanks for the fix! But this is a breaking change. Can we document this in the migration guide?

[wengh]

Thanks for the fix! But this is a breaking change. Can we document this in the migration guide?

@allisonwang-db #50722