workbench 0.8.177__py3-none-any.whl → 0.8.227__py3-none-any.whl

workbench/core/artifacts/feature_set_core.py

@@ -7,7 +7,6 @@ from datetime import datetime, timezone
 import botocore.exceptions
 import pandas as pd
 import awswrangler as wr
-import numpy as np
 
 from sagemaker.feature_store.feature_group import FeatureGroup
 from sagemaker.feature_store.feature_store import FeatureStore
@@ -17,7 +16,7 @@ from workbench.core.artifacts.artifact import Artifact
 from workbench.core.artifacts.data_source_factory import DataSourceFactory
 from workbench.core.artifacts.athena_source import AthenaSource
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Dict, Union
 
 from workbench.utils.aws_utils import aws_throttle
 
@@ -194,24 +193,24 @@ class FeatureSetCore(Artifact):
 
         return View(self, view_name)
 
-    def set_display_columns(self, diplay_columns: list[str]):
+    def set_display_columns(self, display_columns: list[str]):
         """Set the display columns for this Data Source
 
         Args:
-            diplay_columns (list[str]): The display columns for this Data Source
+            display_columns (list[str]): The display columns for this Data Source
         """
         # Check mismatch of display columns to computation columns
         c_view = self.view("computation")
         computation_columns = c_view.columns
-        mismatch_columns = [col for col in diplay_columns if col not in computation_columns]
+        mismatch_columns = [col for col in display_columns if col not in computation_columns]
         if mismatch_columns:
             self.log.monitor(f"Display View/Computation mismatch: {mismatch_columns}")
 
-        self.log.important(f"Setting Display Columns...{diplay_columns}")
+        self.log.important(f"Setting Display Columns...{display_columns}")
         from workbench.core.views import DisplayView
 
         # Create a NEW display view
-        DisplayView.create(self, source_table=c_view.table, column_list=diplay_columns)
+        DisplayView.create(self, source_table=c_view.table, column_list=display_columns)
 
     def set_computation_columns(self, computation_columns: list[str], reset_display: bool = True):
         """Set the computation columns for this FeatureSet
@@ -246,7 +245,7 @@ class FeatureSetCore(Artifact):
 
         # Set the compressed features in our FeatureSet metadata
         self.log.important(f"Setting Compressed Columns...{compressed_columns}")
-        self.upsert_workbench_meta({"comp_features": compressed_columns})
+        self.upsert_workbench_meta({"compressed_features": compressed_columns})
 
     def get_compressed_features(self) -> list[str]:
         """Get the compressed features for this FeatureSet
@@ -255,7 +254,7 @@ class FeatureSetCore(Artifact):
             list[str]: The compressed columns for this FeatureSet
         """
         # Get the compressed features from our FeatureSet metadata
-        return self.workbench_meta().get("comp_features", [])
+        return self.workbench_meta().get("compressed_features", [])
 
     def num_columns(self) -> int:
         """Return the number of columns of the Feature Set"""
@@ -482,18 +481,6 @@ class FeatureSetCore(Artifact):
             time.sleep(1)
         cls.log.info(f"FeatureSet {feature_group.name} successfully deleted")
 
-    def set_training_holdouts(self, holdout_ids: list[str]):
-        """Set the hold out ids for the training view for this FeatureSet
-
-        Args:
-            holdout_ids (list[str]): The list of holdout ids.
-        """
-        from workbench.core.views import TrainingView
-
-        # Create a NEW training view
-        self.log.important(f"Setting Training Holdouts: {len(holdout_ids)} ids...")
-        TrainingView.create(self, id_column=self.id_column, holdout_ids=holdout_ids)
-
     def get_training_holdouts(self) -> list[str]:
         """Get the hold out ids for the training view for this FeatureSet
 
@@ -509,6 +496,177 @@ class FeatureSetCore(Artifact):
         ].tolist()
         return hold_out_ids
 
