Fix Validation Step (#302)

Author: AlvaroMarquesAndrade

butterfree/load/writers/historical_feature_store_writer.py

@@ -141,14 +141,7 @@ def write(
         """
         dataframe = self._create_partitions(dataframe)
 
-        partition_df = self._apply_transformations(dataframe)
-
-        if self.debug_mode:
-            dataframe = partition_df
-        else:
-            dataframe = self.check_schema(
-                spark_client, partition_df, feature_set.name, self.database
-            )
+        dataframe = self._apply_transformations(dataframe)
 
         if self.interval_mode:
             if self.debug_mode:

butterfree/load/writers/online_feature_store_writer.py

@@ -7,7 +7,7 @@
 from pyspark.sql.functions import col, row_number
 from pyspark.sql.streaming import StreamingQuery
 
-from butterfree.clients import CassandraClient, SparkClient
+from butterfree.clients import SparkClient
 from butterfree.configs.db import AbstractWriteConfig, CassandraConfig
 from butterfree.constants.columns import TIMESTAMP_COLUMN
 from butterfree.hooks import Hook
@@ -182,22 +182,6 @@ def write(
         """
         table_name = feature_set.entity if self.write_to_entity else feature_set.name
 
-        if not self.debug_mode:
-            config = (
-                self.db_config
-                if self.db_config == CassandraConfig
-                else CassandraConfig()
-            )
-
-            cassandra_client = CassandraClient(
-                host=[config.host],
-                keyspace=config.keyspace,
-                user=config.username,
-                password=config.password,
-            )
-
-            dataframe = self.check_schema(cassandra_client, dataframe, table_name)
-
         if dataframe.isStreaming:
             dataframe = self._apply_transformations(dataframe)
             if self.debug_mode:

tests/integration/butterfree/load/test_sink.py

@@ -9,7 +9,7 @@
 )
 
 
-def test_sink(input_dataframe, feature_set, mocker):
+def test_sink(input_dataframe, feature_set):
     # arrange
     client = SparkClient()
     client.conn.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")
@@ -44,10 +44,6 @@ def test_sink(input_dataframe, feature_set, mocker):
     )
     online_writer = OnlineFeatureStoreWriter(db_config=online_config)
 
-    online_writer.check_schema_hook = mocker.stub("check_schema_hook")
-    online_writer.check_schema_hook.run = mocker.stub("run")
-    online_writer.check_schema_hook.run.return_value = feature_set_df
-
     writers = [historical_writer, online_writer]
     sink = Sink(writers)
 

tests/integration/butterfree/pipelines/test_feature_set_pipeline.py

@@ -4,7 +4,6 @@
 from pyspark.sql import DataFrame
 from pyspark.sql import functions as F
 
-from butterfree.clients import SparkClient
 from butterfree.configs import environment
 from butterfree.configs.db import MetastoreConfig
 from butterfree.constants import DataType

tests/unit/butterfree/load/test_sink.py

@@ -120,7 +120,7 @@ def test_flush_with_writers_list_empty(self):
         with pytest.raises(ValueError):
             Sink(writers=writer)
 
-    def test_flush_streaming_df(self, feature_set, mocker):
+    def test_flush_streaming_df(self, feature_set):
         """Testing the return of the streaming handlers by the sink."""
         # arrange
         spark_client = SparkClient()
@@ -141,14 +141,6 @@ def test_flush_streaming_df(self, feature_set, mocker):
             write_to_entity=True
         )
 
