[SPARK-43055][CONNECT][PYTHON] Support duplicated nested field names

### What changes were proposed in this pull request?

Supports duplicated nested field names when `spark.createDataFrame` or `df.collect`.

### Why are the changes needed?

If there are duplicated nested field names, the following error is raised:

```py
>>> from pyspark.sql.types import *
>>>
>>> data = [Row(Row("a", 1), Row(2, 3, "b", 4, "c")), Row(Row("x", 6), Row(7, 8, "y", 9, "z"))]
>>> schema = (
...     StructType()
...     .add("struct", StructType().add("x", StringType()).add("x", IntegerType()))
...     .add(
...         "struct",
...         StructType()
...         .add("a", IntegerType())
...         .add("x", IntegerType())
...         .add("x", StringType())
...         .add("y", IntegerType())
...         .add("y", StringType()),
...     )
... )
>>> df = spark.createDataFrame(data, schema=schema)
Traceback (most recent call last):
...
pyarrow.lib.ArrowTypeError: Expected bytes, got a 'int' object
```

### Does this PR introduce _any_ user-facing change?

The duplicated nested field names will be available.

### How was this patch tested?

Added a test.

Closes apache#40692 from ueshin/issues/SPARK-43055/duplicate_fields.

Authored-by: Takuya UESHIN <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: ueshin

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala

@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, ExpressionEncode
 import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.UnboundRowEncoder
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder.Deserializer
 import org.apache.spark.sql.connect.client.util.{AutoCloseables, Cleanable}
+import org.apache.spark.sql.connect.common.DataTypeProtoConverter
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.ArrowUtils
 import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch, ColumnVector}
