workbench 0.8.201__py3-none-any.whl → 0.8.204__py3-none-any.whl

workbench/api/df_store.py

@@ -1,35 +1,32 @@
 """DFStore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy"""
 
-from datetime import datetime
 from typing import Union
-import logging
-import pandas as pd
 
 # Workbench Imports
-from workbench.core.cloud_platform.aws.aws_df_store import AWSDFStore
+from workbench.core.artifacts.df_store_core import DFStoreCore
 
 
-class DFStore(AWSDFStore):
+class DFStore(DFStoreCore):
     """DFStore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy
 
-    Common Usage:
-        ```python
-        df_store = DFStore()
+        Common Usage:
+    ```python
+            df_store = DFStore()
 
-        # List Data
-        df_store.list()
+            # List Data
+            df_store.list()
 
-        # Add DataFrame
-        df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
-        df_store.upsert("/test/my_data", df)
+            # Add DataFrame
+            df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
+            df_store.upsert("/test/my_data", df)
 
-        # Retrieve DataFrame
-        df = df_store.get("/test/my_data")
-        print(df)
+            # Retrieve DataFrame
+            df = df_store.get("/test/my_data")
+            print(df)
 
-        # Delete Data
-        df_store.delete("/test/my_data")
-        ```
+            # Delete Data
+            df_store.delete("/test/my_data")
+    ```
     """
 
     def __init__(self, path_prefix: Union[str, None] = None):
@@ -38,101 +35,13 @@ class DFStore(AWSDFStore):
         Args:
             path_prefix (Union[str, None], optional): Add a path prefix to storage locations (Defaults to None)
         """
-        self.log = logging.getLogger("workbench")
-
-        # Initialize the SuperClass
         super().__init__(path_prefix=path_prefix)
 
-    def list(self, include_cache: bool = False) -> list:
-        """List all the objects in the data_store prefix.
-
-        Args:
-            include_cache (bool, optional): Include cache objects in the list (Defaults to False).
-
-        Returns:
-            list: A list of all the objects in the data_store prefix.
-        """
-        return super().list(include_cache=include_cache)
-
-    def summary(self, include_cache: bool = False) -> pd.DataFrame:
-        """Return a nicely formatted summary of object locations, sizes (in MB), and modified dates.
-
-        Args:
-            include_cache (bool, optional): Include cache objects in the summary (Defaults to False).
-
-        Returns:
-            pd.DataFrame: A formatted DataFrame with the summary details.
-        """
-        return super().summary(include_cache=include_cache)
-
-    def details(self, include_cache: bool = False) -> pd.DataFrame:
-        """Return a DataFrame with detailed metadata for all objects in the data_store prefix.
-
-        Args:
-            include_cache (bool, optional): Include cache objects in the details (Defaults to False).
-
-        Returns:
-            pd.DataFrame: A DataFrame with detailed metadata for all objects in the data_store prefix.
-        """
-        return super().details(include_cache=include_cache)
-
-    def check(self, location: str) -> bool:
-        """Check if a DataFrame exists at the specified location
-
-        Args:
-            location (str): The location of the data to check.
-
-        Returns:
-            bool: True if the data exists, False otherwise.
-        """
-        return super().check(location)
-
-    def get(self, location: str) -> Union[pd.DataFrame, None]:
-        """Retrieve a DataFrame from AWS S3.
-
-        Args:
-            location (str): The location of the data to retrieve.
-
-        Returns:
-            pd.DataFrame: The retrieved DataFrame or None if not found.
-        """
-        _df = super().get(location)
-        if _df is None:
-            self.log.error(f"Dataframe not found at location: {location}")
-        return _df
-
-    def upsert(self, location: str, data: Union[pd.DataFrame, pd.Series]):
-        """Insert or update a DataFrame or Series in the AWS S3.
-
-        Args:
-            location (str): The location of the data.
-            data (Union[pd.DataFrame, pd.Series]): The data to be stored.
-        """
-        super().upsert(location, data)
-
-    def last_modified(self, location: str) -> Union[datetime, None]:
-        """Get the last modified date of the DataFrame at the specified location.
-
-        Args:
-            location (str): The location of the data to check.
-
-        Returns:
-            Union[datetime, None]: The last modified date of the DataFrame or None if not found.
-        """
-        return super().last_modified(location)
-
-    def delete(self, location: str):
-        """Delete a DataFrame from the AWS S3.
-
-        Args:
-            location (str): The location of the data to delete.
-        """
-        super().delete(location)
-
 
 if __name__ == "__main__":
     """Exercise the DFStore Class"""
     import time
+    import pandas as pd
 
     # Create a DFStore manager
     df_store = DFStore()

workbench/api/feature_set.py

@@ -58,10 +58,7 @@ class FeatureSet(FeatureSetCore):
             include_aws_columns (bool): Include the AWS columns in the DataFrame (default: False)
 
         Returns:
-            pd.DataFrame: A DataFrame of ALL the data from this FeatureSet
-
-        Note:
-            Obviously this is not recommended for large datasets :)
+            pd.DataFrame: A DataFrame of all the data from this FeatureSet up to the limit
         """
 
         # Get the table associated with the data
