Release/1.4.1 (#375)

* fix: performance improvements (#374)

CHANGELOG.md

@@ -5,6 +5,9 @@ Preferably use **Added**, **Changed**, **Removed** and **Fixed** topics in each
 
 ## [Unreleased]
 
+## [1.4.1](https://github.com/quintoandar/butterfree/releases/tag/1.4.1)
+* Performance Improvements ([#374](https://github.com/quintoandar/butterfree/pull/374))
+
 ## [1.4.0](https://github.com/quintoandar/butterfree/releases/tag/1.4.0)
 * Add Delta support ([#370](https://github.com/quintoandar/butterfree/pull/370))
 

butterfree/_cli/migrate.py

@@ -4,7 +4,7 @@
 import os
 import pkgutil
 import sys
-from typing import Set
+from typing import Set, Type
 
 import boto3
 import setuptools
@@ -90,8 +90,18 @@ def __fs_objects(path: str) -> Set[FeatureSetPipeline]:
 
             instances.add(value)
 
+    def create_instance(cls: Type[FeatureSetPipeline]) -> FeatureSetPipeline:
+        sig = inspect.signature(cls.__init__)
+        parameters = sig.parameters
+
+        if "run_date" in parameters:
+            run_date = datetime.datetime.today().strftime("%y-%m-%d")
+            return cls(run_date)
+
+        return cls()
+
     logger.info("Creating instances...")
-    return set(value() for value in instances)  # type: ignore
+    return set(create_instance(value) for value in instances)  # type: ignore
 
 
 PATH = typer.Argument(

butterfree/extract/source.py

@@ -3,6 +3,7 @@
 from typing import List, Optional
 
 from pyspark.sql import DataFrame
+from pyspark.storagelevel import StorageLevel
 
 from butterfree.clients import SparkClient
 from butterfree.extract.readers.reader import Reader
@@ -95,16 +96,21 @@ def construct(
             DataFrame with the query result against all readers.
 
         """
+        # Step 1: Build temporary views for each reader
         for reader in self.readers:
-            reader.build(
-                client=client, start_date=start_date, end_date=end_date
-            )  # create temporary views for each reader
+            reader.build(client=client, start_date=start_date, end_date=end_date)
 
+        # Step 2: Execute SQL query on the combined readers
         dataframe = client.sql(self.query)
 
+        # Step 3: Cache the dataframe if necessary, using memory and disk storage
         if not dataframe.isStreaming and self.eager_evaluation:
-            dataframe.cache().count()
+            # Persist to ensure the DataFrame is stored in mem and disk (if necessary)
+            dataframe.persist(StorageLevel.MEMORY_AND_DISK)
+            # Trigger the cache/persist operation by performing an action
+            dataframe.count()
 
+        # Step 4: Run post-processing hooks on the dataframe
         post_hook_df = self.run_post_hooks(dataframe)
 
         return post_hook_df

butterfree/pipelines/feature_set_pipeline.py

@@ -2,6 +2,8 @@
 
 from typing import List, Optional
 
+from pyspark.storagelevel import StorageLevel
+
 from butterfree.clients import SparkClient
 from butterfree.dataframe_service import repartition_sort_df
 from butterfree.extract import Source
@@ -209,35 +211,47 @@ def run(
         soon. Use only if strictly necessary.
 
         """
+
+        # Step 1: Construct input dataframe from the source.
         dataframe = self.source.construct(
             client=self.spark_client,
             start_date=self.feature_set.define_start_date(start_date),
             end_date=end_date,
         )
 
+        # Step 2: Repartition and sort if required, avoid if not necessary.
         if partition_by:
             order_by = order_by or partition_by
-            dataframe = repartition_sort_df(
-                dataframe, partition_by, order_by, num_processors
-            )
-
-        dataframe = self.feature_set.construct(
+            current_partitions = dataframe.rdd.getNumPartitions()
+            optimal_partitions = num_processors or current_partitions
+            if current_partitions != optimal_partitions:
+                dataframe = repartition_sort_df(
+                    dataframe, partition_by, order_by, num_processors
+                )
+
+        # Step 3: Construct the feature set dataframe using defined transformations.
+        transformed_dataframe = self.feature_set.construct(
             dataframe=dataframe,
             client=self.spark_client,
             start_date=start_date,
             end_date=end_date,
             num_processors=num_processors,
         )
 
+        if dataframe.storageLevel != StorageLevel.NONE:
+            dataframe.unpersist()  # Clear the data from the cache (disk and memory)
+
+        # Step 4: Load the data into the configured sink.
         self.sink.flush(
-            dataframe=dataframe,
+            dataframe=transformed_dataframe,
             feature_set=self.feature_set,
             spark_client=self.spark_client,
         )
 
-        if not dataframe.isStreaming:
+        # Step 5: Validate the output if not streaming and data volume is reasonable.
+        if not transformed_dataframe.isStreaming:
             self.sink.validate(
-                dataframe=dataframe,
+                dataframe=transformed_dataframe,
                 feature_set=self.feature_set,
                 spark_client=self.spark_client,
             )

butterfree/transform/aggregated_feature_set.py

@@ -387,6 +387,7 @@ def _aggregate(
         ]
 
         groupby = self.keys_columns.copy()
+
         if window is not None:
             dataframe = dataframe.withColumn("window", window.get())
             groupby.append("window")
@@ -410,19 +411,23 @@ def _aggregate(
                 "keep_rn", functions.row_number().over(partition_window)
             ).filter("keep_rn = 1")
 
-        # repartition to have all rows for each group at the same partition
-        # by doing that, we won't have to shuffle data on grouping by id
-        dataframe = repartition_df(
-            dataframe,
-            partition_by=groupby,
-            num_processors=num_processors,
-        )
+        current_partitions = dataframe.rdd.getNumPartitions()
+        optimal_partitions = num_processors or current_partitions
+
+        if current_partitions != optimal_partitions:
+            dataframe = repartition_df(
+                dataframe,
+                partition_by=groupby,
+                num_processors=optimal_partitions,
+            )
+
         grouped_data = dataframe.groupby(*groupby)
 
-        if self._pivot_column:
+        if self._pivot_column and self._pivot_values:
             grouped_data = grouped_data.pivot(self._pivot_column, self._pivot_values)
 
         aggregated = grouped_data.agg(*aggregations)
+
         return self._with_renamed_columns(aggregated, features, window)
 
     def _with_renamed_columns(
@@ -637,12 +642,13 @@ def construct(
         output_df = output_df.select(*self.columns).replace(  # type: ignore
             float("nan"), None
         )
-        if not output_df.isStreaming:
-            if self.deduplicate_rows:
-                output_df = self._filter_duplicated_rows(output_df)
-            if self.eager_evaluation:
-                output_df.cache().count()
+
+        if not output_df.isStreaming and self.deduplicate_rows:
+            output_df = self._filter_duplicated_rows(output_df)
 
         post_hook_df = self.run_post_hooks(output_df)
 
+        if not output_df.isStreaming and self.eager_evaluation:
+            post_hook_df.cache().count()
+
         return post_hook_df

butterfree/transform/feature_set.py

@@ -436,11 +436,8 @@ def construct(
             pre_hook_df,
         ).select(*self.columns)
 
-        if not output_df.isStreaming:
-            if self.deduplicate_rows:
-                output_df = self._filter_duplicated_rows(output_df)
-            if self.eager_evaluation:
-                output_df.cache().count()
+        if not output_df.isStreaming and self.deduplicate_rows:
+            output_df = self._filter_duplicated_rows(output_df)
 
         output_df = self.incremental_strategy.filter_with_incremental_strategy(
             dataframe=output_df, start_date=start_date, end_date=end_date

docs/source/butterfree.configs.rst

@@ -23,26 +23,6 @@ butterfree.configs.environment module
 butterfree.configs.logger module
 --------------------------------
 
-.. automodule:: butterfree.configs.logger
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. automodule:: butterfree.configs.logger
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. automodule:: butterfree.configs.logger
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. automodule:: butterfree.configs.logger
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 .. automodule:: butterfree.configs.logger
    :members:
    :undoc-members:

docs/source/butterfree.constants.rst

@@ -31,28 +31,6 @@ butterfree.constants.migrations module
 butterfree.constants.spark\_constants module
 --------------------------------------------
 
-.. automodule:: butterfree.constants.migrations
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-
-.. automodule:: butterfree.constants.migrations
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-
-.. automodule:: butterfree.constants.migrations
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-
-.. automodule:: butterfree.constants.migrations
-   :members:
-   :undoc-members:
-   :show-inheritance:
 
 .. automodule:: butterfree.constants.spark_constants
    :members:
@@ -62,26 +40,6 @@ butterfree.constants.spark\_constants module
 butterfree.constants.window\_definitions module
 -----------------------------------------------
 
-.. automodule:: butterfree.constants.window_definitions
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. automodule:: butterfree.constants.window_definitions
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. automodule:: butterfree.constants.window_definitions
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
-.. automodule:: butterfree.constants.window_definitions
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 .. automodule:: butterfree.constants.window_definitions
    :members:
    :undoc-members:

docs/source/butterfree.dataframe_service.rst

@@ -20,14 +20,6 @@ butterfree.dataframe\_service.partitioning module
    :undoc-members:
    :show-inheritance:
 
-butterfree.dataframe\_service.repartition module
-------------------------------------------------
-
-.. automodule:: butterfree.dataframe_service.repartition
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 .. automodule:: butterfree.dataframe_service.repartition
    :members:
    :undoc-members:

setup.py

@@ -1,7 +1,7 @@
 from setuptools import find_packages, setup
 
 __package_name__ = "butterfree"
-__version__ = "1.4.0"
+__version__ = "1.4.1"
 __repository_url__ = "https://github.com/quintoandar/butterfree"
 
 with open("requirements.txt") as f: