update

Author: nanne-aben

tests/_utils/test_load_table.py

@@ -231,6 +231,9 @@ def test_create_schema_with_invalid_column_name(spark: SparkSession):
     df = spark.createDataFrame([("Alice", 24), ("Bob", 25)], ["first-name", "age"])
     ds, schema = create_schema(df)
 
+    df2 = ds.to_dataframe()
+    assert_df_equality(df, df2)
+
 
 def test_create_schema_with_invalid_column_name_in_a_structtype(spark: SparkSession):
     data = [
@@ -261,6 +264,9 @@ def test_create_schema_with_invalid_column_name_in_a_structtype(spark: SparkSess
     df = spark.createDataFrame(data)
     ds, schema = create_schema(df)
 
+    df2 = ds.to_dataframe()
+    assert_df_equality(df, df2)
+
 
 def test_create_schema_with_invalid_column_name_in_a_nested_structtype(spark: SparkSession):
     data = [
@@ -298,3 +304,6 @@ def test_create_schema_with_invalid_column_name_in_a_nested_structtype(spark: Sp
 
     df = spark.createDataFrame(data)
     ds, schema = create_schema(df)
+
+    df2 = ds.to_dataframe()
+    assert_df_equality(df, df2)

typedspark/_core/dataset.py

@@ -22,7 +22,7 @@
 from pyspark.sql import DataFrame
 from typing_extensions import Concatenate, ParamSpec
 
-from typedspark._core.rename_columns import rename_columns
+from typedspark._core.rename_columns import rename_columns, rename_columns_2
 from typedspark._core.validate_schema import validate_schema
 from typedspark._schema.schema import Schema
 from typedspark._transforms.transform_to_schema import transform_to_schema
@@ -91,7 +91,7 @@ class Person(Schema):
 
         schema = cls._schema_annotations  # type: ignore
 
-        df = rename_columns(df, schema.get_structtype())
+        df = rename_columns(df, schema)
         df = transform_to_schema(df, schema)
         if register_to_schema:
             schema = register_schema_to_dataset(df, schema)
@@ -114,14 +114,12 @@ class Person(Schema):
         df = cast(DataFrame, self)
         df.__class__ = DataFrame
 
-        for column in self._schema_annotations.get_structtype().fields:
-            if column.metadata:
-                df = df.withColumnRenamed(
-                    column.name, column.metadata.get("external_name", column.name)
-                )
+        df = rename_columns_2(df, self._schema_annotations)
 
         return df
 
+        # return rename_columns(df, schema)
+
     """The following functions are equivalent to their parents in ``DataFrame``, but
     since they don't affect the ``Schema``, we can add type annotations here.
 

typedspark/_core/rename_columns.py

@@ -1,17 +1,19 @@
 """Helper functions to rename columns from their external name (defined in
 `ColumnMeta(external_name=...)`) to their internal name."""
 
-from typing import Optional
+from typing import Optional, Type
 
 from pyspark.sql import Column, DataFrame
 from pyspark.sql.functions import col, lit, struct, when
 from pyspark.sql.types import StructField, StructType
 
+from typedspark._schema.schema import Schema
 
-def rename_columns(df: DataFrame, schema: StructType) -> DataFrame:
+
+def rename_columns(df: DataFrame, schema: Type[Schema]) -> DataFrame:
     """Helper functions to rename columns from their external name (defined in
-    `ColumnMeta(external_name=...)`) to their internal name."""
-    for field in schema.fields:
+    `ColumnMeta(external_name=...)`) to their internal name (as used in the Schema)."""
+    for field in schema.get_structtype().fields:
         internal_name = field.name
 
         if field.metadata and "external_name" in field.metadata:
@@ -25,6 +27,23 @@ def rename_columns(df: DataFrame, schema: StructType) -> DataFrame:
     return df
 
 
+def rename_columns_2(df: DataFrame, schema: Type[Schema]) -> DataFrame:
+    """Helper functions to rename columns from their internal name (as used in the
+    Schema) to their external name (defined in `ColumnMeta(external_name=...)`)."""
+    for field in schema.get_structtype().fields:
+        internal_name = field.name
+
+        if field.metadata and "external_name" in field.metadata:
+            external_name = field.metadata["external_name"]
+            df = df.withColumnRenamed(internal_name, external_name)  # swap
+
+        if isinstance(field.dataType, StructType):
+            structtype = _create_renamed_structtype_2(field.dataType, internal_name)
+            df = df.withColumn(external_name, structtype)  # swap
+
+    return df
+
+
 def _create_renamed_structtype(
     schema: StructType,
     parent: str,
@@ -35,7 +54,7 @@ def _create_renamed_structtype(
 
     mapping = []
     for field in schema.fields:
-        external_name = _get_external_name(field, full_parent_path)
+        external_name = _get_updated_parent_path(full_parent_path, field)
 
         if isinstance(field.dataType, StructType):
             mapping += [
@@ -51,11 +70,44 @@ def _create_renamed_structtype(
     return _produce_nested_structtype(mapping, parent, full_parent_path)
 
 
-def _get_external_name(field: StructField, full_parent_path: str) -> str:
+def _create_renamed_structtype_2(
+    schema: StructType,
+    parent: str,
+    full_parent_path: Optional[str] = None,
+) -> Column:
+    if not full_parent_path:
+        full_parent_path = f"`{parent}`"
+
+    mapping = []
+    for field in schema.fields:
+        internal_name = field.name
+        external_name = field.metadata.get("external_name", internal_name)
+
+        updated_parent_path = _get_updated_parent_path_2(full_parent_path, internal_name)  # swap
+
+        if isinstance(field.dataType, StructType):
+            mapping += [
+                _create_renamed_structtype(
+                    field.dataType,
+                    parent=external_name,  # swap
+                    full_parent_path=updated_parent_path,
+                )
+            ]
+        else:
+            mapping += [col(updated_parent_path).alias(external_name)]  # swap
+
+    return _produce_nested_structtype(mapping, parent, full_parent_path)
+
+
+def _get_updated_parent_path(full_parent_path: str, field: StructField) -> str:
     external_name = field.metadata.get("external_name", field.name)
     return f"{full_parent_path}.`{external_name}`"
 
 
+def _get_updated_parent_path_2(full_parent_path: str, field: str) -> str:
+    return f"{full_parent_path}.`{field}`"
+
+
 def _produce_nested_structtype(
     mapping: list[Column],
     parent: str,