-        online_feature_store_writer_on_entity.check_schema_hook = mocker.stub(
-            "check_schema_hook"
-        )
-        online_feature_store_writer_on_entity.check_schema_hook.run = mocker.stub("run")
-        online_feature_store_writer_on_entity.check_schema_hook.run.return_value = (
-            mocked_stream_df
-        )
-
         sink = Sink(
             writers=[
                 online_feature_store_writer,
@@ -171,7 +163,7 @@ def test_flush_streaming_df(self, feature_set, mocker):
             assert isinstance(handler, StreamingQuery)
 
     def test_flush_with_multiple_online_writers(
-        self, feature_set, feature_set_dataframe, mocker
+        self, feature_set, feature_set_dataframe
     ):
         """Testing the flow of writing to a feature-set table and to an entity table."""
         # arrange
@@ -187,14 +179,6 @@ def test_flush_with_multiple_online_writers(
             write_to_entity=True
         )
 
-        online_feature_store_writer_on_entity.check_schema_hook = mocker.stub(
-            "check_schema_hook"
-        )
-        online_feature_store_writer_on_entity.check_schema_hook.run = mocker.stub("run")
-        online_feature_store_writer_on_entity.check_schema_hook.run.return_value = (
-            feature_set_dataframe
-        )
-
         sink = Sink(
             writers=[online_feature_store_writer, online_feature_store_writer_on_entity]
         )

tests/unit/butterfree/load/writers/test_historical_feature_store_writer.py

@@ -23,11 +23,6 @@ def test_write(
         spark_client.write_table = mocker.stub("write_table")
         writer = HistoricalFeatureStoreWriter()
 
-        schema_dataframe = writer._create_partitions(feature_set_dataframe)
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = schema_dataframe
-
         # when
         writer.write(
             feature_set=feature_set,
@@ -107,7 +102,6 @@ def test_write_in_debug_mode(
         historical_feature_set_dataframe,
         feature_set,
         spark_session,
-        mocker,
     ):
         # given
         spark_client = SparkClient()
@@ -310,12 +304,6 @@ def test_write_with_transform(
 
         writer = HistoricalFeatureStoreWriter().with_(json_transform)
 
-        schema_dataframe = writer._create_partitions(feature_set_dataframe)
-        json_dataframe = writer._apply_transformations(schema_dataframe)
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = json_dataframe
-
         # when
         writer.write(
             feature_set=feature_set,

tests/unit/butterfree/load/writers/test_online_feature_store_writer.py

@@ -68,10 +68,6 @@ def test_write(
         spark_client.write_dataframe = mocker.stub("write_dataframe")
         writer = OnlineFeatureStoreWriter(cassandra_config)
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = feature_set_dataframe
-
         # when
         writer.write(feature_set, feature_set_dataframe, spark_client)
 
@@ -104,10 +100,6 @@ def test_write_in_debug_mode(
         spark_client = SparkClient()
         writer = OnlineFeatureStoreWriter(debug_mode=True)
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = feature_set_dataframe
-
         # when
         writer.write(
             feature_set=feature_set,
@@ -132,10 +124,6 @@ def test_write_in_debug_and_stream_mode(self, feature_set, spark_session):
 
         writer = OnlineFeatureStoreWriter(debug_mode=True)
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = mocked_stream_df
-
         # act
         handler = writer.write(
             feature_set=feature_set,
@@ -151,7 +139,7 @@ def test_write_in_debug_and_stream_mode(self, feature_set, spark_session):
         assert isinstance(handler, StreamingQuery)
 
     @pytest.mark.parametrize("has_checkpoint", [True, False])
-    def test_write_stream(self, feature_set, has_checkpoint, monkeypatch, mocker):
+    def test_write_stream(self, feature_set, has_checkpoint, monkeypatch):
         # arrange
         spark_client = SparkClient()
         spark_client.write_stream = Mock()
@@ -174,10 +162,6 @@ def test_write_stream(self, feature_set, has_checkpoint, monkeypatch, mocker):
         writer = OnlineFeatureStoreWriter(cassandra_config)
         writer.filter_latest = Mock()
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = dataframe
-
         # act
         stream_handler = writer.write(feature_set, dataframe, spark_client)
 
@@ -201,7 +185,7 @@ def test_get_db_schema(self, cassandra_config, test_feature_set, expected_schema
 
         assert schema == expected_schema
 
-    def test_write_stream_on_entity(self, feature_set, monkeypatch, mocker):
+    def test_write_stream_on_entity(self, feature_set, monkeypatch):
         """Test write method with stream dataframe and write_to_entity enabled.
 
         The main purpose of this test is assert the correct setup of stream checkpoint
@@ -224,10 +208,6 @@ def test_write_stream_on_entity(self, feature_set, monkeypatch, mocker):
 
         writer = OnlineFeatureStoreWriter(write_to_entity=True)
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = dataframe
-
         # act
         stream_handler = writer.write(feature_set, dataframe, spark_client)
 
@@ -256,10 +236,6 @@ def test_write_with_transform(
         spark_client.write_dataframe = mocker.stub("write_dataframe")
         writer = OnlineFeatureStoreWriter(cassandra_config).with_(json_transform)
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = feature_set_dataframe
-
         # when
         writer.write(feature_set, feature_set_dataframe, spark_client)
 
@@ -293,10 +269,6 @@ def test_write_with_kafka_config(
         kafka_config = KafkaConfig()
         writer = OnlineFeatureStoreWriter(kafka_config).with_(json_transform)
 
-        writer.check_schema_hook = mocker.stub("check_schema_hook")
-        writer.check_schema_hook.run = mocker.stub("run")
-        writer.check_schema_hook.run.return_value = feature_set_dataframe
-
         # when
         writer.write(feature_set, feature_set_dataframe, spark_client)
 
@@ -320,10 +292,6 @@ def test_write_with_custom_kafka_config(
             json_transform
         )
 
-        custom_writer.check_schema_hook = mocker.stub("check_schema_hook")
-        custom_writer.check_schema_hook.run = mocker.stub("run")
-        custom_writer.check_schema_hook.run.return_value = feature_set_dataframe
-
         # when
         custom_writer.write(feature_set, feature_set_dataframe, spark_client)