arize 8.0.0a22__py3-none-any.whl → 8.0.0b0__py3-none-any.whl

arize/ml/bounded_executor.py

@@ -0,0 +1,47 @@
+"""Bounded thread pool executor with queue size limits."""
+
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from threading import BoundedSemaphore
+
+
+class BoundedExecutor:
+    """BoundedExecutor behaves as a ThreadPoolExecutor which will block on calls to submit().
+
+    Blocks once the limit given as "bound" work items are queued for execution.
+
+    :param bound: Integer - the maximum number of items in the work queue
+    :param max_workers: Integer - the size of the thread pool
+    """
+
+    def __init__(self, bound: int, max_workers: int) -> None:
+        """Initialize the bounded executor.
+
+        Args:
+            bound: Maximum number of items in the work queue.
+            max_workers: Size of the thread pool.
+        """
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)
+        self.semaphore = BoundedSemaphore(bound + max_workers)
+
+    """See concurrent.futures.Executor#submit"""
+
+    def submit(
+        self, fn: Callable[..., object], *args: object, **kwargs: object
+    ) -> object:
+        """Submit a callable to be executed with bounded concurrency."""
+        self.semaphore.acquire()
+        try:
+            future = self.executor.submit(fn, *args, **kwargs)
+        except Exception:
+            self.semaphore.release()
+            raise
+        else:
+            future.add_done_callback(lambda _: self.semaphore.release())
+            return future
+
+    """See concurrent.futures.Executor#shutdown"""
+
+    def shutdown(self, wait: bool = True) -> None:
+        """Shutdown the executor, optionally waiting for pending tasks to complete."""
+        self.executor.shutdown(wait)

arize/{models → ml}/casting.py

@@ -1,27 +1,45 @@
+"""Type casting utilities for ML model data conversion."""
+
 # type: ignore[pb2]
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING, List, Tuple, Union
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from arize.logging import log_a_list
-from arize.types import ArizeTypes, Schema, TypedColumns, TypedValue, is_list_of
+from arize.ml.types import (
+    ArizeTypes,
+    Schema,
+    TypedColumns,
+    TypedValue,
+)
+from arize.utils.types import is_list_of
 
 if TYPE_CHECKING:
     import pandas as pd
 
 
 class CastingError(Exception):
+    """Raised when type casting fails for a value."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, error_msg: str, typed_value: TypedValue) -> None:
+        """Initialize the exception with type casting context.
+
+        Args:
+            error_msg: Description of the casting failure.
+            typed_value: The TypedValue that failed to cast.
+        """
         self.error_msg = error_msg
         self.typed_value = typed_value
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"Failed to cast value {self.typed_value.value} of type {type(self.typed_value.value)} "
             f"to type {self.typed_value.type}. "
@@ -30,7 +48,10 @@ class CastingError(Exception):
 
 
 class ColumnCastingError(Exception):
+    """Raised when type casting fails for a column."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(
@@ -39,11 +60,19 @@ class ColumnCastingError(Exception):
         attempted_columns: str,
         attempted_type: TypedColumns,
     ) -> None:
+        """Initialize the exception with column casting context.
+
+        Args:
+            error_msg: Description of the casting failure.
+            attempted_columns: Columns that failed to cast.
+            attempted_type: The TypedColumns type that was attempted.
+        """
         self.error_msg = error_msg
         self.attempted_casting_columns = attempted_columns
         self.attempted_casting_type = attempted_type
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return (
             f"Failed to cast to type {self.attempted_casting_type} "
             f"for columns: {log_a_list(self.attempted_casting_columns, 'and')}. "
@@ -52,60 +81,70 @@ class ColumnCastingError(Exception):
 
 
 class InvalidTypedColumnsError(Exception):
+    """Raised when typed columns are invalid or incorrectly specified."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, field_name: str, reason: str) -> None:
+        """Initialize the exception with typed columns validation context.
+
+        Args:
+            field_name: Name of the schema field with invalid typed columns.
+            reason: Description of why the typed columns are invalid.
+        """
         self.field_name = field_name
         self.reason = reason
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return f"The {self.field_name} TypedColumns object {self.reason}."
 
 
 class InvalidSchemaFieldTypeError(Exception):
+    """Raised when schema field has invalid or unexpected type."""
+
     def __str__(self) -> str:
+        """Return a human-readable error message."""
         return self.error_message()
 
     def __init__(self, msg: str) -> None:
+        """Initialize the exception with schema field type error message.
+
+        Args:
+            msg: Error message describing the schema field type issue.
+        """
         self.msg = msg
 
     def error_message(self) -> str:
+        """Return the error message for this exception."""
         return self.msg
 
 
 def cast_typed_columns(
     dataframe: pd.DataFrame,
     schema: Schema,
-) -> Tuple[pd.DataFrame, Schema]:
-    """
-    Cast feature and tag columns in the dataframe to the types specified in each TypedColumns config.
-    This optional feature provides a simple way for users to prevent
-    type drift within a column across many SDK uploads.
-
-    Arguments:
-    ---------
-        dataframe: pd.DataFrame
-            A deepcopy of the user's dataframe.
-        schema: Schema
-            The schema, which may include feature and tag column names
+) -> tuple[pd.DataFrame, Schema]:
+    """Cast feature and tag columns in the dataframe to the types specified in each TypedColumns config.
+
+    This optional feature provides a simple way for users to prevent type drift within
+    a column across many SDK uploads.
+
+    Args:
+        dataframe (pd.DataFrame): A deepcopy of the user's dataframe.
+        schema (Schema): The schema, which may include feature and tag column names
             in a TypedColumns object or a List[string].
 
     Returns:
-    -------
-        dataframe: pd.DataFrame
-            The dataframe, with columns cast to the specified types.
-        schema: Schema
-            A new Schema object, with feature and tag column names converted to the List[string] format
-            expected in downstream validation.
+        tuple[pd.DataFrame, Schema]: A tuple containing:
+            - dataframe: The dataframe, with columns cast to the specified types.
+            - schema: A new Schema object, with feature and tag column names converted
+                to the List[string] format expected in downstream validation.
 
     Raises:
-    ------
-        ColumnCastingError
-            If casting fails.
-        InvalidTypedColumnsError
-            If the TypedColumns object is invalid.
-
+        ColumnCastingError: If casting fails.
+        InvalidTypedColumnsError: If the TypedColumns object is invalid.
     """
     typed_column_fields = schema.typed_column_fields()
     feature_field = "feature_column_names"