+    # ---- Public methods for training configuration ----
+    def set_training_config(
+        self,
+        holdout_ids: List[Union[str, int]] = None,
+        weight_dict: Dict[Union[str, int], float] = None,
+        default_weight: float = 1.0,
+        exclude_zero_weights: bool = True,
+    ):
+        """Configure training view with holdout IDs and/or sample weights.
+
+        This method creates a training view that can include both:
+        - A 'training' column (True/False) based on holdout IDs
+        - A 'sample_weight' column for weighted training
+
+        Args:
+            holdout_ids: List of IDs to mark as training=False (validation/holdout set)
+            weight_dict: Mapping of ID to sample weight
+                - weight > 1.0: oversample/emphasize
+                - weight = 1.0: normal (default)
+                - 0 < weight < 1.0: downweight/de-emphasize
+                - weight = 0.0: exclude from training (filtered out if exclude_zero_weights=True)
+            default_weight: Weight for IDs not in weight_dict (default: 1.0)
+            exclude_zero_weights: If True, filter out rows with sample_weight=0 (default: True)
+
+        Example:
+            # Temporal split with sample weights
+            fs.set_training_config(
+                holdout_ids=temporal_hold_out_ids,  # IDs after cutoff date
+                weight_dict={'compound_42': 0.0, 'compound_99': 2.0},  # exclude/upweight
+            )
+        """
+        from workbench.core.views.training_view import TrainingView
+
+        # If neither is provided, create a standard training view
+        if not holdout_ids and not weight_dict:
+            self.log.important("No holdouts or weights specified, creating standard training view")
+            TrainingView.create(self, id_column=self.id_column)
+            return
+
+        # If only holdout_ids, delegate to set_training_holdouts
+        if holdout_ids and not weight_dict:
+            self.set_training_holdouts(holdout_ids)
+            return
+
+        # If only weight_dict, delegate to set_sample_weights
+        if weight_dict and not holdout_ids:
+            self.set_sample_weights(weight_dict, default_weight, exclude_zero_weights)
+            return
+
+        # Both holdout_ids and weight_dict provided - build combined view
+        self.log.important(f"Setting training config: {len(holdout_ids)} holdouts, {len(weight_dict)} weights")
+
+        # Get column list (excluding AWS-generated columns)
+        from workbench.core.views.view_utils import get_column_list
+
+        aws_cols = ["write_time", "api_invocation_time", "is_deleted", "event_time"]
+        source_columns = get_column_list(self.data_source, self.table)
+        column_list = [col for col in source_columns if col not in aws_cols]
+        sql_columns = ", ".join([f'"{column}"' for column in column_list])
+
+        # Build inner query with both columns
+        training_case = self._build_holdout_case(holdout_ids)
+        weight_case = self._build_weight_case(weight_dict, default_weight)
+        inner_sql = f"SELECT {sql_columns}, {training_case}, {weight_case} FROM {self.table}"
+
+        # Optionally filter out zero weights
+        if exclude_zero_weights:
+            zero_count = sum(1 for w in weight_dict.values() if w == 0.0)
+            if zero_count:
+                self.log.important(f"Filtering out {zero_count} rows with sample_weight = 0")
+            sql_query = f"SELECT * FROM ({inner_sql}) WHERE sample_weight > 0"
+        else:
+            sql_query = inner_sql
+
+        self._create_training_view(sql_query)
+
+    def set_training_holdouts(self, holdout_ids: list[str]):
+        """Set the hold out ids for the training view for this FeatureSet
+
+        Args:
+            holdout_ids (list[str]): The list of holdout ids.
+        """
+        from workbench.core.views import TrainingView
+
+        self.log.important(f"Setting Training Holdouts: {len(holdout_ids)} ids...")
+        TrainingView.create(self, id_column=self.id_column, holdout_ids=holdout_ids)
+
+    def set_sample_weights(
+        self,
+        weight_dict: Dict[Union[str, int], float],
+        default_weight: float = 1.0,
+        exclude_zero_weights: bool = True,
+    ):
+        """Configure training view with sample weights for each ID.
+
+        Args:
+            weight_dict: Mapping of ID to sample weight
+                - weight > 1.0: oversample/emphasize
+                - weight = 1.0: normal (default)
+                - 0 < weight < 1.0: downweight/de-emphasize
+                - weight = 0.0: exclude from training
+            default_weight: Weight for IDs not in weight_dict (default: 1.0)
+            exclude_zero_weights: If True, filter out rows with sample_weight=0 (default: True)
+
+        Example:
+            weights = {
+                'compound_42': 3.0,  # oversample 3x
+                'compound_99': 0.1,  # noisy, downweight
+                'compound_123': 0.0, # exclude from training
+            }
+            fs.set_sample_weights(weights)  # zeros automatically excluded
+            fs.set_sample_weights(weights, exclude_zero_weights=False)  # keep zeros
+        """
+        from workbench.core.views import TrainingView
+
+        if not weight_dict:
+            self.log.important("Empty weight_dict, creating standard training view")
+            TrainingView.create(self, id_column=self.id_column)
+            return
+
+        self.log.important(f"Setting sample weights for {len(weight_dict)} IDs")
+
+        # Build inner query with sample weights
+        weight_case = self._build_weight_case(weight_dict, default_weight)
+        inner_sql = f"SELECT *, {weight_case} FROM {self.table}"
+
+        # Optionally filter out zero weights
+        if exclude_zero_weights:
+            zero_count = sum(1 for w in weight_dict.values() if w == 0.0)
+            self.log.important(f"Filtering out {zero_count} rows with sample_weight = 0")
+            sql_query = f"SELECT * FROM ({inner_sql}) WHERE sample_weight > 0"
+        else:
+            sql_query = inner_sql
+
+        TrainingView.create_with_sql(self, sql_query=sql_query, id_column=self.id_column)
+
+    # ---- Internal helpers for training view SQL generation ----
+    @staticmethod
+    def _format_id_for_sql(id_val: Union[str, int]) -> str:
+        """Format an ID value for use in SQL."""
+        return repr(id_val)
+
+    def _build_holdout_case(self, holdout_ids: List[Union[str, int]]) -> str:
+        """Build SQL CASE statement for training column based on holdout IDs."""
+        if all(isinstance(id_val, str) for id_val in holdout_ids):
+            formatted_ids = ", ".join(f"'{id_val}'" for id_val in holdout_ids)
+        else:
+            formatted_ids = ", ".join(map(str, holdout_ids))
+        return f"""CASE
+            WHEN {self.id_column} IN ({formatted_ids}) THEN False
+            ELSE True
+        END AS training"""
+
+    def _build_weight_case(self, weight_dict: Dict[Union[str, int], float], default_weight: float) -> str:
+        """Build SQL CASE statement for sample_weight column."""
+        conditions = [
+            f"WHEN {self.id_column} = {self._format_id_for_sql(id_val)} THEN {weight}"
+            for id_val, weight in weight_dict.items()
+        ]
+        case_body = "\n            ".join(conditions)
+        return f"""CASE
+            {case_body}
+            ELSE {default_weight}
+        END AS sample_weight"""
+
+    def _create_training_view(self, sql_query: str):
+        """Create the training view directly from a SQL query."""
+        view_table = f"{self.table}___training"
+        create_view_query = f"CREATE OR REPLACE VIEW {view_table} AS\n{sql_query}"
+        self.data_source.execute_statement(create_view_query)
+
     @classmethod
     def delete_views(cls, table: str, database: str):
         """Delete any views associated with this FeatureSet
@@ -558,20 +716,6 @@ class FeatureSetCore(Artifact):
         """
         return self.data_source.smart_sample()
 