@@ -83,7 +80,7 @@ class FeatureSet(FeatureSetCore):
         tags: list = None,
         description: str = None,
         feature_list: list = None,
-        target_column: str = None,
+        target_column: Union[str, list[str]] = None,
         model_class: str = None,
         model_import_str: str = None,
         custom_script: Union[str, Path] = None,
@@ -103,7 +100,7 @@ class FeatureSet(FeatureSetCore):
             tags (list, optional): Set the tags for the model.  If not given tags will be generated.
             description (str, optional): Set the description for the model. If not give a description is generated.
             feature_list (list, optional): Set the feature list for the model. If not given a feature list is generated.
-            target_column (str, optional): The target column for the model (use None for unsupervised model)
+            target_column (str or list[str], optional): Target column(s) for the model (use None for unsupervised model)
             model_class (str, optional): Model class to use (e.g. "KMeans", default: None)
             model_import_str (str, optional): The import for the model (e.g. "from sklearn.cluster import KMeans")
             custom_script (str, optional): The custom script to use for the model (default: None)
@@ -157,6 +154,24 @@ class FeatureSet(FeatureSetCore):
         # Return the Model
         return Model(name)
 
+    def prox_model(self, target: str, features: list) -> "Proximity":  # noqa: F821
+        """Create a local Proximity Model for this Model
+
+        Args:
+           target (str): The target column name
+           features (list): The list of feature column names
+
+        Returns:
+           Proximity: A local Proximity Model
+        """
+        from workbench.algorithms.dataframe.proximity import Proximity  # noqa: F401 (avoid circular import)
+
+        # Create the Proximity Model from the full FeatureSet dataframe
+        full_df = self.pull_dataframe()
+
+        # Create and return the Proximity Model
+        return Proximity(full_df, self.id_column, features, target, track_columns=features)
+
 
 if __name__ == "__main__":
     """Exercise the FeatureSet Class"""
@@ -167,5 +182,24 @@ if __name__ == "__main__":
     pprint(my_features.summary())
     pprint(my_features.details())
 
+    # Pull the full DataFrame
+    df = my_features.pull_dataframe()
+    print(df.head())
+
+    # Create a Proximity Model from the FeatureSet
+    features = ["height", "weight", "age", "iq_score", "likes_dogs", "food"]
+    my_prox = my_features.prox_model(target="salary", features=features)
+    neighbors = my_prox.neighbors(42)
+    print("Neighbors for ID 42:")
+    print(neighbors)
+
     # Create a Model from the FeatureSet