@@ -60,13 +61,20 @@ private[sql] class SparkResult[T](
   private def processResponses(stopOnFirstNonEmptyResponse: Boolean): Boolean = {
     while (responses.hasNext) {
       val response = responses.next()
-      if (response.hasArrowBatch) {
+      if (response.hasSchema) {
+        // The original schema should arrive before ArrowBatches.
+        structType =
+          DataTypeProtoConverter.toCatalystType(response.getSchema).asInstanceOf[StructType]
+      } else if (response.hasArrowBatch) {
         val ipcStreamBytes = response.getArrowBatch.getData
         val reader = new ArrowStreamReader(ipcStreamBytes.newInput(), allocator)
         try {
           val root = reader.getVectorSchemaRoot
           if (batches.isEmpty) {
-            structType = ArrowUtils.fromArrowSchema(root.getSchema)
+            if (structType == null) {
+              // If the schema is not available yet, fallback to the schema from Arrow.
+              structType = ArrowUtils.fromArrowSchema(root.getSchema)
+            }
             // TODO: create encoders that directly operate on arrow vectors.
             boundEncoder = createEncoder(structType).resolveAndBind(structType.toAttributes)
           }

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.connect.service
 
+import java.util.concurrent.atomic.AtomicInteger
+
 import scala.collection.JavaConverters._
 
 import com.google.protobuf.ByteString
@@ -38,7 +40,7 @@ import org.apache.spark.sql.connect.service.SparkConnectStreamHandler.processAsA
 import org.apache.spark.sql.execution.{SparkPlan, SQLExecution}
 import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, QueryStageExec}
 import org.apache.spark.sql.execution.arrow.ArrowConverters
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructField, StructType, UserDefinedType}
 import org.apache.spark.util.{ThreadUtils, Utils}
 
 class SparkConnectStreamHandler(responseObserver: StreamObserver[ExecutePlanResponse])
@@ -120,8 +122,38 @@ object SparkConnectStreamHandler {
       sessionId: String,
       dataframe: DataFrame,
       responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
+
+    def deduplicateFieldNames(dt: DataType): DataType = dt match {
+      case udt: UserDefinedType[_] => deduplicateFieldNames(udt.sqlType)
+      case st @ StructType(fields) =>
+        val newNames = if (st.names.toSet.size == st.names.length) {
+          st.names
+        } else {
+          val genNawName = st.names.groupBy(identity).map {
+            case (name, names) if names.length > 1 =>
+              val i = new AtomicInteger()
+              name -> { () => s"${name}_${i.getAndIncrement()}" }
+            case (name, _) => name -> { () => name }
+          }
+          st.names.map(genNawName(_)())
+        }
+        val newFields =
+          fields.zip(newNames).map { case (StructField(_, dataType, nullable, metadata), name) =>
+            StructField(name, deduplicateFieldNames(dataType), nullable, metadata)
+          }
+        StructType(newFields)
+      case ArrayType(elementType, containsNull) =>
+        ArrayType(deduplicateFieldNames(elementType), containsNull)
+      case MapType(keyType, valueType, valueContainsNull) =>
+        MapType(
+          deduplicateFieldNames(keyType),
+          deduplicateFieldNames(valueType),
+          valueContainsNull)
+      case _ => dt
+    }
+
     val spark = dataframe.sparkSession
-    val schema = dataframe.schema
+    val schema = deduplicateFieldNames(dataframe.schema).asInstanceOf[StructType]
     val maxRecordsPerBatch = spark.sessionState.conf.arrowMaxRecordsPerBatch
     val timeZoneId = spark.sessionState.conf.sessionLocalTimeZone
     // Conservatively sets it 70% because the size is not accurate but estimated.

python/pyspark/sql/connect/client.py

@@ -678,12 +678,13 @@ def to_pandas(self, plan: pb2.Plan) -> "pd.DataFrame":
         req.plan.CopyFrom(plan)
         table, schema, metrics, observed_metrics, _ = self._execute_and_fetch(req)
         assert table is not None
-        pdf = table.rename_columns([f"col_{i}" for i in range(len(table.column_names))]).to_pandas()
-        pdf.columns = table.column_names
 
         schema = schema or types.from_arrow_schema(table.schema)
         assert schema is not None and isinstance(schema, StructType)
 
+        pdf = table.to_pandas()
+        pdf.columns = schema.fieldNames()
+
         for field, pa_field in zip(schema, table.schema):
             if isinstance(field.dataType, TimestampType):
                 assert pa_field.type.tz is not None

python/pyspark/sql/connect/conversion.py

@@ -19,6 +19,7 @@
 check_dependencies(__name__)
 
 import array
+import itertools
 import datetime
 import decimal
 
@@ -31,20 +32,23 @@
     TimestampType,
     TimestampNTZType,
     MapType,
+    StructField,
     StructType,
     ArrayType,
     BinaryType,
     NullType,
     DecimalType,
     StringType,
     UserDefinedType,
+    cast,
 )
 
 from pyspark.sql.connect.types import to_arrow_schema
 
 from typing import (
     Any,
     Callable,
+    Dict,
     Sequence,
     List,
 )
@@ -99,10 +103,10 @@ def _create_converter(dataType: DataType) -> Callable:
 
             field_names = dataType.fieldNames()
 
-            field_convs = {
-                field.name: LocalDataToArrowConversion._create_converter(field.dataType)
+            field_convs = [
+                LocalDataToArrowConversion._create_converter(field.dataType)
                 for field in dataType.fields
-            }
+            ]
 
             def convert_struct(value: Any) -> Any:
                 if value is None:
@@ -113,24 +117,15 @@ def convert_struct(value: Any) -> Any:
                     ), f"{type(value)} {value}"
 
                     _dict = {}
-                    if isinstance(value, dict):
-                        for k, v in value.items():
-                            assert isinstance(k, str)
-                            _dict[k] = field_convs[k](v)
-                    elif isinstance(value, Row) and hasattr(value, "__fields__"):
-                        for k, v in value.asDict(recursive=False).items():
-                            assert isinstance(k, str)
-                            _dict[k] = field_convs[k](v)
-                    elif not isinstance(value, Row) and hasattr(value, "__dict__"):
-                        for k, v in value.__dict__.items():
-                            assert isinstance(k, str)
-                            _dict[k] = field_convs[k](v)
-                    else:
-                        i = 0
-                        for v in value:
-                            field_name = field_names[i]
-                            _dict[field_name] = field_convs[field_name](v)
-                            i += 1
+                    if not isinstance(value, Row) and hasattr(value, "__dict__"):
+                        value = value.__dict__
+                    for i, field in enumerate(field_names):
+                        if isinstance(value, dict):
+                            v = value.get(field)
+                        else:
+                            v = value[i]
+
+                        _dict[f"col_{i}"] = field_convs[i](v)
 
                     return _dict
 
@@ -255,8 +250,6 @@ def convert(data: Sequence[Any], schema: StructType) -> "pa.Table":
 
         assert schema is not None and isinstance(schema, StructType)
 
-        pa_schema = to_arrow_schema(schema)
-
         column_names = schema.fieldNames()
 
         column_convs = [
@@ -276,6 +269,27 @@ def convert(data: Sequence[Any], schema: StructType) -> "pa.Table":
 
                 pylist[i].append(column_convs[i](value))
 
+        def normalize(dt: DataType) -> DataType:
+            if isinstance(dt, StructType):
+                return StructType(
+                    [
+                        StructField(f"col_{i}", normalize(field.dataType), nullable=field.nullable)
+                        for i, field in enumerate(dt.fields)
+                    ]
+                )
+            elif isinstance(dt, ArrayType):
+                return ArrayType(normalize(dt.elementType), containsNull=dt.containsNull)
+            elif isinstance(dt, MapType):
+                return MapType(
+                    normalize(dt.keyType),
+                    normalize(dt.valueType),
+                    valueContainsNull=dt.valueContainsNull,
+                )
+            else:
+                return dt
+
+        pa_schema = to_arrow_schema(cast(StructType, normalize(schema)))
+
         return pa.Table.from_arrays(pylist, schema=pa_schema)
 
 
@@ -319,28 +333,42 @@ def _create_converter(dataType: DataType) -> Callable:
 
         elif isinstance(dataType, StructType):
 
-            field_convs = {
-                f.name: ArrowTableToRowsConversion._create_converter(f.dataType)
-                for f in dataType.fields
-            }
-            need_conv = any(
-                ArrowTableToRowsConversion._need_converter(f.dataType) for f in dataType.fields
-            )
+            field_names = dataType.names
+
+            if len(set(field_names)) == len(field_names):
+                dedup_field_names = field_names
+            else:
+                gen_new_name: Dict[str, Callable[[], str]] = {}
+                for name, group in itertools.groupby(dataType.names):
+                    if len(list(group)) > 1:
+
+                        def _gen(_name: str) -> Callable[[], str]:
+                            _i = itertools.count()
+                            return lambda: f"{_name}_{next(_i)}"
+
+                    else:
+
+                        def _gen(_name: str) -> Callable[[], str]:
+                            return lambda: _name
+
+                    gen_new_name[name] = _gen(name)
+                dedup_field_names = [gen_new_name[name]() for name in dataType.names]
+
+            field_convs = [
+                ArrowTableToRowsConversion._create_converter(f.dataType) for f in dataType.fields
+            ]
 
             def convert_struct(value: Any) -> Any:
                 if value is None:
                     return None
                 else:
                     assert isinstance(value, dict)
 
-                    if need_conv:
-                        _dict = {}
-                        for k, v in value.items():
-                            assert isinstance(k, str)
-                            _dict[k] = field_convs[k](v)
-                        return Row(**_dict)
-                    else:
-                        return Row(**value)
+                    _values = [
+                        field_convs[i](value.get(name, None))
+                        for i, name in enumerate(dedup_field_names)
+                    ]
+                    return _create_row(field_names, _values)
 
             return convert_struct
 
@@ -425,13 +453,10 @@ def convert(table: "pa.Table", schema: StructType) -> List[Row]:
             ArrowTableToRowsConversion._create_converter(f.dataType) for f in schema.fields
         ]
 
-        # table.to_pylist() automatically remove columns with duplicated names,
-        # to avoid this, use columnar lists here.
-        # TODO: support duplicated field names in the one struct. e.g. SF.struct("a", "a")
         columnar_data = [column.to_pylist() for column in table.columns]
 
         rows: List[Row] = []
         for i in range(0, table.num_rows):
-            values = [field_converters[j](columnar_data[j][i]) for j in range(0, table.num_columns)]
-            rows.append(_create_row(fields=table.column_names, values=values))
+            values = [field_converters[j](columnar_data[j][i]) for j in range(table.num_columns)]
+            rows.append(_create_row(fields=schema.fieldNames(), values=values))
         return rows

python/pyspark/sql/tests/test_dataframe.py

@@ -1700,6 +1700,26 @@ def test_where(self):
             message_parameters={"arg_name": "condition", "arg_type": "int"},
         )
 
+    def test_duplicate_field_names(self):
+        data = [Row(Row("a", 1), Row(2, 3, "b", 4, "c")), Row(Row("x", 6), Row(7, 8, "y", 9, "z"))]
+        schema = (
+            StructType()
+            .add("struct", StructType().add("x", StringType()).add("x", IntegerType()))
+            .add(
+                "struct",
+                StructType()
+                .add("a", IntegerType())
+                .add("x", IntegerType())
+                .add("x", StringType())
+                .add("y", IntegerType())
+                .add("y", StringType()),
+            )
+        )
+        df = self.spark.createDataFrame(data, schema=schema)
+
+        self.assertEqual(df.schema, schema)
+        self.assertEqual(df.collect(), data)
+
 
 class QueryExecutionListenerTests(unittest.TestCase, SQLTestUtils):
     # These tests are separate because it uses 'spark.sql.queryExecutionListeners' which is