@@ -120,7 +159,7 @@ def cast_typed_columns(
         )
 
     # Make sure no other schema fields have this type.
-    if any({f for f in typed_column_fields if f not in allowed_fields}):
+    if any(f for f in typed_column_fields if f not in allowed_fields):
         raise InvalidSchemaFieldTypeError(
             "Only the feature_column_names and tag_column_names Schema fields can be of type "
             "TypedColumns. Fields with type TypedColumns:"
@@ -130,10 +169,7 @@ def cast_typed_columns(
     for field_name in typed_column_fields:
         f = getattr(schema, field_name)
         if f:
-            try:
-                _validate_typed_columns(field_name, f)
-            except InvalidTypedColumnsError:
-                raise
+            _validate_typed_columns(field_name, f)
             dataframe = _cast_columns(dataframe, f)
 
     # Now that the dataframe values have been cast to the specified types:
@@ -144,6 +180,14 @@ def cast_typed_columns(
 
 
 def cast_dictionary(d: dict) -> dict:
+    """Cast TypedValue entries in a dictionary to their appropriate Python types.
+
+    Args:
+        d: Dictionary that may contain TypedValue objects as values.
+
+    Returns:
+        Dictionary with TypedValue objects cast to their native Python types.
+    """
     cast_dict = {}
     for k, v in d.items():
         if isinstance(v, TypedValue):
@@ -154,47 +198,38 @@ def cast_dictionary(d: dict) -> dict:
 
 def _cast_value(
     typed_value: TypedValue,
-) -> Union[str, int, float, List[str], None]:
-    """
-    Casts a TypedValue to its provided type, preserving all null values as None or float('nan').
+) -> str | int | float | list[str] | None:
+    """Casts a TypedValue to its provided type, preserving all null values as None or float('nan').
 
-    Arguments:
-    ---------
-    typed_value: TypedValue
-        The TypedValue to cast.
+    Args:
+        typed_value (TypedValue): The TypedValue to cast.
 
     Returns:
-    -------
-    Union[str, int, float, List[str], None]
-        The cast value.
+        str | int | float | list[str] | None: The cast value.
 
     Raises:
-    ------
-    CastingError
-        If the value cannot be cast to the provided type.
-
+        CastingError: If the value cannot be cast to the provided type.
     """
     if typed_value.value is None:
         return None
 
     if typed_value.type == ArizeTypes.FLOAT:
         return _cast_to_float(typed_value)
-    elif typed_value.type == ArizeTypes.INT:
+    if typed_value.type == ArizeTypes.INT:
         return _cast_to_int(typed_value)
-    elif typed_value.type == ArizeTypes.STR:
+    if typed_value.type == ArizeTypes.STR:
         return _cast_to_str(typed_value)
-    else:
-        raise CastingError("Unknown casting type", typed_value)
+    raise CastingError("Unknown casting type", typed_value)
 
 
-def _cast_to_float(typed_value: TypedValue) -> Union[float, None]:
+def _cast_to_float(typed_value: TypedValue) -> float | None:
     try:
         return float(typed_value.value)
     except Exception as e:
         raise CastingError(str(e), typed_value) from e
 
 
-def _cast_to_int(typed_value: TypedValue) -> Union[int, None]:
+def _cast_to_int(typed_value: TypedValue) -> int | None:
     # a NaN float can't be cast to an int. Proactively return None instead.
     if isinstance(typed_value.value, float) and math.isnan(typed_value.value):
         return None
@@ -214,7 +249,7 @@ def _cast_to_int(typed_value: TypedValue) -> Union[int, None]:
         raise CastingError(str(e), typed_value) from e
 
 
-def _cast_to_str(typed_value: TypedValue) -> Union[str, None]:
+def _cast_to_str(typed_value: TypedValue) -> str | None:
     # a NaN float can't be cast to a string. Proactively return None instead.
     if isinstance(typed_value.value, float) and math.isnan(typed_value.value):
         return None
@@ -227,21 +262,15 @@ def _cast_to_str(typed_value: TypedValue) -> Union[str, None]:
 def _validate_typed_columns(
     field_name: str, typed_columns: TypedColumns
 ) -> None:
-    """
-    Validate a TypedColumns object.
+    """Validate a TypedColumns object.
 
-    Arguments:
-    ---------
-        field_name: str
-            The name of the Schema field that the TypedColumns object is associated with.
-        typed_columns: TypedColumns
-            The TypedColumns object to validate.
+    Args:
+        field_name (str): The name of the Schema field that the TypedColumns object
+            is associated with.
+        typed_columns (TypedColumns): The TypedColumns object to validate.
 
     Raises:
-    ------
-        InvalidTypedColumnsError
-            If the TypedColumns object is invalid.
-
+        InvalidTypedColumnsError: If the TypedColumns object is invalid.
     """
     if typed_columns.is_empty():
         raise InvalidTypedColumnsError(field_name=field_name, reason="is empty")
@@ -256,28 +285,20 @@ def _validate_typed_columns(
 def _cast_columns(
     dataframe: pd.DataFrame, columns: TypedColumns
 ) -> pd.DataFrame:
-    """
-    Cast columns corresponding to a single TypedColumns object and a single Arize Schema field.
+    """Cast columns corresponding to a single TypedColumns object and a single Arize Schema field.
+
     (feature_column_names or tag_column_names)
 
-    Arguments:
-    ---------
-        dataframe: pd.DataFrame
-            A deepcopy of the user's dataframe.
-        columns: TypedColumns
-            The TypedColumns object, which specifies the columns to cast
-            (and/or to not cast) and their target types.
+    Args:
+        dataframe (pd.DataFrame): A deepcopy of the user's dataframe.
+        columns (TypedColumns): The TypedColumns object, which specifies the columns
+            to cast (and/or to not cast) and their target types.
 
     Returns:
-    -------
-        dataframe: pd.DataFrame
-            The dataframe with columns cast to the specified types.
+        pd.DataFrame: The dataframe with columns cast to the specified types.
 
     Raises:
-    ------
-        ColumnCastingError
-            If casting fails.
-
+        ColumnCastingError: If casting fails.
     """
     if columns.to_str:
         try:
@@ -324,52 +345,41 @@ def _cast_columns(
 
 
 def _cast_df(
-    df: pd.DataFrame, cols: List[str], target_type_str: str
+    df: pd.DataFrame, cols: list[str], target_type_str: str
 ) -> pd.DataFrame:
-    """
-    Arguments:
-    ---------
-        df: pd.DataFrame
-            A deepcopy of the user's dataframe.
-        cols: List[str]
-            The list of column names to cast.
-        target_type_str: str
-            The target type to cast to.
+    """Cast columns in a dataframe to the specified type.
+
+    Args:
+        df (pd.DataFrame): A deepcopy of the user's dataframe.
+        cols (list[str]): The list of column names to cast.
+        target_type_str (str): The target type to cast to.
 
     Returns:
-    -------
-        df: pd.DataFrame
-            The dataframe with columns cast to the specified types.
+        pd.DataFrame: The dataframe with columns cast to the specified types.
 
     Raises:
-    ------
-        Exception
-            If casting fails. Common exceptions raised by astype() are TypeError and ValueError.
-
+        Exception: If casting fails. Common exceptions raised by astype() are
+            TypeError and ValueError.
     """
     nan_mapping = {"nan": np.nan, "NaN": np.nan}
     df = df.replace(nan_mapping)
 
     # None or NaN-based values (including np.nan) are automatically converted to pandas pd.NA type
-    return df.astype({col: target_type_str for col in cols})
+    return df.astype(dict.fromkeys(cols, target_type_str))
 
 
 def _convert_schema_field_types(
     schema: Schema,
 ) -> Schema:
-    """
-    Arguments:
-    ---------
-        schema: Schema
-            The schema, which may include feature and tag column names
+    """Convert schema field types from TypedColumns to List[string] format.
+
+    Args:
+        schema (Schema): The schema, which may include feature and tag column names
             in a TypedColumns object or a List[string].
 
     Returns:
-    -------
-        schema: Schema
-            A Schema, with feature and tag column names
-            converted to the List[string] format expected in downstream validation.
-
+        Schema: A Schema, with feature and tag column names converted to the
+            List[string] format expected in downstream validation.
     """
     feature_column_names_list = (
         schema.feature_column_names