feat: optional row count validation (#340)

quintoandar · Aug 18, 2023 · 2a19009 · 2a19009 · chip-n-dale · Aug 18, 2023
1 parent 6b78a50
commit 2a19009
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/butterfree/load/sink.py b/butterfree/load/sink.py
@@ -69,14 +69,15 @@ def validate(
         """
         failures = []
         for writer in self.writers:
-            try:
-                writer.validate(
-                    feature_set=feature_set,
-                    dataframe=dataframe,
-                    spark_client=spark_client,
-                )
-            except AssertionError as e:
-                failures.append(e)
+            if writer.row_count_validation:
+                try:
+                    writer.validate(
+                        feature_set=feature_set,
+                        dataframe=dataframe,
+                        spark_client=spark_client,
+                    )
+                except AssertionError as e:
+                    failures.append(e)
 
         if failures:
             raise RuntimeError(

diff --git a/butterfree/load/writers/writer.py b/butterfree/load/writers/writer.py
@@ -26,13 +26,15 @@ def __init__(
         debug_mode: bool = False,
         interval_mode: bool = False,
         write_to_entity: bool = False,
+        row_count_validation: bool = True,
     ) -> None:
         super().__init__()
         self.db_config = db_config
         self.transformations: List[Dict[str, Any]] = []
         self.debug_mode = debug_mode
         self.interval_mode = interval_mode
         self.write_to_entity = write_to_entity
+        self.row_count_validation = row_count_validation
 
     def with_(
         self, transformer: Callable[..., DataFrame], *args: Any, **kwargs: Any