arize 8.0.0b1__py3-none-any.whl → 8.0.0b4__py3-none-any.whl

arize/ml/bounded_executor.py

@@ -24,12 +24,26 @@ class BoundedExecutor:
         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."""
+        """Submit a callable to be executed with bounded concurrency.
+
+        This method blocks if the work queue is full (at the bound limit) until
+        space becomes available. Compatible with concurrent.futures.Executor.submit().
+
+        Args:
+            fn: The callable to execute.
+            *args: Positional arguments to pass to the callable.
+            **kwargs: Keyword arguments to pass to the callable.
+
+        Returns:
+            concurrent.futures.Future: A Future representing the pending execution.
+
+        Raises:
+            Exception: Any exception raised during submission is re-raised after
+                releasing the semaphore.
+        """
         self.semaphore.acquire()
         try:
             future = self.executor.submit(fn, *args, **kwargs)
@@ -40,8 +54,13 @@ class BoundedExecutor:
             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."""
+        """Shutdown the executor, optionally waiting for pending tasks to complete.
+
+        Compatible with concurrent.futures.Executor.shutdown().
+
+        Args:
+            wait: If True, blocks until all pending tasks complete. If False,
+                returns immediately without waiting. Defaults to True.
+        """
         self.executor.shutdown(wait)

arize/ml/casting.py

@@ -1,6 +1,5 @@
 """Type casting utilities for ML model data conversion."""
 
-# type: ignore[pb2]
 from __future__ import annotations
 
 import math
@@ -14,8 +13,8 @@ from arize.ml.types import (
     Schema,
     TypedColumns,
     TypedValue,
+    _normalize_column_names,
 )
-from arize.utils.types import is_list_of
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -25,7 +24,11 @@ class CastingError(Exception):
     """Raised when type casting fails for a value."""
 
     def __str__(self) -> str:
-        """Return a human-readable error message."""
+        """Return a human-readable error message.
+
+        Returns:
+            str: The formatted error message describing the casting failure.
+        """
         return self.error_message()
 
     def __init__(self, error_msg: str, typed_value: TypedValue) -> None:
@@ -39,7 +42,11 @@ class CastingError(Exception):
         self.typed_value = typed_value
 
     def error_message(self) -> str:
