Skip to content

Inference Store

Inference Storage

InferenceStore manages captured inference results using AWS S3/Parquet/Snappy with Athena queries on top.

The InferenceStore class provides a queryable home for inference results across models. Use it to list the models that have captured inference, count rows, and run Athena queries against the stored predictions.

Workbench REPL

Experiment with the InferenceStore() class in the Workbench REPL.

InferenceStore: Manage Inference Results using AWS S3/Parquet/Snappy with Athena Queries

InferenceStore

InferenceStore: Manage Inference Results using AWS S3/Parquet/Snappy with Athena Queries

Common Usage 
inf_store = InferenceStore()

# List all Models in the Inference Store
inf_store.models()

# List the total number of rows in the Inference Store
inf_store.total_rows()
Source code in src/workbench/api/inference_store.py 
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
class InferenceStore:
    """InferenceStore: Manage Inference Results using AWS S3/Parquet/Snappy with Athena Queries

    Common Usage:
        ```python
        inf_store = InferenceStore()

        # List all Models in the Inference Store
        inf_store.models()

        # List the total number of rows in the Inference Store
        inf_store.total_rows()
        ```
    """

    def __init__(self, catalog_db: str = "inference_store", table_name: str = "inference_store"):
        """Initialize InferenceStore with S3 path and auto-register Glue/Athena table"""
        self.log = logging.getLogger("workbench")
        self.catalog_db = catalog_db
        self.table_name = table_name
        self.boto3_session = get_boto3_session()
        self.schema = ["id", "model", "pred_label", "pred_value", "tags", "meta", "timestamp"]

        # Ensure the Table exists
        if not wr.catalog.does_table_exist(self.catalog_db, self.table_name, self.boto3_session):
            self.log.warning(f"Table {self.table_name} does not exist in database {self.catalog_db}.")
            self.log.important(f"Creating table {self.table_name}...")
            self._create_empty_table()

    def add_inference_results(self, df: pd.DataFrame, schema_map: dict = None, meta_fields: list = None):
        """Add inference results to the Inference Store

        Args:
            df (pd.DataFrame): The DataFrame containing inference results.
            schema_map (dict, optional): A mapping of DataFrame columns to Inference Store schema columns.
            meta_fields (list, optional): Additional metadata fields to include in the Inference Store.
        """

        # Apply schema mapping if provided
        if schema_map:
            self.log.info(f"Applying schema mapping: {schema_map}")
            # Drop existing columns that would conflict with renamed columns
            df = df.drop(columns=[col for col in schema_map.values() if col in df.columns], errors="ignore")
            df = df.rename(columns=schema_map)

        # Check for existing meta data
        if "meta" in df.columns:
            # Check for any blanks or NaNs in the meta column
            self.log.info("Using existing 'meta' column for metadata.")
            df["meta"] = df["meta"].apply(lambda x: "{}" if pd.isna(x) or x == "" else x)
            if meta_fields:
                self.log.warning("Both 'meta' column and 'meta_fields' provided. Ignoring 'meta_fields'.")

        # Convert all meta fields to a combined JSON string and put into the 'meta' column
        elif meta_fields:
            self.log.info(f"Combining metadata fields: {meta_fields}")
            df["meta"] = df[meta_fields].apply(lambda row: row.to_json(), axis=1)
            df.drop(columns=meta_fields, inplace=True)
        else:
            df["meta"] = pd.Series("{}", index=df.index, dtype="string")

        # If "pred_label" or "pred_value" columns are missing, create them with default values
        if "pred_label" not in df.columns:
            df["pred_label"] = pd.Series([None] * len(df), dtype="string")
        if "pred_value" not in df.columns:
            df["pred_value"] = pd.Series([None] * len(df), dtype="float64")

        # Verify that we have all the schema columns
        missing_columns = set(self.schema) - set(df.columns)
        if missing_columns:
            raise ValueError(
                f"DataFrame is missing required columns: {missing_columns}. Expected schema: {self.schema}"
            )

        # Enforce proper data types for all columns
        df = self._enforce_schema_types(df)

        # Subset the DataFrame to only include the schema columns
        df = df[self.schema]

        # Add the results to the Inference Store
        self.log.info(f"Adding inference results to {self.catalog_db}.{self.table_name}")
        dataframe_to_table(df, self.catalog_db, self.table_name)
        self.log.info("Inference results added successfully.")

    def total_rows(self) -> int:
        """Return the total number of rows in the Inference Store

        Returns:
            int: The total number of rows in the Inference Store.
        """
        self.log.info(f"Retrieving total rows from {self.catalog_db}.{self.table_name}")
        df = self.query(f"SELECT COUNT(*) FROM {self.table_name}")
        return df.iloc[0, 0]

    def query(self, athena_query: str) -> pd.DataFrame:
        """Run a query against the Inference Store"""
        self.log.info(f"Running query: {athena_query}")
        start_time = time.time()

        try:
            df = wr.athena.read_sql_query(
                sql=athena_query,
                database=self.catalog_db,
                ctas_approach=False,
                boto3_session=self.boto3_session,
            )
            execution_time = time.time() - start_time
            self.log.info(f"Query completed in {execution_time:.2f} seconds")

            # Convert tags column from string to list
            if "tags" in df.columns:
                df["tags"] = df["tags"].str.strip("[]").str.split(", ")

            # Convert timestamp columns to UTC if they are naive
            for col in df.select_dtypes(include=["datetime64[ns]"]).columns:
                if df[col].dt.tz is None:
                    df[col] = df[col].dt.tz_localize("UTC")
                elif df[col].dt.tz.zone != "UTC":
                    df[col] = df[col].dt.tz_convert("UTC")

            return df
        except Exception as e:
            self.log.error(f"Failed to run query: {e}")
            return pd.DataFrame()

    def delete_all_data(self):
        """Delete all data in the Inference Store"""
        self.log.info(f"Deleting all data from {self.catalog_db}.{self.table_name}")
        delete_table(self.table_name, self.catalog_db)

    def _enforce_schema_types(self, df: pd.DataFrame) -> pd.DataFrame:
        """Enforce proper data types for all columns according to the schema

        Args:
            df (pd.DataFrame): The DataFrame to type-convert

        Returns:
            pd.DataFrame: DataFrame with properly typed columns
        """
        self.log.info("Enforcing schema types...")

        # Type enforcement mapping
        type_mapping = {
            "id": "string",
            "timestamp": "datetime64[ms]",
            "model": "string",
            "pred_label": "string",
            "pred_value": "float64",
            "tags": "object",  # Special handling for lists
            "meta": "string",
        }

        for col, dtype in type_mapping.items():
            if col in df.columns:
                try:
                    if col == "tags":
                        # Ensure tags is always a list, even if empty
                        df[col] = df[col].apply(lambda x: x if isinstance(x, list) else [] if pd.isna(x) else [str(x)])
                    elif col == "timestamp":
                        # Handle timestamp conversion more carefully
                        if not pd.api.types.is_datetime64_any_dtype(df[col]):
                            df[col] = pd.to_datetime(df[col])
                    else:
                        # Standard type conversion
                        df[col] = df[col].astype(dtype)

                except Exception as e:
                    raise ValueError(f"Failed to convert column '{col}' to type '{dtype}': {e}")

        # Return the DataFrame with enforced types
        return df

    def _create_empty_table(self):
        """Create an empty table directly using AWS Data Wrangler with proper column types"""
        self.log.important(f"Creating empty table {self.table_name} in database {self.catalog_db}")

        # Define the table schema with proper types
        table_schema = {
            "id": "string",
            "timestamp": "timestamp",
            "model": "string",
            "pred_label": "string",
            "pred_value": "double",
            "tags": "array<string>",
            "meta": "string",
        }
        s3_path = table_s3_path(database=self.catalog_db, table_name=self.table_name)
        wr.catalog.create_parquet_table(
            database=self.catalog_db,
            table=self.table_name,
            path=s3_path,
            columns_types=table_schema,
            boto3_session=self.boto3_session,
        )

    def __repr__(self):
        """Return a string representation of the InferenceStore object."""
        return f"InferenceStore(catalog_db={self.catalog_db}, table_name={self.table_name})"

