Schema pruning

spark_auto_mapper/automappers/complex.py

@@ -52,7 +52,7 @@ def __init__(
             column_name: str
             mapper: AutoMapperDataTypeBase
             for column_name, mapper in entity.get_child_mappers().items():
-                if column_name == "extension":
+                if not enable_schema_pruning and column_name == "extension":
                     extension_schema: Union[StructType, DataType, None]
                     # since there is a column called extension then get the schema with extension
                     extension_schema = mapper.get_schema(

spark_auto_mapper/automappers/with_column_base.py

@@ -11,6 +11,7 @@
 from spark_auto_mapper.automappers.automapper_base import AutoMapperBase
 from spark_auto_mapper.automappers.check_schema_result import CheckSchemaResult
 from spark_auto_mapper.data_types.data_type_base import AutoMapperDataTypeBase
+from spark_auto_mapper.schema_pruning.schema_pruner import SchemaPruner
 from spark_auto_mapper.type_definitions.defined_types import AutoMapperAnyDataType
 from spark_auto_mapper.helpers.value_parser import AutoMapperValueParser
 
@@ -53,11 +54,23 @@ def get_column_spec(self, source_df: Optional[DataFrame]) -> Column:
             child: AutoMapperDataTypeBase = self.value
             if self.column_schema:
                 if self.enable_schema_pruning:
-                    self.value.set_schema(
+                    self.value.mark_used_fields_in_schema(
                         column_name=self.dst_column,
                         column_path=self.dst_column,
+                        field=self.column_schema,
                         column_data_type=self.column_schema.dataType,
                     )
+                    # remove all fields without "used" property
+                    SchemaPruner.prune_schema(
+                        field=self.column_schema,
+                        field_data_type=self.column_schema.dataType,
+                    )
+
+                    # self.value.set_schema(
+                    #     column_name=self.dst_column,
+                    #     column_path=self.dst_column,
+                    #     column_data_type=self.column_schema.dataType,
+                    # )
             column_spec = child.get_column_spec(
                 source_df=source_df, current_column=None, parent_columns=None
             )
@@ -76,16 +89,16 @@ def get_column_spec(self, source_df: Optional[DataFrame]) -> Column:
             # if the type has a schema then apply it
             if self.column_schema:
                 column_data_type: DataType = self.column_schema.dataType
-                if self.enable_schema_pruning:
-                    # first disable generation of null properties since we are doing schema reduction
-                    self.value.include_null_properties(include_null_properties=False)
-                    # second ask the mapper to reduce schema that is not used
-                    column_data_type = self.value.filter_schema_by_fields_present(
-                        column_name=self.dst_column,
-                        column_path=self.dst_column,
-                        column_data_type=column_data_type,
-                        skip_null_properties=True,
-                    )
+                # if self.enable_schema_pruning:
+                #     # first disable generation of null properties since we are doing schema reduction
+                #     self.value.include_null_properties(include_null_properties=False)
+                #     # second ask the mapper to reduce schema that is not used
+                #     column_data_type = self.value.filter_schema_by_fields_present(
+                #         column_name=self.dst_column,
+                #         column_path=self.dst_column,
+                #         column_data_type=column_data_type,
+                #         skip_null_properties=True,
+                #     )
                 column_spec = column_spec.cast(column_data_type)
             # if dst_column already exists in source_df then prepend with ___ to make it unique
             if source_df is not None and self.dst_column in source_df.columns:
@@ -143,7 +156,8 @@ def check_schema(
                     desired_schema=desired_schema,
                 )
                 return CheckSchemaResult(result=result)
-            except AnalysisException:
+            except AnalysisException as e:
+                print(e)
                 return None
         else:
             return None

spark_auto_mapper/data_types/data_type_base.py

@@ -1024,6 +1024,152 @@ def set_schema(
             column_data_type=column_data_type,
         )
 
+    def _mark_used_fields_in_schema_for_array(
+        self,
+        *,
+        column_name: Optional[str],
+        column_path: Optional[str],
+        field: StructField,
+        column_data_type: DataType,
+    ) -> None:
+        assert isinstance(field, StructField)
+        assert isinstance(column_data_type, DataType)
+        assert isinstance(
+            column_data_type, ArrayType
+        ), f"{type(column_data_type)} should be ArrayType for {column_name} with path {column_path}"
+
+        element_type = column_data_type.elementType
+        # self.set_children_schema(element_type)
+        children: Union[
+            AutoMapperDataTypeBase, List[AutoMapperDataTypeBase]
+        ] = self.children
+        assert isinstance(children, list), f"{type(children)} should be a list"
+        if len(children) > 0:
+            child: AutoMapperDataTypeBase
+            for child in children:
+                child.mark_used_fields_in_schema(
+                    column_name=column_name,
+                    column_path=column_path,
+                    field=field,
+                    column_data_type=element_type,
+                )
+
+    # noinspection PyUnusedLocal
+    def _mark_used_fields_in_schema_for_struct(
+        self,
+        *,
+        column_name: Optional[str],
+        column_path: Optional[str],
+        field: StructField,
+        column_data_type: DataType,
+    ) -> None:
+        assert isinstance(field, StructField)
+        assert isinstance(column_data_type, DataType)
+
+        assert isinstance(
+            column_data_type, StructType
+        ), f"{type(column_data_type)} should be StructType for {column_name} with path {column_path}"
+
+        children: Union[
+            "AutoMapperDataTypeBase", List["AutoMapperDataTypeBase"]
+        ] = self.children
+        if isinstance(children, list) and len(children) > 0:
+            child: "AutoMapperDataTypeBase"
+            for index, child in enumerate(children):
+                assert isinstance(child, AutoMapperDataTypeBase), f"{type(child)}"
+                if not child.column_name:
+                    continue
+                assert child.column_name, f"No column name for {child}"
+                clean_child_name: str = PythonKeywordCleaner.from_python_safe(
+                    child.column_name
+                )
+                matching_fields = [
+                    f for f in column_data_type.fields if f.name == clean_child_name
+                ]
+                if len(matching_fields) == 0:
+                    pass
+                assert len(matching_fields) == 1, (
+                    f"Schema match failed for column {column_path}.{clean_child_name}"
+                    f" in schema fields"
+                    f": [{','.join([f.name for f in column_data_type.fields])}]"
+                )
+                child_field: StructField = matching_fields[0]
+                child.mark_used_fields_in_schema(
+                    column_name=child_field.name,
+                    column_path=f"{column_path}.{child_field.name}",
+                    field=child_field,
+                    column_data_type=child_field.dataType,
+                )
+        elif not isinstance(children, list) and children is not None:
+            child = children
+            assert child.column_name
+            clean_child_name = PythonKeywordCleaner.from_python_safe(child.column_name)
+            matching_fields = [
+                f for f in column_data_type.fields if f.name == clean_child_name
+            ]
+            assert len(matching_fields) == 1, (
+                f"Schema match failed for column {column_path}.{clean_child_name}"
+                f" in schema fields"
+                f": [{','.join([f.name for f in column_data_type.fields])}]"
+            )
+            child_field = matching_fields[0]
+            child.mark_used_fields_in_schema(
+                column_name=child_field.name,
+                column_path=f"{column_path}.{child_field.name}",
+                field=child_field,
+                column_data_type=child_field.dataType,
+            )
+
+    def mark_used_fields_in_schema(
+        self,
+        *,
+        column_name: Optional[str],
+        column_path: Optional[str],
+        field: StructField,
+        column_data_type: DataType,
+    ) -> None:
+        """
+        Sets the fields as used schema for this AutoMapper type
+
+
+        :param column_name: column name
+        :param column_path: full path to column
+        :param field: schema field for this mapper
+        :param column_data_type: schema for this mapper
+        """
+
+        assert isinstance(field, StructField)
+        assert isinstance(column_data_type, DataType)
+
+        # self.schema = column_data_type
+
+        # if this is a basic type so nothing to do
+        if not isinstance(column_data_type, StructType) and not isinstance(
+            column_data_type, ArrayType
+        ):
+            setattr(field, "used", True)
+            return
+
+        if isinstance(column_data_type, ArrayType):
+            self._mark_used_fields_in_schema_for_array(
+                column_name=column_name,
+                column_path=column_path,
+                field=field,
+                column_data_type=column_data_type,
+            )
+            return
+
+        assert isinstance(
+            column_data_type, StructType
+        ), f"{type(column_data_type)} should be StructType for {column_name} with path {column_path}"
+
+        self._mark_used_fields_in_schema_for_struct(
+            column_name=column_name,
+            column_path=column_path,
+            field=field,
+            column_data_type=column_data_type,
+        )
+
     @property
     @abstractmethod
     def children(

spark_auto_mapper/schema_pruning/annotated_struct_type.py

@@ -0,0 +1,12 @@
+from typing import List
+
+from pyspark.sql.types import StructType, StructField
+
+
+class AnnotatedStructType:
+    def __init__(self, fields: List[StructField]) -> None:
+        pass
+
+    @staticmethod
+    def parse(struct_type: StructType) -> "AnnotatedStructType":
+        pass

spark_auto_mapper/schema_pruning/schema_pruner.py

@@ -0,0 +1,41 @@
+from pyspark.sql.types import StructType, StructField, DataType, ArrayType
+
+
+class SchemaPruner:
+    @staticmethod
+    def _prune_schema_array(*, field: StructField, field_data_type: DataType) -> None:
+        assert isinstance(field_data_type, ArrayType)
+
+        SchemaPruner.prune_schema(
+            field=field, field_data_type=field_data_type.elementType
+        )
+
+    @staticmethod
+    def _prune_schema_struct(*, field: StructField, field_data_type: DataType) -> None:
+        assert isinstance(field_data_type, StructType)
+
+        # remove any fields that don't have the "used" field
+        field_data_type.fields = [
+            f for f in field_data_type.fields if hasattr(f, "used")
+        ]
+        field_data_type.names = [
+            n
+            for n in field_data_type.names
+            if n in [f.name for f in field_data_type.fields]
+        ]
+        # remove the used tag
+        for f in field_data_type.fields:
+            SchemaPruner.prune_schema(field=f, field_data_type=field_data_type)
+            delattr(f, "used")
+
+    @staticmethod
+    def prune_schema(*, field: StructField, field_data_type: DataType) -> None:
+        if isinstance(field.dataType, StructType):
+            SchemaPruner._prune_schema_struct(
+                field=field, field_data_type=field_data_type
+            )
+
+        if isinstance(field.dataType, ArrayType):
+            SchemaPruner._prune_schema_array(
+                field=field, field_data_type=field_data_type
+            )

tests/schema_pruner/test_schema_pruner.py

@@ -0,0 +1,17 @@
+from pyspark.sql.types import StructField, StructType, StringType, IntegerType
+
+from spark_auto_mapper.schema_pruning.schema_pruner import SchemaPruner
+
+
+def test_schema_pruner() -> None:
+    field: StructField = StructField(
+        "foo",
+        StructType(
+            [
+                StructField("prop1", StringType()),
+                StructField("prop2", IntegerType()),
+            ]
+        ),
+    )
+
+    SchemaPruner.prune_schema(field=field, field_data_type=field.dataType)

tests/schema_pruning/test_automapper_schema_pruning_with_defined_class.py

@@ -53,8 +53,7 @@ def test_auto_mapper_schema_pruning_with_defined_class(
 
     # Act
     mapper = AutoMapper(
-        view="members",
-        source_view="patients",
+        view="members", source_view="patients", enable_schema_pruning=True
     ).complex(MyClass(name=A.column("last_name"), age=A.number(A.column("my_age"))))
 
     assert isinstance(mapper, AutoMapper)