arize 8.0.0a14__py3-none-any.whl → 8.0.0a16__py3-none-any.whl

arize/_generated/api_client/test/test_experiments_runs_list200_response.py

@@ -0,0 +1,56 @@
+# coding: utf-8
+
+"""
+    Arize REST API
+
+    API specification for the backend data server. The API is hosted globally at https://api.arize.com/v2 or in your own environment. You can access the OpenAPI spec for this API at https://api.arize.com/v2/spec.yaml 
+
+    The version of the OpenAPI document: 0.0.1
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+import unittest
+
+from arize._generated.api_client.models.experiments_runs_list200_response import ExperimentsRunsList200Response
+
+class TestExperimentsRunsList200Response(unittest.TestCase):
+    """ExperimentsRunsList200Response unit test stubs"""
+
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def make_instance(self, include_optional) -> ExperimentsRunsList200Response:
+        """Test ExperimentsRunsList200Response
+            include_optional is a boolean, when False only required
+            params are included, when True both required and
+            optional params are included """
+        # uncomment below to create an instance of `ExperimentsRunsList200Response`
+        """
+        model = ExperimentsRunsList200Response()
+        if include_optional:
+            return ExperimentsRunsList200Response(
+                experiment_runs = [
+                    { }
+                    ]
+            )
+        else:
+            return ExperimentsRunsList200Response(
+                experiment_runs = [
+                    { }
+                    ],
+        )
+        """
+
+    def testExperimentsRunsList200Response(self):
+        """Test ExperimentsRunsList200Response"""
+        # inst_req_only = self.make_instance(include_optional=False)
+        # inst_req_and_optional = self.make_instance(include_optional=True)
+
+if __name__ == '__main__':
+    unittest.main()

arize/_generated/api_client_README.md

@@ -1,5 +1,5 @@
 # Arize
-API specification for the backend data server. The API is hosted globally at https://app.arize.com/api/v1 or in your own environment. You can access the OpenAPI spec for this API at https://app.arize.com/api/v1/spec.yaml 
+API specification for the backend data server. The API is hosted globally at https://api.arize.com/v2 or in your own environment. You can access the OpenAPI spec for this API at https://api.arize.com/v2/spec.yaml 
 
 The `arize._generated.api_client` package is automatically generated by the [OpenAPI Generator](https://openapi-generator.tech) project:
 
@@ -74,13 +74,16 @@ All URIs are relative to *http://localhost*
 
 Class | Method | HTTP request | Description
 ------------ | ------------- | ------------- | -------------
-*DatasetsApi* | [**datasets_create**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_create) | **POST** /api/v1/datasets | Create a new dataset with JSON examples
-*DatasetsApi* | [**datasets_delete**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_delete) | **DELETE** /api/v1/datasets/{datasetId} | Delete a dataset by ID
-*DatasetsApi* | [**datasets_get**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_get) | **GET** /api/v1/datasets/{datasetId} | Get dataset by ID
-*DatasetsApi* | [**datasets_list**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_list) | **GET** /api/v1/datasets | List datasets the user has access to
-*DatasetsApi* | [**datasets_list_examples**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_list_examples) | **GET** /api/v1/datasets/{datasetId}/examples | List examples for a dataset
-*ExperimentsApi* | [**experiments_delete**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_delete) | **DELETE** /api/v1/experiments/{experimentId} | Delete an experiment by ID
-*ExperimentsApi* | [**experiments_list**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_list) | **GET** /api/v1/datasets/{datasetId}/experiments | List experiments for a given dataset
+*DatasetsApi* | [**datasets_create**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_create) | **POST** /v2/datasets | Create a new dataset with JSON examples
+*DatasetsApi* | [**datasets_delete**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_delete) | **DELETE** /v2/datasets/{datasetId} | Delete a dataset by ID
+*DatasetsApi* | [**datasets_get**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_get) | **GET** /v2/datasets/{datasetId} | Get dataset by ID
+*DatasetsApi* | [**datasets_list**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_list) | **GET** /v2/datasets | List datasets the user has access to
+*DatasetsApi* | [**datasets_list_examples**](arize/_generated/api_client/docs/DatasetsApi.md#datasets_list_examples) | **GET** /v2/datasets/{datasetId}/examples | List examples for a dataset
+*ExperimentsApi* | [**experiments_create**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_create) | **POST** /v2/experiments | Create a new experiment with run data
+*ExperimentsApi* | [**experiments_delete**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_delete) | **DELETE** /v2/experiments/{experimentId} | Delete an experiment by ID
+*ExperimentsApi* | [**experiments_get**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_get) | **GET** /v2/experiments/{experimentId} | Get experiment by ID
+*ExperimentsApi* | [**experiments_list**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_list) | **GET** /v2/experiments | List experiments
+*ExperimentsApi* | [**experiments_runs_list**](arize/_generated/api_client/docs/ExperimentsApi.md#experiments_runs_list) | **GET** /v2/experiments/{experimentId}/runs | List experiment runs for a specific experiment
 
 
 ## Documentation For Models
@@ -92,7 +95,9 @@ Class | Method | HTTP request | Description
  - [DatasetsListExamples200Response](arize/_generated/api_client/docs/DatasetsListExamples200Response.md)
  - [Error](arize/_generated/api_client/docs/Error.md)
  - [Experiment](arize/_generated/api_client/docs/Experiment.md)
+ - [ExperimentsCreateRequest](arize/_generated/api_client/docs/ExperimentsCreateRequest.md)
  - [ExperimentsList200Response](arize/_generated/api_client/docs/ExperimentsList200Response.md)
+ - [ExperimentsRunsList200Response](arize/_generated/api_client/docs/ExperimentsRunsList200Response.md)
 
 
 <a id="documentation-for-authorization"></a>

arize/client.py

@@ -12,6 +12,10 @@ if TYPE_CHECKING:
     from arize.spans.client import SpansClient
 
 
+# TODO(Kiko): Protobuf versioning is too old
+# TODO(Kiko): Make sure the client has same options as SDKConfiguration
+# TODO(Kiko): It does not make any sense to require space ID in run_experiment, dataset ID should suffice
+# TODO(Kiko): Should probably wrap every single method of gen clients so that we can add nice docstrings
 # TODO(Kiko): Add flight max_chunksize opt to write_table. In config?
 # TODO(Kiko): experimental/datasets must be adapted into the datasets subclient
 # TODO(Kiko): experimental/prompt hub is missing
@@ -65,11 +69,21 @@ class ArizeClient(LazySubclientsMixin):
         # Gate only the generated-backed ones
         "datasets": (
             "datasets-experiments",
-            ("pydantic",),
+            (
+                "pydantic",
+                "openinference.semconv",
+            ),
         ),
         "experiments": (
             "datasets-experiments",
-            ("pydantic",),
+            (
+                "pydantic",
+                "wrapt",
+                # "numpy",
+                # "openinference.semconv",
+                # "opentelemetry.sdk",
+                # "opentelemetry.exporter.otlp.proto.grpc.trace_exporter",
+            ),
         ),
         "spans": (
             "spans",
@@ -98,6 +112,7 @@ class ArizeClient(LazySubclientsMixin):
         flight_server_host: str | None = None,
         flight_server_port: int | None = None,
         flight_scheme: str | None = None,
+        pyarrow_max_chunksize: int | None = None,
         request_verify: bool | None = None,
         stream_max_workers: int | None = None,
         stream_max_queue_bound: int | None = None,
@@ -115,6 +130,8 @@ class ArizeClient(LazySubclientsMixin):
             cfg_kwargs["flight_server_port"] = flight_server_port
         if flight_scheme is not None:
             cfg_kwargs["flight_scheme"] = flight_scheme
+        if pyarrow_max_chunksize is not None:
+            cfg_kwargs["pyarrow_max_chunksize"] = pyarrow_max_chunksize
         if request_verify is not None:
             cfg_kwargs["request_verify"] = request_verify
         if stream_max_workers is not None:

arize/config.py

@@ -7,23 +7,30 @@ from typing import Any, Dict
 
 from arize.constants.config import (
     DEFAULT_API_HOST,
-    DEFAULT_API_INSECURE,
     DEFAULT_FLIGHT_HOST,
     DEFAULT_FLIGHT_PORT,
     DEFAULT_FLIGHT_TRANSPORT_SCHEME,
+    DEFAULT_INSECURE,
+    DEFAULT_MAX_HTTP_PAYLOAD_SIZE_MB,
+    DEFAULT_OTLP_HOST,
+    DEFAULT_PYARROW_MAX_CHUNKSIZE,
     DEFAULT_REQUEST_VERIFY,
     DEFAULT_STREAM_MAX_QUEUE_BOUND,
     DEFAULT_STREAM_MAX_WORKERS,
     ENV_API_HOST,
-    ENV_API_INSECURE,
     ENV_API_KEY,
     ENV_FLIGHT_HOST,
     ENV_FLIGHT_PORT,
     ENV_FLIGHT_TRANSPORT_SCHEME,
+    ENV_INSECURE,
+    ENV_MAX_HTTP_PAYLOAD_SIZE_MB,
+    ENV_OTLP_HOST,
+    ENV_PYARROW_MAX_CHUNKSIZE,
     ENV_REQUEST_VERIFY,
     ENV_STREAM_MAX_QUEUE_BOUND,
     ENV_STREAM_MAX_WORKERS,
 )
+from arize.constants.pyarrow import MAX_CHUNKSIZE
 from arize.exceptions.auth import MissingAPIKeyError
 from arize.version import __version__
 
@@ -45,7 +52,7 @@ def _api_host_factory() -> str:
 
 
 def _api_scheme_factory() -> str:
-    insecure = os.getenv(ENV_API_INSECURE, DEFAULT_API_INSECURE)
+    insecure = os.getenv(ENV_INSECURE, DEFAULT_INSECURE)
     if insecure:
         return "http"
     return "https"
@@ -65,6 +72,17 @@ def _flight_scheme_factory() -> str:
     )
 
 
+def _pyarrow_max_chunksize() -> int:
+    max_chunksize = int(
+        os.getenv(ENV_PYARROW_MAX_CHUNKSIZE, DEFAULT_PYARROW_MAX_CHUNKSIZE)
+    )
+    if max_chunksize <= 0 or max_chunksize > MAX_CHUNKSIZE:
+        raise ValueError(
+            f"Pyarrow max_chunksize must be between 1 and {MAX_CHUNKSIZE}, got {max_chunksize}"
+        )
+    return max_chunksize
+
+
 def _verify_factory() -> bool:
     return _parse_bool(os.getenv(ENV_REQUEST_VERIFY, DEFAULT_REQUEST_VERIFY))
 
@@ -79,6 +97,25 @@ def _stream_max_queue_bound_factory() -> int:
     )
 
 
+def _otlp_scheme_factory() -> str:
+    insecure = os.getenv(ENV_INSECURE, DEFAULT_INSECURE)
+    if insecure:
+        return "http"
+    return "https"
+
+
+def _otlp_host_factory() -> str:
+    return os.getenv(ENV_OTLP_HOST, DEFAULT_OTLP_HOST)
+
+
+def _max_http_payload_size_mb_factory() -> float:
+    return float(
+        os.getenv(
+            ENV_MAX_HTTP_PAYLOAD_SIZE_MB, DEFAULT_MAX_HTTP_PAYLOAD_SIZE_MB
+        )
+    )
+
+
 def _mask_secret(secret: str, N: int = 4) -> str:
     """Show first N chars then '***'; empty string if empty."""
     return f"{secret[:N]}***"
@@ -99,11 +136,17 @@ class SDKConfiguration:
     flight_server_host: str = field(default_factory=_flight_host_factory)
     flight_server_port: int = field(default_factory=_flight_port_factory)
     flight_scheme: str = field(default_factory=_flight_scheme_factory)
+    pyarrow_max_chunksize: int = field(default_factory=_pyarrow_max_chunksize)
     request_verify: bool = field(default_factory=_verify_factory)
     stream_max_workers: int = field(default_factory=_stream_max_workers_factory)
     stream_max_queue_bound: int = field(
         default_factory=_stream_max_queue_bound_factory
     )
+    otlp_host: str = field(default_factory=_otlp_host_factory)
+    otlp_scheme: str = field(default_factory=_otlp_scheme_factory)
+    max_http_payload_size_mb: float = field(
+        default_factory=_max_http_payload_size_mb_factory
+    )
 
     # Private, excluded from comparisons & repr
     _headers: Dict[str, str] = field(init=False, repr=False, compare=False)
@@ -121,6 +164,10 @@ class SDKConfiguration:
     def api_url(self) -> str:
         return _endpoint(self.api_scheme, self.api_host)
 
+    @property
+    def otlp_url(self) -> str:
+        return _endpoint(self.otlp_scheme, self.otlp_host, "/v1")
+
     @property
     def files_url(self) -> str:
         return _endpoint(self.api_scheme, self.api_host, "/v1/pandas_arrow")

arize/constants/config.py

@@ -3,19 +3,25 @@ ENV_API_KEY = "ARIZE_API_KEY"
 
 # Server configuration env vars
 ENV_API_HOST = "ARIZE_API_HOST"
-ENV_API_INSECURE = "ARIZE_API_INSECURE"
+ENV_OTLP_HOST = "ARIZE_OTLP_HOST"
 ENV_FLIGHT_HOST = "ARIZE_FLIGHT_HOST"
 ENV_FLIGHT_PORT = "ARIZE_FLIGHT_PORT"
 ENV_FLIGHT_TRANSPORT_SCHEME = "ARIZE_FLIGHT_TRANSPORT_SCHEME"
+ENV_PYARROW_MAX_CHUNKSIZE = "ARIZE_MAX_CHUNKSIZE"
 ENV_REQUEST_VERIFY = "ARIZE_REQUEST_VERIFY"
+ENV_INSECURE = "ARIZE_INSECURE"
+ENV_MAX_HTTP_PAYLOAD_SIZE_MB = "ARIZE_MAX_HTTP_PAYLOAD_SIZE_MB"
 
 # Server configuration default values
 DEFAULT_API_HOST = "api.arize.com"  # NOTE: Must not prefix with https://
-DEFAULT_API_INSECURE = False
+DEFAULT_OTLP_HOST = "otlp.arize.com"  # NOTE: Must not prefix with https://
 DEFAULT_FLIGHT_HOST = "flight.arize.com"  # NOTE: Must not prefix with https://
 DEFAULT_FLIGHT_PORT = 443
 DEFAULT_FLIGHT_TRANSPORT_SCHEME = "grpc+tls"
+DEFAULT_PYARROW_MAX_CHUNKSIZE = 10_000
 DEFAULT_REQUEST_VERIFY = True
+DEFAULT_INSECURE = False
+DEFAULT_MAX_HTTP_PAYLOAD_SIZE_MB = 100
 
 # ML Streaming configuration
 ENV_STREAM_MAX_WORKERS = "ARIZE_STREAM_MAX_WORKERS"

arize/constants/openinference.py

@@ -0,0 +1,14 @@
+import openinference.semconv.trace as oinf
+
+OPEN_INFERENCE_JSON_STR_TYPES = frozenset(
+    [
+        oinf.DocumentAttributes.DOCUMENT_METADATA,
+        oinf.SpanAttributes.LLM_FUNCTION_CALL,
+        oinf.SpanAttributes.LLM_INVOCATION_PARAMETERS,
+        oinf.SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES,
+        oinf.MessageAttributes.MESSAGE_FUNCTION_CALL_ARGUMENTS_JSON,
+        oinf.SpanAttributes.METADATA,
+        oinf.SpanAttributes.TOOL_PARAMETERS,
+        oinf.ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON,
+    ]
+)

arize/constants/pyarrow.py

@@ -0,0 +1 @@
+MAX_CHUNKSIZE = 100_000

arize/datasets/__init__.py

@@ -1,70 +0,0 @@
-from collections.abc import Mapping
-
-from arize._generated.api_client import models
-
-
-def make_to_df(field_name: str):
-    def to_df(
-        self,
-        by_alias: bool = False,
-        exclude_none: str | bool = False,
-        json_normalize: bool = False,
-        convert_dtypes: bool = True,
-    ):
-        """
-        Convert a list of objects to a pandas DataFrame.
-
-        Behavior:
-          - If an item is a Pydantic v2 model, use `.model_dump(by_alias=...)`.
-          - If an item is a mapping (dict-like), use it as-is.
-          - Otherwise, raise a ValueError (unsupported row type).
-
-        Parameters:
-          by_alias: Use field aliases when dumping Pydantic models.
-          exclude_none:
-            - False: keep Nones as-is
-            - "all": drop columns where *all* values are None/NaN
-            - "any": drop columns where *any* value is None/NaN
-            - True: alias for "all"
-          json_normalize: If True, flatten nested dicts via `pandas.json_normalize`.
-          convert_dtypes: If True, call `DataFrame.convert_dtypes()` at the end.
-
-        Returns:
-          pandas.DataFrame
-        """
-        import pandas as pd
-
-        items = getattr(self, field_name, []) or []
-
-        rows = []
-        for it in items:
-            if hasattr(it, "model_dump"):  # Pydantic v2 object
-                rows.append(it.model_dump(by_alias=by_alias))
-
-            elif isinstance(it, Mapping):  # Plain mapping
-                rows.append(it)
-            else:
-                raise ValueError(
-                    f"Cannot convert item of type {type(it)} to DataFrame row"
-                )
-
-        df = (
-            pd.json_normalize(rows, sep=".")
-            if json_normalize
-            else pd.DataFrame(rows)
-        )
-
-        # Drop None/NaN columns if requested
-        if exclude_none in ("any", "all", True):
-            drop_how = "all" if exclude_none is True else exclude_none
-            df.dropna(axis=1, how=drop_how, inplace=True)
-
-        if convert_dtypes:
-            df = df.convert_dtypes()
-        return df
-
-    return to_df
-
-
-models.DatasetsList200Response.to_df = make_to_df("datasets")  # type: ignore[attr-defined]
-models.DatasetsListExamples200Response.to_df = make_to_df("examples")  # type: ignore[attr-defined]

arize/datasets/client.py

@@ -1,18 +1,29 @@
 from __future__ import annotations
 
 import logging
+import time
+import uuid
 from typing import Any, Dict, List
 
 import pandas as pd
 import pyarrow as pa
 
 from arize._flight.client import ArizeFlightClient
+from arize._generated.api_client import models
 from arize.config import SDKConfiguration
+from arize.datasets.validation import validate_dataset_df
 from arize.exceptions.base import INVALID_ARROW_CONVERSION_MSG
+from arize.utils.openinference_conversion import (
+    convert_boolean_columns_to_str,
+    convert_datetime_columns_to_int,
+    convert_default_columns_to_json_str,
+)
+from arize.utils.size import get_payload_size_mb
 
 logger = logging.getLogger(__name__)
 
-REST_LIMIT_DATASET_EXAMPLES = 3
+# TODO(Kiko): Decide based on size of payload instead
+REST_LIMIT_DATASET_EXAMPLES = 0
 
 
 class DatasetsClient:
@@ -29,11 +40,60 @@ class DatasetsClient:
         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.list_examples = self._list_examples
         self.create = self._create_dataset
 
+    def _list_examples(
+        self,
+        dataset_id: str,
+        dataset_version_id: str = "",
+        limit: int = 100,
+        all: bool = False,
+    ):
+        if not all:
+            return self._api.datasets_list_examples(
+                dataset_id=dataset_id,
+                dataset_version_id=dataset_version_id,
+                limit=limit,
+            )
+
+        # TODO(Kiko): Space ID should not be needed,
+        # should work on server tech debt to remove this
+        dataset = self.get(dataset_id=dataset_id)
+        space_id = dataset.space_id
+
+        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,
+            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
+        ) as flight_client:
+            try:
+                response = flight_client.get_dataset_examples(
+                    space_id=space_id,
+                    dataset_id=dataset_id,
+                    dataset_version_id=dataset_version_id,
+                )
+            except Exception as e:
+                msg = f"Error during 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 request"
+            logger.error(msg)
+            raise RuntimeError(msg)
+        # The response from flightserver is the dataset ID. To return the dataset
+        # object we make a GET query
+        return models.DatasetsListExamples200Response(
+            examples=response.to_dict(orient="records")
+        )
+
     def _create_dataset(
         self,
         name: str,
@@ -45,7 +105,11 @@ class DatasetsClient:
             raise TypeError(
                 "Examples must be a list of dicts or a pandas DataFrame"
             )
-        if len(examples) <= REST_LIMIT_DATASET_EXAMPLES or force_http:
+        below_threshold = (
+            get_payload_size_mb(examples)
+            <= self._sdk_config.max_http_payload_size_mb
+        )
+        if below_threshold or force_http:
             from arize._generated import api_client as gen
 
             data = (
@@ -69,7 +133,9 @@ class DatasetsClient:
             "gRPC + Flight."
         )
         data = (
-            pd.DataFrame(examples) if isinstance(examples, list) else examples
+            examples
+            if isinstance(examples, pd.DataFrame)
+            else pd.DataFrame(examples)
         )
         return self._create_dataset_via_flight(
             name=name,
@@ -83,26 +149,21 @@ class DatasetsClient:
         space_id: str,
         examples: pd.DataFrame,
     ):
+        data = examples.copy()
         # 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]
-        #     )
+        data = convert_datetime_columns_to_int(data)
+        data = convert_boolean_columns_to_str(data)
+        data = _set_default_columns_for_dataset(data)
+        data = convert_default_columns_to_json_str(data)
+
+        validation_errors = validate_dataset_df(data)
+        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)
+            pa_table = pa.Table.from_pandas(data, preserve_index=False)
         except pa.ArrowInvalid as e:
             logger.error(f"{INVALID_ARROW_CONVERSION_MSG}: {str(e)}")
             raise pa.ArrowInvalid(
@@ -119,6 +180,7 @@ class DatasetsClient:
             port=self._sdk_config.flight_server_port,
             scheme=self._sdk_config.flight_scheme,
             request_verify=self._sdk_config.request_verify,
+            max_chunksize=self._sdk_config.pyarrow_max_chunksize,
         ) as flight_client:
             try:
                 response = flight_client.create_dataset(
@@ -140,3 +202,28 @@ class DatasetsClient:
         # object we make a GET query
         dataset = self.get(dataset_id=response)
         return dataset
+
+
+def _set_default_columns_for_dataset(df: pd.DataFrame) -> pd.DataFrame:
+    current_time = int(time.time() * 1000)
+    if "created_at" in df.columns:
+        if df["created_at"].isnull().values.any():
+            df["created_at"].fillna(current_time, inplace=True)
+    else:
+        df["created_at"] = current_time
+
+    if "updated_at" in df.columns:
+        if df["updated_at"].isnull().values.any():
+            df["updated_at"].fillna(current_time, inplace=True)
+    else:
+        df["updated_at"] = current_time
+
+    if "id" in df.columns:
+        if df["id"].isnull().values.any():
+            df["id"] = df["id"].apply(
+                lambda x: str(uuid.uuid4()) if pd.isnull(x) else x
+            )
+    else:
+        df["id"] = [str(uuid.uuid4()) for _ in range(len(df))]
+
+    return df

arize/datasets/errors.py

@@ -0,0 +1,61 @@
+from abc import ABC, abstractmethod
+
+
+class DatasetError(Exception, ABC):
+    def __str__(self) -> str:
+        return self.error_message()
+
+    @abstractmethod
+    def __repr__(self) -> str:
+        pass
+
+    @abstractmethod
+    def error_message(self) -> str:
+        pass
+
+
+class InvalidSessionError(DatasetError):
+    def error_message(self) -> str:
+        return (
+            "Credentials not provided or invalid. Please pass in the correct api_key when "
+            "initiating a new ArizeExportClient. Alternatively, you can set up credentials "
+            "in a profile or as an environment variable"
+        )
+
+    def __repr__(self) -> str:
+        return "InvalidSessionError()"
+
+
+class InvalidConfigFileError(DatasetError):
+    def error_message(self) -> str:
+        return "Invalid/Misconfigured Configuration File"
+
+    def __repr__(self) -> str:
+        return "InvalidConfigFileError()"
+
+
+class IDColumnUniqueConstraintError(DatasetError):
+    def error_message(self) -> str:
+        return "'id' column must contain unique values"
+
+    def __repr__(self) -> str:
+        return "IDColumnUniqueConstraintError()"
+
+
+class RequiredColumnsError(DatasetError):
+    def __init__(self, missing_columns: set) -> None:
+        self.missing_columns = missing_columns
+
+    def error_message(self) -> str:
+        return f"Missing required columns: {self.missing_columns}"
+
+    def __repr__(self) -> str:
+        return f"RequiredColumnsError({self.missing_columns})"
+
+
+class EmptyDatasetError(DatasetError):
+    def error_message(self) -> str:
+        return "DataFrame must have at least one row in it."
+
+    def __repr__(self) -> str:
+        return "EmptyDatasetError()"

arize/datasets/validation.py

@@ -0,0 +1,46 @@
+from typing import List
+
+import pandas as pd
+
+from arize.datasets import errors as err
+
+
+def validate_dataset_df(
+    df: pd.DataFrame,
+) -> List[err.DatasetError]:
+    ## check all require columns are present
+    required_columns_errors = _check_required_columns(df)
+    if required_columns_errors:
+        return required_columns_errors
+
+    ## check id column is unique
+    id_column_unique_constraint_error = _check_id_column_is_unique(df)
+    if id_column_unique_constraint_error:
+        return id_column_unique_constraint_error
+
+    # check DataFrame has at least one row in it
+    emtpy_dataframe_error = _check_empty_dataframe(df)
+    if emtpy_dataframe_error:
+        return emtpy_dataframe_error
+
+    return []
+
+
+def _check_required_columns(df: pd.DataFrame) -> List[err.DatasetError]:
+    required_columns = ["id", "created_at", "updated_at"]
+    missing_columns = set(required_columns) - set(df.columns)
+    if missing_columns:
+        return [err.RequiredColumnsError(missing_columns)]
+    return []
+
+
+def _check_id_column_is_unique(df: pd.DataFrame) -> List[err.DatasetError]:
+    if not df["id"].is_unique:
+        return [err.IDColumnUniqueConstraintError()]
+    return []
+
+
+def _check_empty_dataframe(df: pd.DataFrame) -> List[err.DatasetError]:
+    if df.empty:
+        return [err.EmptyDatasetError()]
+    return []