-    my_model = my_features.to_model(name="test-model", model_type=ModelType.REGRESSOR, target_column="iq_score")
+    """
+    my_model = my_features.to_model(
+        name="test-model",
+        model_type=ModelType.REGRESSOR,
+        target_column="salary",
+        feature_list=features
+    )
+    pprint(my_model.summary())
+    """

workbench/api/parameter_store.py

@@ -1,13 +1,10 @@
 """ParameterStore: Manages Workbench parameters in a Cloud Based Parameter Store."""
 
-from typing import Union
-import logging
-
 # Workbench Imports
-from workbench.core.cloud_platform.aws.aws_parameter_store import AWSParameterStore
+from workbench.core.artifacts.parameter_store_core import ParameterStoreCore
 
 
-class ParameterStore(AWSParameterStore):
+class ParameterStore(ParameterStoreCore):
     """ParameterStore: Manages Workbench parameters in a Cloud Based Parameter Store.
 
     Common Usage:
@@ -43,56 +40,10 @@ class ParameterStore(AWSParameterStore):
 
     def __init__(self):
         """ParameterStore Init Method"""
-        self.log = logging.getLogger("workbench")
 
-        # Initialize the SuperClass
+        # Initialize parent class
         super().__init__()
 
-    def list(self, prefix: str = None) -> list:
-        """List all parameters in the AWS Parameter Store, optionally filtering by a prefix.
-
-        Args:
-            prefix (str, optional): A prefix to filter the parameters by. Defaults to None.
-
-        Returns:
-            list: A list of parameter names and details.
-        """
-        return super().list(prefix=prefix)
-
-    def get(self, name: str, warn: bool = True, decrypt: bool = True) -> Union[str, list, dict, None]:
-        """Retrieve a parameter value from the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter to retrieve.
-            warn (bool): Whether to log a warning if the parameter is not found.
-            decrypt (bool): Whether to decrypt secure string parameters.
-
-        Returns:
-            Union[str, list, dict, None]: The value of the parameter or None if not found.
-        """
-        return super().get(name=name, warn=warn, decrypt=decrypt)
-
-    def upsert(self, name: str, value):
-        """Insert or update a parameter in the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter.
-            value (str | list | dict): The value of the parameter.
-        """
-        super().upsert(name=name, value=value)
-
-    def delete(self, name: str):
-        """Delete a parameter from the AWS Parameter Store.
-
-        Args:
-            name (str): The name of the parameter to delete.
-        """
-        super().delete(name=name)
-
-    def __repr__(self):
-        """Return a string representation of the ParameterStore object."""
-        return super().__repr__()
-
 
 if __name__ == "__main__":
     """Exercise the ParameterStore Class"""

workbench/core/artifacts/artifact.py

@@ -8,8 +8,8 @@ from typing import Union
 
 # Workbench Imports
 from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
-from workbench.core.cloud_platform.aws.aws_parameter_store import AWSParameterStore as ParameterStore
-from workbench.core.cloud_platform.aws.aws_df_store import AWSDFStore as DFStore
+from workbench.core.artifacts.parameter_store_core import ParameterStoreCore
+from workbench.core.artifacts.df_store_core import DFStoreCore
 from workbench.utils.aws_utils import dict_to_aws_tags
 from workbench.utils.config_manager import ConfigManager, FatalConfigError
 from workbench.core.cloud_platform.cloud_meta import CloudMeta
@@ -48,11 +48,11 @@ class Artifact(ABC):
     tag_delimiter = "::"
 
     # Grab our Dataframe Cache Storage
-    df_cache = DFStore(path_prefix="/workbench/dataframe_cache")
+    df_cache = DFStoreCore(path_prefix="/workbench/dataframe_cache")
 
     # Artifact may want to use the Parameter Store or Dataframe Store
-    param_store = ParameterStore()
-    df_store = DFStore()
+    param_store = ParameterStoreCore()
+    df_store = DFStoreCore()
 
     def __init__(self, name: str, use_cached_meta: bool = False):
         """Initialize the Artifact Base Class

