feat: add orphan and group rejections to spark backend

src/dve/core_engine/backends/base/reader.py

@@ -11,7 +11,7 @@
 from dve.core_engine.backends.exceptions import (
     CriticalMessageBearingError,
     MessageBearingError,
-    ReaderLacksEntityTypeSupport
+    ReaderLacksEntityTypeSupport,
 )
 from dve.core_engine.backends.types import EntityName, EntityType
 from dve.core_engine.configuration.v1 import (

src/dve/core_engine/backends/base/rules.py

@@ -427,7 +427,13 @@ def process_node(node: HierarchyNode) -> bool:
 
             self.cache_entity(node.entity_name, entities)
 
-            return len(_messages) > 0
+            _orph_count = len(_messages)
+
+            self.logger.info(
+                f"Found {_orph_count} orphan records between {node.entity_name} and {node.parent_entity}"  # pylint: disable=C0301
+            )
+
+            return _orph_count > 0
 
         entity_issues_found: dict[EntityName, bool] = {}
 
@@ -454,7 +460,7 @@ def process_node(node: HierarchyNode) -> bool:
                 return False
 
             self.logger.info(
-                f"Identifying that mandatory entity `{node.parent_entity}` has at least 1 valid child record"  # pylint: disable=C0301
+                f"Identifying that mandatory entity `{node.parent_entity}` has at least 1 valid child record in {node.entity_name}"  # pylint: disable=C0301
             )
 
             join_expr = " AND ".join(
@@ -493,7 +499,14 @@ def process_node(node: HierarchyNode) -> bool:
                 ]
                 msg_writer.write_queue.put(_messages)
             self.cache_entity(node.parent_entity, entities)
-            return len(_messages) > 0
+
+            _no_valid_child_records: int = len(_messages)
+
+            self.logger.info(
+                f"Found {_no_valid_child_records} records with no valid children in {node.parent_entity}."  # pylint: disable=C0301
+            )
+
+            return _no_valid_child_records > 0
 
         entity_issues_found: dict[EntityName, bool] = {}
 

src/dve/core_engine/backends/exceptions.py

@@ -32,6 +32,7 @@ def __init__(self, *args: object, messages: Messages) -> None:
         self.messages = messages
         """The messages to be returned as part of the error."""
 
+
 class CriticalMessageBearingError(BackendError):
     """
     A backend error that comes with a pre-created message.
@@ -44,6 +45,7 @@ def __init__(self, *args: object, message: FeedbackMessage) -> None:
         self.message = message
         """The message to be returned as part of the error."""
 
+
 class UnableToParseCSVError(CriticalMessageBearingError):
     """An error raised when unable to parse a CSV file"""
 
@@ -60,7 +62,8 @@ def __init__(
                 failure_type="submission",
                 is_informational=False,
                 error_type="csv read",
-                error_message=error_message or "Unable to parse the CSV file. Please check the structure of your CSV.",  # pylint: disable=C0301
+                error_message=error_message
+                or "Unable to parse the CSV file. Please check the structure of your CSV.",  # pylint: disable=C0301
                 error_code=error_code or "MalformedCSV",
             )
         )

src/dve/core_engine/backends/implementations/duckdb/readers/csv.py

@@ -123,7 +123,8 @@ def read_to_relation(  # pylint: disable=unused-argument
             raise UnableToParseCSVError(
                 entity_name="csv_structure",
                 error_code=self.ft_error_code,
-                error_message=self.ft_error_message or "Unable to parse CSV file. Structure is likely malformed.",  # pylint: disable=C0301
+                error_message=self.ft_error_message
+                or "Unable to parse CSV file. Structure is likely malformed.",  # pylint: disable=C0301
             ) from exc
 
         if self.null_empty_strings:
@@ -183,7 +184,8 @@ def read_to_relation(  # pylint: disable=unused-argument
             raise UnableToParseCSVError(
                 entity_name="csv_structure",
                 error_code=self.ft_error_code,
-                error_message=self.ft_error_message or "Unable to parse CSV file. Structure is likely malformed.",  # pylint: disable=C0301
+                error_message=self.ft_error_message
+                or "Unable to parse CSV file. Structure is likely malformed.",  # pylint: disable=C0301
             ) from exc
 
         if self.null_empty_strings:
@@ -200,7 +202,8 @@ def read_to_relation(  # pylint: disable=unused-argument
             raise UnableToParseCSVError(
                 entity_name="csv_structure",
                 error_code=self.ft_error_code,
-                error_message=self.ft_error_message or "Found zero records after loading CSV. File is likely malformed.",  # pylint: disable=C0301
+                error_message=self.ft_error_message
+                or "Found zero records after loading CSV. File is likely malformed.",  # pylint: disable=C0301
             )
 
         return entity

src/dve/core_engine/backends/implementations/duckdb/rules.py

@@ -421,12 +421,6 @@ def identify_orphans(
             .set_alias("orphan")
         )
 
-        if relation_is_empty(orphaned_rel):
-            self.logger.info(
-                f"Found 0 orphan records between {config.entity_name} and {config.target_name}"
-            )
-            return []
-
         message_rel = (
             entities[config.entity_name]
             .set_alias(config.entity_name)
@@ -474,15 +468,6 @@ def check_mandatory_group(
         missing_children_rel = joined_rel.filter("fk IS NULL")
         filtered_rel = joined_rel.filter("fk IS NOT NULL").select(StarExpression(exclude=["fk"]))
 
-        _no_valid_child_records: tuple[int] = missing_children_rel.count("*").fetchone()  # type: ignore # pylint: disable=C0301
-        if _no_valid_child_records:
-            _no_valid_children = _no_valid_child_records[0]
-        else:
-            _no_valid_children = 0
-        self.logger.info(
-            f"Found {_no_valid_children} records with no valid children in {config.entity_name}."
-        )  # pylint: disable=C0301
-
         entities[config.entity_name] = filtered_rel
 
         return duckdb_rel_to_dictionaries(missing_children_rel)

src/dve/core_engine/backends/implementations/spark/readers/xml.py

@@ -44,14 +44,15 @@ def read_to_dataframe(
         resource: URI,
         entity_name: EntityName,
         schema: type[BaseModel],
+        all_model_fields: Optional[set[str]] = None,
     ) -> DataFrame:
         """Stream an XML file into a Spark data frame"""
         if not self.spark:
             self.spark = SparkSession.builder.getOrCreate()  # type: ignore
         spark_schema = get_type_from_annotation(schema)
         return self.add_record_index(
             self.spark.createDataFrame(  # type: ignore
-                list(self.read_to_py_iterator(resource, entity_name, schema)),
+                list(self.read_to_py_iterator(resource, entity_name, schema, all_model_fields)),
                 schema=spark_schema,
             )
         )

src/dve/core_engine/backends/implementations/spark/rules.py

@@ -1,7 +1,7 @@
 """Step implementations in Spark."""
 
 # pylint: disable=R0801
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterable, Iterator
 from typing import Optional
 from uuid import uuid4
 
@@ -13,6 +13,7 @@
 from dve.core_engine.backends.exceptions import ConstraintError
 from dve.core_engine.backends.implementations.spark.spark_helpers import (
     create_udf,
+    df_is_empty,
     get_all_registered_udfs,
     object_to_spark_literal,
     spark_filter_contract_errors,
@@ -48,6 +49,7 @@
     SemiJoin,
     TableUnion,
 )
+from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.functions import implementations as functions
 from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.templating import template_object
@@ -341,16 +343,107 @@ def union(self, entities: SparkEntities, *, config: TableUnion) -> Messages:
         return []
 
     def identify_orphans(
-        self, entities: SparkEntities, *, config: OrphanIdentification
-    ) -> tuple[Messages, int]:
-        # TODO - adjust this to new setup of identify and remove orphans
-        raise NotImplementedError
+        self,
+        entities: SparkEntities,
+        *,
+        config: OrphanIdentification,
+    ) -> Iterable:
+        """Identify records in an entity which don't have at least one corresponding
+        match in the target. A new boolean column will be added to `entity` ('IsOrphaned')
+        indicating whether the condition matched.
+
+        If there is already an 'IsOrphaned' column in the entity, this will be set to the
+        logical OR of its current value and the value it would have been set to otherwise.
+
+        """
+        source_df: DataFrame = entities[config.entity_name]
+        source_df = source_df.alias(config.entity_name)
+
+        if df_is_empty(source_df):
+            self.logger.info(f"{config.entity_name} is empty. Skipping orphan check.")
+            return
+
+        target_df: DataFrame = entities[config.target_name]
+        match_name = f"matched_{uuid4().hex}"
+        target_df = target_df.select("*", sf.lit(1).alias(match_name)).alias(config.target_name)
+
+        orphaned_df: DataFrame = (
+            source_df.join(target_df, on=sf.expr(config.join_condition), how="left")
+            .groupBy(f"{config.entity_name}.{config.id}")
+            .agg(
+                sf.first(f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME}").alias(
+                    RECORD_INDEX_COLUMN_NAME
+                ),  # pylint: disable=C0301
+                (sf.coalesce(sf.count(match_name), sf.lit(0)) == sf.lit(0)).alias("IsOrphaned"),
+            )
+            .filter(sf.col("IsOrphaned"))
+            .select(RECORD_INDEX_COLUMN_NAME)
+            .alias("orphan")
+        )
+
+        message_df = (
+            entities[config.entity_name]
+            .alias(config.entity_name)
+            .join(
+                orphaned_df,
+                sf.expr(
+                    f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME} = orphan.{RECORD_INDEX_COLUMN_NAME}",  # pylint: disable=C0301
+                ),
+                "semi",
+            )
+        )
+        filtered_rel = (
+            entities[config.entity_name]
+            .alias(config.entity_name)
+            .join(
+                orphaned_df,
+                sf.expr(
+                    f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME} = orphan.{RECORD_INDEX_COLUMN_NAME}",  # pylint: disable=C0301
+                ),
+                "anti",
+            )
+        )
+
+        entities[config.entity_name] = filtered_rel
+
+        for r in message_df.toLocalIterator():
+            yield r.asDict()
 
     def check_mandatory_group(
         self, entities: SparkEntities, *, config: GroupIdentification
     ) -> Iterator:
-        # TODO - implement for spark
-        raise NotImplementedError
+        """
+        Check that a mandatory key in an entity has at least one valid entry in the all the
+        child entities.
+        """
+        source_df: DataFrame = entities[config.entity_name]
+        source_df = source_df.alias(config.entity_name)
+        target_df: DataFrame = entities[config.target_name]
+        target_df = target_df.alias(config.target_name)
+
+        source_columns = [f"{config.entity_name}.{c.strip()}" for c in source_df.columns]
+        _pk, fk = config.join_condition.split("=")
+
+        joined_df = source_df.join(target_df, sf.expr(config.join_condition), "left").select(
+            *source_columns,
+            sf.col(fk.strip()).alias("fk"),
+        )
+
+        missing_children_df = joined_df.filter("fk IS NULL")
+        filtered_df = joined_df.filter("fk IS NOT NULL").select("*").drop(sf.col("fk"))
+
+        if not df_is_empty(missing_children_df):
+            _no_valid_children = missing_children_df.count()
+        else:
+            _no_valid_children = 0
+        self.logger.info(
+            f"Found {_no_valid_children} records with no valid children in {config.entity_name}."
+        )
+
+        entities[config.entity_name] = filtered_df
+
+        for r in missing_children_df.toLocalIterator():
+            yield r.asDict()
 
     def filter(self, entities: SparkEntities, *, config: ImmediateFilter) -> Messages:
         """Filter an entity immediately, and do not emit any messages.

src/dve/core_engine/backends/readers/utilities.py

@@ -66,7 +66,7 @@ def raise_message_bearing_error_on_header_differences(
                 reporting_field="csv_header",
                 error_code=field_check_error_code,
                 error_message=field_check_error_message,
-            )
+            ),
         )
 
 