-    def anomalies(self) -> pd.DataFrame:
-        """Get a set of anomalous data from the underlying DataSource
-        Returns:
-            pd.DataFrame: A dataframe of anomalies from the underlying DataSource
-        """
-
-        # FIXME: Mock this for now
-        anom_df = self.sample().copy()
-        anom_df["anomaly_score"] = np.random.rand(anom_df.shape[0])
-        anom_df["cluster"] = np.random.randint(0, 10, anom_df.shape[0])
-        anom_df["x"] = np.random.rand(anom_df.shape[0])
-        anom_df["y"] = np.random.rand(anom_df.shape[0])
-        return anom_df
-
     def value_counts(self) -> dict:
         """Get the value counts for the string columns of the underlying DataSource
 
@@ -667,7 +811,7 @@ if __name__ == "__main__":
     pd.set_option("display.width", 1000)
 
     # Grab a FeatureSet object and pull some information from it
-    my_features = LocalFeatureSetCore("test_features")
+    my_features = LocalFeatureSetCore("abalone_features")
     if not my_features.exists():
         print("FeatureSet not found!")
         sys.exit(1)
@@ -707,7 +851,7 @@ if __name__ == "__main__":
 
     # Test getting the holdout ids
     print("Getting the hold out ids...")
-    holdout_ids = my_features.get_training_holdouts("id")
+    holdout_ids = my_features.get_training_holdouts()
     print(f"Holdout IDs: {holdout_ids}")
 
     # Get a sample of the data
@@ -727,18 +871,78 @@ if __name__ == "__main__":
     # Set the holdout ids for the training view
     print("Setting hold out ids...")
     table = my_features.view("training").table
-    df = my_features.query(f'SELECT id, name FROM "{table}"')
-    my_holdout_ids = [id for id in df["id"] if id < 20]
-    my_features.set_training_holdouts("id", my_holdout_ids)
-
-    # Test the hold out set functionality with strings
-    print("Setting hold out ids (strings)...")
-    my_holdout_ids = [name for name in df["name"] if int(name.split(" ")[1]) > 80]
-    my_features.set_training_holdouts("name", my_holdout_ids)
+    df = my_features.query(f'SELECT auto_id, length FROM "{table}"')
+    my_holdout_ids = [id for id in df["auto_id"] if id < 20]
+    my_features.set_training_holdouts(my_holdout_ids)
 
     # Get the training data
     print("Getting the training data...")
     training_data = my_features.get_training_data()
+    print(f"Training Data: {training_data.shape}")
+
+    # Test set_sample_weights
+    print("\n--- Testing set_sample_weights ---")
+    sample_ids = df["auto_id"].tolist()[:5]
+    weight_dict = {sample_ids[0]: 0.0, sample_ids[1]: 0.5, sample_ids[2]: 2.0}
+    my_features.set_sample_weights(weight_dict)
+    training_view = my_features.view("training")
+    training_df = training_view.pull_dataframe()
+    print(f"Training view shape after set_sample_weights: {training_df.shape}")
+    print(f"Columns: {training_df.columns.tolist()}")
+    assert "sample_weight" in training_df.columns, "sample_weight column missing!"
+    assert "training" in training_df.columns, "training column missing!"
+    # Verify zero-weight row was excluded
+    assert sample_ids[0] not in training_df["auto_id"].values, "Zero-weight ID should be excluded!"
+    print("set_sample_weights test passed!")
+
+    # Test set_training_config with both holdouts and weights
+    print("\n--- Testing set_training_config (combined) ---")
+    holdout_ids = [id for id in df["auto_id"] if id >= 100 and id < 120]
+    weight_dict = {sample_ids[3]: 0.0, sample_ids[4]: 3.0}  # exclude one, upweight another
+    my_features.set_training_config(holdout_ids=holdout_ids, weight_dict=weight_dict)
+    training_view = my_features.view("training")
+    training_df = training_view.pull_dataframe()
+    print(f"Training view shape after set_training_config: {training_df.shape}")
+    print(f"Columns: {training_df.columns.tolist()}")
+    assert "sample_weight" in training_df.columns, "sample_weight column missing!"
+    assert "training" in training_df.columns, "training column missing!"
+    # Verify holdout IDs are marked as training=False
+    holdout_rows = training_df[training_df["auto_id"].isin(holdout_ids)]
+    assert all(holdout_rows["training"] == False), "Holdout IDs should have training=False!"  # noqa: E712
+    # Verify zero-weight row was excluded
+    assert sample_ids[3] not in training_df["auto_id"].values, "Zero-weight ID should be excluded!"
+    # Verify upweighted row has correct weight
+    upweight_row = training_df[training_df["auto_id"] == sample_ids[4]]
+    assert upweight_row["sample_weight"].iloc[0] == 3.0, "Upweighted ID should have weight=3.0!"
+    print("set_training_config (combined) test passed!")
+
+    # Test set_training_config with only holdouts (should delegate to set_training_holdouts)
+    print("\n--- Testing set_training_config (holdouts only) ---")
+    my_features.set_training_config(holdout_ids=holdout_ids)
+    training_view = my_features.view("training")
+    training_df = training_view.pull_dataframe()
+    assert "training" in training_df.columns, "training column missing!"
+    holdout_rows = training_df[training_df["auto_id"].isin(holdout_ids)]
+    assert all(holdout_rows["training"] == False), "Holdout IDs should have training=False!"  # noqa: E712
+    print("set_training_config (holdouts only) test passed!")
+
+    # Test set_training_config with only weights (should delegate to set_sample_weights)
+    print("\n--- Testing set_training_config (weights only) ---")
+    my_features.set_training_config(weight_dict={sample_ids[0]: 0.5, sample_ids[1]: 2.0})
+    training_view = my_features.view("training")
+    training_df = training_view.pull_dataframe()
+    assert "sample_weight" in training_df.columns, "sample_weight column missing!"
+    print("set_training_config (weights only) test passed!")
+
+    # Test set_training_config with neither (should create standard training view)
+    print("\n--- Testing set_training_config (neither) ---")
+    my_features.set_training_config()
+    training_view = my_features.view("training")
+    training_df = training_view.pull_dataframe()
+    assert "training" in training_df.columns, "training column missing!"
+    print("set_training_config (neither) test passed!")
+
+    print("\n=== All training config tests passed! ===")
 
     # Now delete the AWS artifacts associated with this Feature Set
     # print("Deleting Workbench Feature Set...")