arize 8.0.0a13__tar.gz → 8.0.0a15__tar.gz

{arize-8.0.0a13 → arize-8.0.0a15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize
-Version: 8.0.0a13
+Version: 8.0.0a15
 Summary: A helper library to interact with Arize AI APIs
 Project-URL: Homepage, https://arize.com
 Project-URL: Documentation, https://docs.arize.com/arize
@@ -424,7 +424,7 @@ examples = [
 ]
 ```
 
-If the number of examples (rows in dataframe, items in list) is too large, the client SDK will try to send the data via Arrow Flight via gRPC for better performance. If you want to force the data transfer to HTTP you can use the `force_http` flag. The response is always a `Dataset` object.
+If the number of examples (rows in dataframe, items in list) is too large, the client SDK will try to send the data via Arrow Flight via gRPC for better performance. If you want to force the data transfer to HTTP you can use the `force_http` flag. The response is a `Dataset` object.
 
 ```python
 created_dataset = client.datasets.create(
@@ -434,11 +434,19 @@ created_dataset = client.datasets.create(
 )
 ```
 
+The `Dataset` object also counts with convenience method similar to `List***` objects:
+
+```python
+# Get the response as a dictionary
+dataset_dict = create_dataset.to_dict()
+# Get the response in JSON format
+dataset_dict = create_dataset.to_json()
+```
 
 
 ### Get Dataset by ID
 
-To get a dataset by its ID use `client.datasets.get()`, you can optionally also pass the version ID of a particular version of interest of the dataset.
+To get a dataset by its ID use `client.datasets.get()`, you can optionally also pass the version ID of a particular version of interest of the dataset. The returned type is `Dataset`.
 
 ```python
 dataset = client.datasets.get(

{arize-8.0.0a13 → arize-8.0.0a15}/README.md

@@ -362,7 +362,7 @@ examples = [
 ]
 ```
 
-If the number of examples (rows in dataframe, items in list) is too large, the client SDK will try to send the data via Arrow Flight via gRPC for better performance. If you want to force the data transfer to HTTP you can use the `force_http` flag. The response is always a `Dataset` object.
+If the number of examples (rows in dataframe, items in list) is too large, the client SDK will try to send the data via Arrow Flight via gRPC for better performance. If you want to force the data transfer to HTTP you can use the `force_http` flag. The response is a `Dataset` object.
 
 ```python
 created_dataset = client.datasets.create(
@@ -372,11 +372,19 @@ created_dataset = client.datasets.create(
 )
 ```
 
+The `Dataset` object also counts with convenience method similar to `List***` objects:
+
+```python
+# Get the response as a dictionary
+dataset_dict = create_dataset.to_dict()
+# Get the response in JSON format
+dataset_dict = create_dataset.to_json()
+```
 
 
 ### Get Dataset by ID
 
-To get a dataset by its ID use `client.datasets.get()`, you can optionally also pass the version ID of a particular version of interest of the dataset.
+To get a dataset by its ID use `client.datasets.get()`, you can optionally also pass the version ID of a particular version of interest of the dataset. The returned type is `Dataset`.
 
 ```python
 dataset = client.datasets.get(

{arize-8.0.0a13 → arize-8.0.0a15}/src/arize/_exporter/client.py

@@ -20,7 +20,6 @@ from arize._generated.protocol.flight import export_pb2
 from arize.logging import CtxAdapter
 from arize.types import Environments, SimilaritySearchParams
 from arize.utils.dataframe import reset_dataframe_index
-from arize.utils.proto import get_pb_flight_doput_request
 
 logger = logging.getLogger(__name__)
 
@@ -131,7 +130,7 @@ class ArizeExportClient:
         reset_dataframe_index(df)
         return df
 
-    def export_model_to_parquet(
+    def export_to_parquet(
         self,
         path: str,
         space_id: str,
@@ -285,7 +284,7 @@ class ArizeExportClient:
             end_time=Timestamp(seconds=int(end_time.timestamp())),
             filter_expression=where,
             similarity_search_params=(
-                get_pb_flight_doput_request(similarity_search_params)
+                _get_pb_similarity_search_params(similarity_search_params)
                 if similarity_search_params
                 else None
             ),
@@ -326,3 +325,19 @@ class ArizeExportClient:
             colour="#008000",
             unit=" row",
         )
+
+
+def _get_pb_similarity_search_params(
+    similarity_params: SimilaritySearchParams,
+) -> export_pb2.SimilaritySearchParams:
+    proto_params = export_pb2.SimilaritySearchParams()
+    proto_params.search_column_name = similarity_params.search_column_name
+    proto_params.threshold = similarity_params.threshold
+    for ref in similarity_params.references:
+        new_ref = proto_params.references.add()
+        new_ref.prediction_id = ref.prediction_id
+        new_ref.reference_column_name = ref.reference_column_name
+        if ref.prediction_timestamp:
+            new_ref.prediction_timestamp.FromDatetime(ref.prediction_timestamp)
+
+    return proto_params

{arize-8.0.0a13 → arize-8.0.0a15}/src/arize/_flight/client.py

@@ -179,7 +179,9 @@ class ArizeFlightClient:
                 return res
         except Exception as e:
             logger.exception(f"Error logging arrow table to Arize: {e}")
-                    raise RuntimeError(f"Error logging arrow table to Arize: {e}") from e
+            raise RuntimeError(
+                f"Error logging arrow table to Arize: {e}"
+            ) from e
 
     # ---------- dataset methods ----------
 
@@ -221,7 +223,9 @@ class ArizeFlightClient:
                 return res
         except Exception as e:
             logger.exception(f"Error logging arrow table to Arize: {e}")
-            raise RuntimeError(f"Error logging arrow table to Arize: {e}") from e
+            raise RuntimeError(
+                f"Error logging arrow table to Arize: {e}"
+            ) from e
 
 
 def append_to_pyarrow_metadata(

arize-8.0.0a15/src/arize/datasets/client.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+import pandas as pd
+import pyarrow as pa
+
+from arize._flight.client import ArizeFlightClient
+from arize.config import SDKConfiguration
+from arize.exceptions.base import INVALID_ARROW_CONVERSION_MSG
+
+logger = logging.getLogger(__name__)
+
+REST_LIMIT_DATASET_EXAMPLES = 3
+
+
+class DatasetsClient:
+    def __init__(self, sdk_config: SDKConfiguration):
+        self._sdk_config = sdk_config
+
+        # Import at runtime so it’s still lazy and extras-gated by the parent
+        from arize._generated import api_client as gen
+
+        # Use the shared generated client from the config
+        self._api = gen.DatasetsApi(self._sdk_config.get_generated_client())
+
+        # Forward methods to preserve exact runtime signatures/docs
+        self.list = self._api.datasets_list
+        self.get = self._api.datasets_get
+        self.delete = self._api.datasets_delete
+        self.list_examples = self._api.datasets_list_examples
+
+        # Custom methods
+        self.create = self._create_dataset
+
+    def _create_dataset(
+        self,
+        name: str,
+        space_id: str,
+        examples: List[Dict[str, Any]] | pd.DataFrame,
+        force_http: bool = False,
+    ):
+        if not isinstance(examples, (list, pd.DataFrame)):
+            raise TypeError(
+                "Examples must be a list of dicts or a pandas DataFrame"
+            )
+        if len(examples) <= REST_LIMIT_DATASET_EXAMPLES or force_http:
+            from arize._generated import api_client as gen
+
+            data = (
+                examples.to_dict(orient="records")
+                if isinstance(examples, pd.DataFrame)
+                else examples
+            )
+
+            body = gen.DatasetsCreateRequest(
+                name=name,
+                spaceId=space_id,
+                examples=data,
+            )
+            return self._api.datasets_create(datasets_create_request=body)
+
+        # If we have too many examples, try to convert to a dataframe
+        # and log via gRPC + flight
+        logger.info(
+            f"Uploading {len(examples)} examples via REST may be slow. "
+            "Trying to convert to DataFrame for more efficient upload via "
+            "gRPC + Flight."
+        )
+        data = (
+            pd.DataFrame(examples) if isinstance(examples, list) else examples
+        )
+        return self._create_dataset_via_flight(
+            name=name,
+            space_id=space_id,
+            examples=data,
+        )
+
+    def _create_dataset_via_flight(
+        self,
+        name: str,
+        space_id: str,
+        examples: pd.DataFrame,
+    ):
+        # Convert datetime columns to int64 (ms since epoch)
+        # TODO(Kiko): Missing validation block
+        # data = _convert_datetime_columns_to_int(data)
+        # df = self._set_default_columns_for_dataset(data)
+        # if convert_dict_to_json:
+        #     df = _convert_default_columns_to_json_str(df)
+        # df = _convert_boolean_columns_to_str(df)
+        # validation_errors = Validator.validate(df)
+        # validation_errors.extend(
+        #     Validator.validate_max_chunk_size(max_chunk_size)
+        # )
+        # if validation_errors:
+        #     raise RuntimeError(
+        #         [e.error_message() for e in validation_errors]
+        #     )
+
+        # Convert to Arrow table
+        try:
+            logger.debug("Converting data to Arrow format")
+            pa_table = pa.Table.from_pandas(examples)
+        except pa.ArrowInvalid as e:
+            logger.error(f"{INVALID_ARROW_CONVERSION_MSG}: {str(e)}")
+            raise pa.ArrowInvalid(
+                f"Error converting to Arrow format: {str(e)}"
+            ) from e
+        except Exception as e:
+            logger.error(f"Unexpected error creating Arrow table: {str(e)}")
+            raise
+
+        response = None
+        with ArizeFlightClient(
+            api_key=self._sdk_config.api_key,
+            host=self._sdk_config.flight_server_host,
+            port=self._sdk_config.flight_server_port,
+            scheme=self._sdk_config.flight_scheme,
+            request_verify=self._sdk_config.request_verify,
+        ) as flight_client:
+            try:
+                response = flight_client.create_dataset(
+                    space_id=space_id,
+                    dataset_name=name,
+                    pa_table=pa_table,
+                )
+            except Exception as e:
+                msg = f"Error during update request: {str(e)}"
+                logger.error(msg)
+                raise RuntimeError(msg) from e
+        if response is None:
+            # This should not happen with proper Flight client implementation,
+            # but we handle it defensively
+            msg = "No response received from flight server during update"
+            logger.error(msg)
+            raise RuntimeError(msg)
+        # The response from flightserver is the dataset ID. To return the dataset
+        # object we make a GET query
+        dataset = self.get(dataset_id=response)
+        return dataset

{arize-8.0.0a13/src/arize/utils → arize-8.0.0a15/src/arize/models}/casting.py

@@ -131,28 +131,28 @@ def cast_typed_columns(
         f = getattr(schema, field_name)
         if f:
             try:
-                validate_typed_columns(field_name, f)
+                _validate_typed_columns(field_name, f)
             except InvalidTypedColumnsError:
                 raise
-            dataframe = cast_columns(dataframe, f)
+            dataframe = _cast_columns(dataframe, f)
 
     # Now that the dataframe values have been cast to the specified types:
     # for downstream validation to work as expected,
     # feature & tag schema field types should be List[string] of column names.
     # Since Schema is a frozen class, we must construct a new instance.
-    return dataframe, convert_schema_field_types(schema)
+    return dataframe, _convert_schema_field_types(schema)
 
 
 def cast_dictionary(d: dict) -> dict:
     cast_dict = {}
     for k, v in d.items():
         if isinstance(v, TypedValue):
-            v = cast_value(v)
+            v = _cast_value(v)
         cast_dict[k] = v
     return cast_dict
 
 
-def cast_value(
+def _cast_value(
     typed_value: TypedValue,
 ) -> Union[str, int, float, List[str], None]:
     """
@@ -224,7 +224,7 @@ def _cast_to_str(typed_value: TypedValue) -> Union[str, None]:
         raise CastingError(str(e), typed_value) from e
 
 
-def validate_typed_columns(
+def _validate_typed_columns(
     field_name: str, typed_columns: TypedColumns
 ) -> None:
     """
@@ -253,7 +253,7 @@ def validate_typed_columns(
         )
 
 
-def cast_columns(
+def _cast_columns(
     dataframe: pd.DataFrame, columns: TypedColumns
 ) -> pd.DataFrame:
     """
@@ -288,7 +288,7 @@ def cast_columns(
             # uses pd.NA for missing values (when storage arg is not configured)
             # In the future, try out pd.convert_dtypes (new in pandas 2.0):
             # https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.convert_dtypes.html
-            dataframe = cast_df(dataframe, columns.to_str, "string")
+            dataframe = _cast_df(dataframe, columns.to_str, "string")
         except Exception as e:
             raise ColumnCastingError(
                 error_msg=str(e),
@@ -300,7 +300,7 @@ def cast_columns(
         # see https://pandas.pydata.org/docs/reference/api/pandas.Int64Dtype.html
         # uses pd.NA for missing values
         try:
-            dataframe = cast_df(dataframe, columns.to_int, "Int64")
+            dataframe = _cast_df(dataframe, columns.to_int, "Int64")
         except Exception as e:
             raise ColumnCastingError(
                 error_msg=str(e),
@@ -312,7 +312,7 @@ def cast_columns(
         # see https://pandas.pydata.org/docs/reference/api/pandas.Float64Dtype.html
         # uses pd.NA for missing values
         try:
-            dataframe = cast_df(dataframe, columns.to_float, "Float64")
+            dataframe = _cast_df(dataframe, columns.to_float, "Float64")
         except Exception as e:
             raise ColumnCastingError(
                 error_msg=str(e),
@@ -323,7 +323,7 @@ def cast_columns(
     return dataframe
 
 
-def cast_df(
+def _cast_df(
     df: pd.DataFrame, cols: List[str], target_type_str: str
 ) -> pd.DataFrame:
     """
@@ -354,7 +354,7 @@ def cast_df(
     return df.astype({col: target_type_str for col in cols})
 
 
-def convert_schema_field_types(
+def _convert_schema_field_types(
     schema: Schema,
 ) -> Schema:
     """