-        """Return the error message for this exception."""
+        """Return the error message for this exception.
+
+        Returns:
+            str: Detailed error message including the value, its type, target type, and failure reason.
+        """
         return (
             f"Failed to cast value {self.typed_value.value} of type {type(self.typed_value.value)} "
             f"to type {self.typed_value.type}. "
@@ -51,14 +58,18 @@ class ColumnCastingError(Exception):
     """Raised when type casting fails for a column."""
 
     def __str__(self) -> str:
-        """Return a human-readable error message."""
+        """Return a human-readable error message.
+
+        Returns:
+            str: The formatted error message describing the column casting failure.
+        """
         return self.error_message()
 
     def __init__(
         self,
         error_msg: str,
-        attempted_columns: str,
-        attempted_type: TypedColumns,
+        attempted_columns: list[str],
+        attempted_type: str,
     ) -> None:
         """Initialize the exception with column casting context.
 
@@ -72,7 +83,11 @@ class ColumnCastingError(Exception):
         self.attempted_casting_type = attempted_type
 
     def error_message(self) -> str:
-        """Return the error message for this exception."""
+        """Return the error message for this exception.
+
+        Returns:
+            str: Detailed error message including the target type, affected columns, and failure reason.
+        """
         return (
             f"Failed to cast to type {self.attempted_casting_type} "
             f"for columns: {log_a_list(self.attempted_casting_columns, 'and')}. "
@@ -84,7 +99,11 @@ class InvalidTypedColumnsError(Exception):
     """Raised when typed columns are invalid or incorrectly specified."""
 
     def __str__(self) -> str:
-        """Return a human-readable error message."""
+        """Return a human-readable error message.
+
+        Returns:
+            str: The formatted error message describing the invalid typed columns.
+        """
         return self.error_message()
 
     def __init__(self, field_name: str, reason: str) -> None:
@@ -98,7 +117,11 @@ class InvalidTypedColumnsError(Exception):
         self.reason = reason
 
     def error_message(self) -> str:
-        """Return the error message for this exception."""
+        """Return the error message for this exception.
+
+        Returns:
+            str: Error message describing which field has invalid typed columns and why.
+        """
         return f"The {self.field_name} TypedColumns object {self.reason}."
 
 
@@ -106,7 +129,11 @@ class InvalidSchemaFieldTypeError(Exception):
     """Raised when schema field has invalid or unexpected type."""
 
     def __str__(self) -> str:
-        """Return a human-readable error message."""
+        """Return a human-readable error message.
+
+        Returns:
+            str: The formatted error message describing the invalid schema field type.
+        """
         return self.error_message()
 
     def __init__(self, msg: str) -> None:
@@ -118,7 +145,11 @@ class InvalidSchemaFieldTypeError(Exception):
         self.msg = msg
 
     def error_message(self) -> str:
-        """Return the error message for this exception."""
+        """Return the error message for this exception.
+
+        Returns:
+            str: The error message describing the schema field type issue.
+        """
         return self.msg
 
 
@@ -132,12 +163,12 @@ def cast_typed_columns(
     a column across many SDK uploads.
 
     Args:
-        dataframe (pd.DataFrame): A deepcopy of the user's dataframe.
+        dataframe (:class:`pandas.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:
-        tuple[pd.DataFrame, Schema]: A tuple containing:
+        tuple[:class:`pandas.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.
@@ -290,12 +321,12 @@ def _cast_columns(
     (feature_column_names or tag_column_names)
 
     Args:
-        dataframe (pd.DataFrame): A deepcopy of the user's dataframe.
+        dataframe (:class:`pandas.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:
-        pd.DataFrame: The dataframe with columns cast to the specified types.
+        :class:`pandas.DataFrame`: The dataframe with columns cast to the specified types.
 
     Raises:
         ColumnCastingError: If casting fails.
@@ -350,12 +381,12 @@ def _cast_df(
     """Cast columns in a dataframe to the specified type.
 
     Args:
-        df (pd.DataFrame): A deepcopy of the user's dataframe.
+        df (:class:`pandas.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:
-        pd.DataFrame: The dataframe with columns cast to the specified types.
+        :class:`pandas.DataFrame`: The dataframe with columns cast to the specified types.
 
     Raises:
         Exception: If casting fails. Common exceptions raised by astype() are
@@ -381,23 +412,10 @@ def _convert_schema_field_types(
         Schema: A Schema, with feature and tag column names converted to the
             List[string] format expected in downstream validation.
     """
-    feature_column_names_list = (
+    feature_column_names_list = _normalize_column_names(
         schema.feature_column_names
-        if is_list_of(schema.feature_column_names, str)
-        else (
-            schema.feature_column_names.get_all_column_names()
-            if schema.feature_column_names
-            else []
-        )
-    )
-
-    tag_column_names_list = (
-        schema.tag_column_names
-        if is_list_of(schema.tag_column_names, str)
-        else schema.tag_column_names.get_all_column_names()
-        if schema.tag_column_names
-        else []
     )
+    tag_column_names_list = _normalize_column_names(schema.tag_column_names)
 
     schema_dict = {
         "feature_column_names": feature_column_names_list,

arize/ml/client.py

@@ -1,12 +1,11 @@
 """Client implementation for managing ML models in the Arize platform."""
 
-# type: ignore[pb2]
 from __future__ import annotations
 
 import copy
 import logging
 import time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, cast
 
 from arize._generated.protocol.rec import public_pb2 as pb2
 from arize._lazy import require
@@ -377,8 +376,12 @@ class MLModelsClient:
 
             if embedding_features or prompt or response:
                 # NOTE: Deep copy is necessary to avoid side effects on the original input dictionary
-                combined_embedding_features = (
-                    embedding_features.copy() if embedding_features else {}
+                combined_embedding_features: dict[str, str | Embedding] = (
+                    cast(
+                        "dict[str, str | Embedding]", embedding_features.copy()
+                    )
+                    if embedding_features
+                    else {}
                 )
                 # Map prompt as embedding features for generative models
                 if prompt is not None:
@@ -395,7 +398,7 @@ class MLModelsClient:
                 p.MergeFrom(embedding_feats)
 
             if tags or llm_run_metadata:
-                joined_tags = copy.deepcopy(tags)
+                joined_tags = copy.deepcopy(tags) if tags is not None else {}
                 if llm_run_metadata:
                     if llm_run_metadata.total_token_count is not None:
                         joined_tags[
@@ -522,7 +525,7 @@ class MLModelsClient:
             record=rec,
             headers=headers,
             timeout=timeout,
-            indexes=None,
+            indexes=None,  # type: ignore[arg-type]
         )
 
     def log(
@@ -542,7 +545,7 @@ class MLModelsClient:
         timeout: float | None = None,
         tmp_dir: str = "",
     ) -> requests.Response:
-        """Log a batch of model predictions and actuals to Arize from a pandas DataFrame.
+        """Log a batch of model predictions and actuals to Arize from a :class:`pandas.DataFrame`.
 
         This method uploads multiple records to Arize in a single batch operation using
         Apache Arrow format for efficient transfer. The dataframe structure is defined
@@ -554,8 +557,8 @@ class MLModelsClient:
             model_type: The type of model. Supported types: BINARY, MULTI_CLASS, REGRESSION,
                 RANKING, OBJECT_DETECTION. Note: GENERATIVE_LLM is not supported; use the
                 spans module instead.
-            dataframe: Pandas DataFrame containing the data to upload. Columns should
-                correspond to the schema field mappings.
+            dataframe (:class:`pandas.DataFrame`): Pandas DataFrame containing the data to
+                upload. Columns should correspond to the schema field mappings.
             schema: Schema object (Schema or CorpusSchema) that defines the mapping between
                 dataframe columns and Arize data fields (e.g., prediction_label_column_name,
                 feature_column_names, etc.).
@@ -668,8 +671,10 @@ class MLModelsClient:
         dataframe = remove_extraneous_columns(df=dataframe, schema=schema)
 
         # always validate pd.Category is not present, if yes, convert to string
+        # Type ignore: pandas.api.types.is_categorical_dtype exists but stubs may be incomplete
         has_cat_col = any(
-            ptypes.is_categorical_dtype(x) for x in dataframe.dtypes
+            ptypes.is_categorical_dtype(x)  # type: ignore[attr-defined]
+            for x in dataframe.dtypes
         )
         if has_cat_col:
             cat_cols = [
@@ -691,14 +696,15 @@ class MLModelsClient:
             from arize.ml.surrogate_explainer.mimic import Mimic
 
             logger.debug("Running surrogate_explainability.")
-            if schema.shap_values_column_names:
+            # Type ignore: schema typed as BaseSchema but runtime is Schema with these attrs
+            if schema.shap_values_column_names:  # type: ignore[attr-defined]
                 logger.info(
                     "surrogate_explainability=True has no effect "
                     "because shap_values_column_names is already specified in schema."
                 )
-            elif schema.feature_column_names is None or (
-                hasattr(schema.feature_column_names, "__len__")
-                and len(schema.feature_column_names) == 0
+            elif schema.feature_column_names is None or (  # type: ignore[attr-defined]
+                hasattr(schema.feature_column_names, "__len__")  # type: ignore[attr-defined]
+                and len(schema.feature_column_names) == 0  # type: ignore[attr-defined]
             ):
                 logger.info(
                     "surrogate_explainability=True has no effect "
@@ -706,7 +712,9 @@ class MLModelsClient:
                 )
             else:
                 dataframe, schema = Mimic.augment(
-                    df=dataframe, schema=schema, model_type=model_type
+                    df=dataframe,
+                    schema=schema,  # type: ignore[arg-type]
+                    model_type=model_type,
                 )
 
         # Convert to Arrow table
@@ -733,8 +741,8 @@ class MLModelsClient:
                 pyarrow_schema=pa_table.schema,
             )
             if errors:
-                for e in errors:
-                    logger.error(e)
+                for error in errors:
+                    logger.error(error)
                 raise ValidationFailure(errors)
         if validate:
             logger.debug("Performing values validation.")
@@ -745,8 +753,8 @@ class MLModelsClient:
                 model_type=model_type,
             )
             if errors:
-                for e in errors:
-                    logger.error(e)
+                for error in errors:
+                    logger.error(error)
                 raise ValidationFailure(errors)
 
         if isinstance(schema, Schema) and not schema.has_prediction_columns():
@@ -759,12 +767,12 @@ class MLModelsClient:
 
         if environment == Environments.CORPUS:
             proto_schema = _get_pb_schema_corpus(
-                schema=schema,
+                schema=schema,  # type: ignore[arg-type]
                 model_id=model_name,
             )
         else:
             proto_schema = _get_pb_schema(
-                schema=schema,
+                schema=schema,  # type: ignore[arg-type]
                 model_id=model_name,
                 model_version=model_version,
                 model_type=model_type,
@@ -811,10 +819,10 @@ class MLModelsClient:
         similarity_search_params: SimilaritySearchParams | None = None,
         stream_chunk_size: int | None = None,
     ) -> pd.DataFrame:
-        """Export model data from Arize to a pandas DataFrame.
+        """Export model data from Arize to a :class:`pandas.DataFrame`.
 
         Retrieves prediction and optional actual data for a model within a specified time
-        range and returns it as a pandas DataFrame for analysis.
+        range and returns it as a :class:`pandas.DataFrame` for analysis.
 
         Args:
             space_id: The space ID where the model resides.
@@ -835,8 +843,9 @@ class MLModelsClient:
             stream_chunk_size: Optional chunk size for streaming large result sets.
 
         Returns:
-            A pandas DataFrame containing the exported data with columns for predictions,
-            actuals (if requested), features, tags, timestamps, and other model metadata.
+            :class:`pandas.DataFrame`: A pandas DataFrame containing the exported data
+                with columns for predictions, actuals (if requested), features, tags,
+                timestamps, and other model metadata.
 
         Raises:
             RuntimeError: If the Flight client request fails or returns no response.
@@ -879,6 +888,7 @@ class MLModelsClient:
     def export_to_parquet(
         self,
         *,
+        path: str,
         space_id: str,
         model_name: str,
         environment: Environments,
@@ -891,13 +901,14 @@ class MLModelsClient:
         columns: list | None = None,
         similarity_search_params: SimilaritySearchParams | None = None,
         stream_chunk_size: int | None = None,
-    ) -> pd.DataFrame:
-        """Export model data from Arize to a Parquet file and return as DataFrame.
+    ) -> None:
+        """Export model data from Arize to a Parquet file.
 
         Retrieves prediction and optional actual data for a model within a specified time
-        range, saves it as a Parquet file, and returns it as a pandas DataFrame.
+        range and writes it directly to a Parquet file at the specified path.
 
         Args:
+            path: The file path where the Parquet file will be written.
             space_id: The space ID where the model resides.
             model_name: The name of the model to export data from.
             environment: The environment to export from (PRODUCTION, TRAINING, or VALIDATION).
@@ -915,16 +926,12 @@ class MLModelsClient:
                 filtering.
             stream_chunk_size: Optional chunk size for streaming large result sets.
 
-        Returns:
-            A pandas DataFrame containing the exported data. The data is also saved to a
-            Parquet file by the underlying export client.
-
         Raises:
             RuntimeError: If the Flight client request fails or returns no response.
 
         Notes:
             - Uses Apache Arrow Flight for efficient data transfer
-            - The Parquet file location is managed by the ArizeExportClient
+            - Data is written directly to the specified path as a Parquet file
             - Large exports may benefit from specifying stream_chunk_size
         """
         require(_BATCH_EXTRA, _BATCH_DEPS)
@@ -942,7 +949,8 @@ class MLModelsClient:
             exporter = ArizeExportClient(
                 flight_client=flight_client,
             )
-            return exporter.export_to_parquet(
+            exporter.export_to_parquet(
+                path=path,
                 space_id=space_id,
                 model_id=model_name,
                 environment=environment,
@@ -981,7 +989,7 @@ class MLModelsClient:
         headers: dict[str, str],
         timeout: float | None,
         indexes: tuple,
-    ) -> object:
+    ) -> cf.Future[Any]:
         """Post a record to Arize via async HTTP request with protobuf JSON serialization."""
         from google.protobuf.json_format import MessageToDict
 

arize/ml/proto.py

@@ -1,6 +1,5 @@
 """Protocol buffer utilities for ML model data serialization."""
 
-# type: ignore[pb2]
 from __future__ import annotations
 
 from google.protobuf.timestamp_pb2 import Timestamp
@@ -26,17 +25,19 @@ from arize.ml.types import (
 from arize.utils.types import is_list_of
 
 
-def get_pb_dictionary(d: dict[object, object] | None) -> dict[str, object]:
+def get_pb_dictionary(d: object | None) -> dict[str, object]:
     """Convert a dictionary to protobuf format with string keys and pb2.Value values.
 
     Args:
-        d: Dictionary to convert, or None.
+        d: Dictionary to convert, or :obj:`None`.
 
     Returns:
-        Dictionary with string keys and protobuf Value objects, or empty dict if input is None.
+        Dictionary with string keys and protobuf Value objects, or empty dict if input is :obj:`None`.
     """
     if d is None:
         return {}
+    if not isinstance(d, dict):
+        return {}
     # Takes a dictionary and
     # - casts the keys as strings
     # - turns the values of the dictionary to our proto values pb2.Value()
@@ -48,7 +49,7 @@ def get_pb_dictionary(d: dict[object, object] | None) -> dict[str, object]:
     return converted_dict
 
 
-def get_pb_value(name: str | int | float, value: pb2.Value) -> pb2.Value:
+def get_pb_value(name: object, value: pb2.Value) -> pb2.Value:
     """Convert a Python value to a protobuf Value object.
 
     Args:
@@ -56,7 +57,7 @@ def get_pb_value(name: str | int | float, value: pb2.Value) -> pb2.Value:
         value: The value to convert to protobuf format.
 
     Returns:
-        A pb2.Value protobuf object, or None if value cannot be converted.
+        A pb2.Value protobuf object, or :obj:`None` if value cannot be converted.
 
     Raises:
         TypeError: If value type is not supported.
@@ -114,7 +115,8 @@ def get_pb_label(
     Raises:
         ValueError: If model_type is not supported.
     """
-    value = convert_element(value)
+    # convert_element preserves value type but returns object for type safety
+    value = convert_element(value)  # type: ignore[assignment]
     if model_type in NUMERIC_MODEL_TYPES:
         return _get_numeric_pb_label(prediction_or_actual, value)
     if (
@@ -129,7 +131,7 @@ def get_pb_label(
     if model_type == ModelTypes.MULTI_CLASS:
         return _get_multi_class_pb_label(value)
     raise ValueError(
-        f"model_type must be one of: {[mt.prediction_or_actual for mt in ModelTypes]} "
+        f"model_type must be one of: {[mt.name for mt in ModelTypes]} "
         f"Got "
         f"{model_type} instead."
     )
@@ -139,10 +141,10 @@ def get_pb_timestamp(time_overwrite: int | None) -> object | None:
     """Convert a Unix timestamp to a protobuf Timestamp object.
 
     Args:
-        time_overwrite: Unix epoch time in seconds, or None.
+        time_overwrite: Unix epoch time in seconds, or :obj:`None`.
 
     Returns:
-        A protobuf Timestamp object, or None if input is None.
+        A protobuf Timestamp object, or :obj:`None` if input is :obj:`None`.
 
     Raises:
         TypeError: If time_overwrite is not an integer.
@@ -197,12 +199,12 @@ def get_pb_embedding(val: Embedding) -> pb2.Embedding:
 
 def _get_numeric_pb_label(
     prediction_or_actual: str,
-    value: int | float,
+    value: object,
 ) -> pb2.PredictionLabel | pb2.ActualLabel:
     if not isinstance(value, (int, float)):
         raise TypeError(
             f"Received {prediction_or_actual}_label = {value}, of type {type(value)}. "
-            + f"{[mt.prediction_or_actual for mt in NUMERIC_MODEL_TYPES]} models accept labels of "
+            + f"{[mt.name for mt in NUMERIC_MODEL_TYPES]} models accept labels of "
             f"type int or float"
         )
     if prediction_or_actual == "prediction":
@@ -214,7 +216,7 @@ def _get_numeric_pb_label(
 
 def _get_score_categorical_pb_label(
     prediction_or_actual: str,
-    value: bool | str | tuple[str, float],
+    value: object,
 ) -> pb2.PredictionLabel | pb2.ActualLabel:
     sc = pb2.ScoreCategorical()
     if isinstance(value, bool):
@@ -229,7 +231,7 @@ def _get_score_categorical_pb_label(
             raise TypeError(
                 f"Received {prediction_or_actual}_label = {value}, of type "
                 f"{type(value)}[{type(value[0])}, None]. "
-                f"{[mt.prediction_or_actual for mt in CATEGORICAL_MODEL_TYPES]} models accept "
+                f"{[mt.name for mt in CATEGORICAL_MODEL_TYPES]} models accept "
                 "values of type str, bool, or Tuple[str, float]"
             )
         if not isinstance(value[0], (bool, str)) or not isinstance(
@@ -238,7 +240,7 @@ def _get_score_categorical_pb_label(
             raise TypeError(
                 f"Received {prediction_or_actual}_label = {value}, of type "
                 f"{type(value)}[{type(value[0])}, {type(value[1])}]. "
-                f"{[mt.prediction_or_actual for mt in CATEGORICAL_MODEL_TYPES]} models accept "
+                f"{[mt.name for mt in CATEGORICAL_MODEL_TYPES]} models accept "
                 "values of type str, bool, or Tuple[str or bool, float]"
             )
         if isinstance(value[0], bool):
@@ -249,7 +251,7 @@ def _get_score_categorical_pb_label(
     else:
         raise TypeError(
             f"Received {prediction_or_actual}_label = {value}, of type {type(value)}. "
-            + f"{[mt.prediction_or_actual for mt in CATEGORICAL_MODEL_TYPES]} models accept values "
+            + f"{[mt.name for mt in CATEGORICAL_MODEL_TYPES]} models accept values "
             f"of type str, bool, int, float or Tuple[str, float]"
         )
     if prediction_or_actual == "prediction":
@@ -261,10 +263,7 @@ def _get_score_categorical_pb_label(
 
 def _get_cv_pb_label(
     prediction_or_actual: str,
-    value: ObjectDetectionLabel
-    | SemanticSegmentationLabel
-    | InstanceSegmentationPredictionLabel
-    | InstanceSegmentationActualLabel,
+    value: object,
 ) -> pb2.PredictionLabel | pb2.ActualLabel:
     if isinstance(value, ObjectDetectionLabel):
         return _get_object_detection_pb_label(prediction_or_actual, value)
@@ -429,7 +428,7 @@ def _get_instance_segmentation_actual_pb_label(
 
 
 def _get_ranking_pb_label(
-    value: RankingPredictionLabel | RankingActualLabel,
+    value: object,
 ) -> pb2.PredictionLabel | pb2.ActualLabel:
     if not isinstance(value, (RankingPredictionLabel, RankingActualLabel)):
         raise InvalidValueType(
@@ -460,7 +459,7 @@ def _get_ranking_pb_label(
 
 
 def _get_multi_class_pb_label(
-    value: MultiClassPredictionLabel | MultiClassActualLabel,
+    value: object,
 ) -> pb2.PredictionLabel | pb2.ActualLabel:
     if not isinstance(
         value, (MultiClassPredictionLabel, MultiClassActualLabel)