__init__(catalog_db='inference_store', table_name='inference_store')

Initialize InferenceStore with S3 path and auto-register Glue/Athena table

Source code in src/workbench/api/inference_store.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
def __init__(self, catalog_db: str = "inference_store", table_name: str = "inference_store"):
    """Initialize InferenceStore with S3 path and auto-register Glue/Athena table"""
    self.log = logging.getLogger("workbench")
    self.catalog_db = catalog_db
    self.table_name = table_name
    self.boto3_session = get_boto3_session()
    self.schema = ["id", "model", "pred_label", "pred_value", "tags", "meta", "timestamp"]

    # Ensure the Table exists
    if not wr.catalog.does_table_exist(self.catalog_db, self.table_name, self.boto3_session):
        self.log.warning(f"Table {self.table_name} does not exist in database {self.catalog_db}.")
        self.log.important(f"Creating table {self.table_name}...")
        self._create_empty_table()

__repr__()

Return a string representation of the InferenceStore object.

Source code in src/workbench/api/inference_store.py 
213
214
215
def __repr__(self):
    """Return a string representation of the InferenceStore object."""
    return f"InferenceStore(catalog_db={self.catalog_db}, table_name={self.table_name})"

add_inference_results(df, schema_map=None, meta_fields=None)

Add inference results to the Inference Store

Parameters:

Name Type Description Default
df DataFrame