src/dve/pipeline/pipeline.py

@@ -764,7 +764,8 @@ def apply_business_rules(  # pylint: disable=R0914,R0915
         fh.remove_prefix(
             fh.joinuri(
                 self.processed_files_path, submission_info.submission_id, "temp_business_rules"
-            )
+            ),
+            recursive=True,
         )
 
         self.check_mandatory_entities_have_records(

tests/features/flights.feature

@@ -286,3 +286,81 @@ Feature: Pipeline tests using the flights dataset
         And the latest audit record for the submission is marked with processing status error_report
         When I run the error report phase
         Then An error report is produced
+
+    Scenario: A flights submission with a rejection on a node with two mandatory nodes (spark)
+        Given I submit the flights file multi_node_file_rejection.xml for processing
+        And A spark pipeline is configured with schema file 'flights_spark.dischema.json'
+        And I add initial audit entries for the submission
+        Then the latest audit record for the submission is marked with processing status file_transformation
+        When I run the file transformation phase
+        Then the country entity is stored as a parquet after the file_transformation phase
+        And the airport entity is stored as a parquet after the file_transformation phase
+        And the flights entity is stored as a parquet after the file_transformation phase
+        And the passengers entity is stored as a parquet after the file_transformation phase
+        And the latest audit record for the submission is marked with processing status data_contract
+        When I run the data contract phase
+        Then there are no file rejections from the data_contract phase
+        And there are no record rejections from the data_contract phase
+        When I run the business rules phase
+        Then there are errors with the following details and associated error_count from the business_rules phase
+            | ErrorType | Status | ErrorCode            | error_count |
+            | record    | error  | StaffIDMissing       | 7           |
+            | record    | error  | AirportHasNoStaff    | 1           |
+            | record    | error  | FlightHasNoAirport   | 1           |
+            | record    | error  | PassengerHasNoFlight | 1           |
+        And the final entities have the following row counts
+            | entity_name | row_count |
+            | country     | 1         |
+            | airport     | 1         |
+            | staff       | 1         |
+            | flights     | 1         |
+            | passengers  | 1         |
+        When I run the error report phase
+        Then An error report is produced
+        And The statistics entry for the submission shows the following information
+            | parameter                    | value |
+            | record_count                 | 1     |
+            | number_submission_rejections | 0     |
+            | number_record_rejections     | 10    |
+            | number_warnings              | 0     |
+
+    Scenario: A flights submission with many types of rejections in a single submission (spark)
+        Given I submit the flights file flights_full_regression.xml for processing
+        And A spark pipeline is configured with schema file 'flights_spark.dischema.json'
+        And I add initial audit entries for the submission
+        Then the latest audit record for the submission is marked with processing status file_transformation
+        When I run the file transformation phase
+        Then the country entity is stored as a parquet after the file_transformation phase
+        And the airport entity is stored as a parquet after the file_transformation phase
+        And the flights entity is stored as a parquet after the file_transformation phase
+        And the passengers entity is stored as a parquet after the file_transformation phase
+        And the latest audit record for the submission is marked with processing status data_contract
+        When I run the data contract phase
+        Then there are errors with the following details and associated error_count from the data_contract phase
+            | FailureType | ErrorCode          | error_count |
+            | record      | AirportIdIsMissing | 1           |
+        When I run the business rules phase
+        Then there are errors with the following details and associated error_count from the business_rules phase
+            | ErrorType | Status | ErrorCode                | error_count |
+            | record    | error  | InvalidFlightDestination | 1           |
+            | record    | error  | PassengerNameMissing     | 1           |
+            | record    | error  | StaffIDMissing           | 4           |
+            | record    | error  | PassengerHasNoFlight     | 3           |
+            | record    | error  | StaffHasNoAirport        | 1           |
+            | record    | error  | FlightHasNoAirport       | 2           |
+            | record    | error  | AirportHasNoStaff        | 1           |
+        And the final entities have the following row counts
+            | entity_name | row_count |
+            | country     | 1         |
+            | airport     | 3         |
+            | staff       | 3         |
+            | flights     | 3         |
+            | passengers  | 2         |
+        When I run the error report phase
+        Then An error report is produced
+        And The statistics entry for the submission shows the following information
+            | parameter                    | value |
+            | record_count                 | 1     |
+            | number_submission_rejections | 0     |
+            | number_record_rejections     | 14    |
+            | number_warnings              | 0     |