workbench/core/artifacts/df_store_core.py

@@ -0,0 +1,114 @@
+"""DFStoreCore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy"""
+
+import logging
+from typing import Union
+
+# Workbench Imports
+from workbench.utils.config_manager import ConfigManager
+from workbench.core.cloud_platform.aws.aws_account_clamp import AWSAccountClamp
+
+# Workbench Bridges Import
+from workbench_bridges.api import DFStore as BridgesDFStore
+
+
+class DFStoreCore(BridgesDFStore):
+    """DFStoreCore: Fast/efficient storage of DataFrames using AWS S3/Parquet/Snappy
+
+        Common Usage:
+    ```python
+            df_store = DFStoreCore()
+
+            # List Data
+            df_store.list()
+
+            # Add DataFrame
+            df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
+            df_store.upsert("/test/my_data", df)
+
+            # Retrieve DataFrame
+            df = df_store.get("/test/my_data")
+            print(df)
+
+            # Delete Data
+            df_store.delete("/test/my_data")
+    ```
+    """
+
+    def __init__(self, path_prefix: Union[str, None] = None):
+        """DFStoreCore Init Method
+
+        Args:
+            path_prefix (Union[str, None], optional): Add a path prefix to storage locations (Defaults to None)
+        """
+        # Get config from workbench's systems
+        bucket = ConfigManager().get_config("WORKBENCH_BUCKET")
+        session = AWSAccountClamp().boto3_session
+
+        # Initialize parent with workbench config
+        super().__init__(path_prefix=path_prefix, s3_bucket=bucket, boto3_session=session)
+        self.log = logging.getLogger("workbench")
+
+
+if __name__ == "__main__":
+    """Exercise the DFStoreCore Class"""
+    import time
+    import pandas as pd
+
+    # Create a DFStoreCore manager
+    df_store = DFStoreCore()
+
+    # Details of the Dataframe Store
+    print("Detailed Data...")
+    print(df_store.details())
+
+    # Add a new DataFrame
+    my_df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
+    df_store.upsert("/testing/test_data", my_df)
+
+    # Get the DataFrame
+    print(f"Getting data 'test_data':\n{df_store.get('/testing/test_data')}")
+
+    # Now let's test adding a Series
+    series = pd.Series([1, 2, 3, 4], name="Series")
+    df_store.upsert("/testing/test_series", series)
+    print(f"Getting data 'test_series':\n{df_store.get('/testing/test_series')}")
+
+    # Summary of the data
+    print("Summary Data...")
+    print(df_store.summary())
+
+    # Repr of the DFStoreCore object
+    print("DFStoreCore Object:")
+    print(df_store)
+
+    # Check if the data exists
+    print("Check if data exists...")
+    print(df_store.check("/testing/test_data"))
+    print(df_store.check("/testing/test_series"))
+
+    # Time the check
+    start_time = time.time()
+    print(df_store.check("/testing/test_data"))
+    print("--- Check %s seconds ---" % (time.time() - start_time))
+
+    # Now delete the test data
+    df_store.delete("/testing/test_data")
+    df_store.delete("/testing/test_series")
+
+    # Check if the data exists
+    print("Check if data exists...")
+    print(df_store.check("/testing/test_data"))
+    print(df_store.check("/testing/test_series"))
+
+    # Add a bunch of dataframes and then test recursive delete
+    for i in range(10):
+        df_store.upsert(f"/testing/data_{i}", pd.DataFrame({"A": [1, 2], "B": [3, 4]}))
+    print("Before Recursive Delete:")
+    print(df_store.summary())
+    df_store.delete_recursive("/testing")
+    print("After Recursive Delete:")
+    print(df_store.summary())
+
+    # Get a non-existent DataFrame
+    print("Getting non-existent data...")
+    print(df_store.get("/testing/no_where"))