The DataFrame containing inference results.

 required
schema_map dict

A mapping of DataFrame columns to Inference Store schema columns.

 None
meta_fields list

Additional metadata fields to include in the Inference Store.

 None
Source code in src/workbench/api/inference_store.py 
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def add_inference_results(self, df: pd.DataFrame, schema_map: dict = None, meta_fields: list = None):
    """Add inference results to the Inference Store

    Args:
        df (pd.DataFrame): The DataFrame containing inference results.
        schema_map (dict, optional): A mapping of DataFrame columns to Inference Store schema columns.
        meta_fields (list, optional): Additional metadata fields to include in the Inference Store.
    """

    # Apply schema mapping if provided
    if schema_map:
        self.log.info(f"Applying schema mapping: {schema_map}")
        # Drop existing columns that would conflict with renamed columns
        df = df.drop(columns=[col for col in schema_map.values() if col in df.columns], errors="ignore")
        df = df.rename(columns=schema_map)

    # Check for existing meta data
    if "meta" in df.columns:
        # Check for any blanks or NaNs in the meta column
        self.log.info("Using existing 'meta' column for metadata.")
        df["meta"] = df["meta"].apply(lambda x: "{}" if pd.isna(x) or x == "" else x)
        if meta_fields:
            self.log.warning("Both 'meta' column and 'meta_fields' provided. Ignoring 'meta_fields'.")

    # Convert all meta fields to a combined JSON string and put into the 'meta' column
    elif meta_fields:
        self.log.info(f"Combining metadata fields: {meta_fields}")
        df["meta"] = df[meta_fields].apply(lambda row: row.to_json(), axis=1)
        df.drop(columns=meta_fields, inplace=True)
    else:
        df["meta"] = pd.Series("{}", index=df.index, dtype="string")

    # If "pred_label" or "pred_value" columns are missing, create them with default values
    if "pred_label" not in df.columns:
        df["pred_label"] = pd.Series([None] * len(df), dtype="string")
    if "pred_value" not in df.columns:
        df["pred_value"] = pd.Series([None] * len(df), dtype="float64")

    # Verify that we have all the schema columns
    missing_columns = set(self.schema) - set(df.columns)
    if missing_columns:
        raise ValueError(
            f"DataFrame is missing required columns: {missing_columns}. Expected schema: {self.schema}"
        )

    # Enforce proper data types for all columns
    df = self._enforce_schema_types(df)

    # Subset the DataFrame to only include the schema columns
    df = df[self.schema]

    # Add the results to the Inference Store
    self.log.info(f"Adding inference results to {self.catalog_db}.{self.table_name}")
    dataframe_to_table(df, self.catalog_db, self.table_name)
    self.log.info("Inference results added successfully.")

delete_all_data()

Delete all data in the Inference Store

Source code in src/workbench/api/inference_store.py 
143
144
145
146
def delete_all_data(self):
    """Delete all data in the Inference Store"""
    self.log.info(f"Deleting all data from {self.catalog_db}.{self.table_name}")
    delete_table(self.table_name, self.catalog_db)

query(athena_query)

Run a query against the Inference Store

Source code in src/workbench/api/inference_store.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def query(self, athena_query: str) -> pd.DataFrame:
    """Run a query against the Inference Store"""
    self.log.info(f"Running query: {athena_query}")
    start_time = time.time()

    try:
        df = wr.athena.read_sql_query(
            sql=athena_query,
            database=self.catalog_db,
            ctas_approach=False,
            boto3_session=self.boto3_session,
        )
        execution_time = time.time() - start_time
        self.log.info(f"Query completed in {execution_time:.2f} seconds")

        # Convert tags column from string to list
        if "tags" in df.columns:
            df["tags"] = df["tags"].str.strip("[]").str.split(", ")

        # Convert timestamp columns to UTC if they are naive
        for col in df.select_dtypes(include=["datetime64[ns]"]).columns:
            if df[col].dt.tz is None:
                df[col] = df[col].dt.tz_localize("UTC")
            elif df[col].dt.tz.zone != "UTC":
                df[col] = df[col].dt.tz_convert("UTC")

        return df
    except Exception as e:
        self.log.error(f"Failed to run query: {e}")
        return pd.DataFrame()

total_rows()

Return the total number of rows in the Inference Store

Returns:

Name Type Description
int int

The total number of rows in the Inference Store.
Source code in src/workbench/api/inference_store.py 
102
103
104
105
106
107
108
109
110
def total_rows(self) -> int:
    """Return the total number of rows in the Inference Store

    Returns:
        int: The total number of rows in the Inference Store.
    """
    self.log.info(f"Retrieving total rows from {self.catalog_db}.{self.table_name}")
    df = self.query(f"SELECT COUNT(*) FROM {self.table_name}")
    return df.iloc[0, 0]