arize-phoenix 3.19.4__py3-none-any.whl → 3.21.0__py3-none-any.whl

phoenix/inferences/inferences.py

@@ -0,0 +1,730 @@
+import logging
+import re
+import uuid
+from copy import deepcopy
+from dataclasses import dataclass, fields, replace
+from enum import Enum
+from itertools import groupby
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+import numpy as np
+import pandas as pd
+from pandas import DataFrame, Series, Timestamp, read_parquet
+from pandas.api.types import (
+    is_numeric_dtype,
+)
+from typing_extensions import TypeAlias
+
+from phoenix.config import DATASET_DIR, GENERATED_DATASET_NAME_PREFIX
+from phoenix.datetime_utils import normalize_timestamps
+from phoenix.utilities.deprecation import deprecated
+
+from . import errors as err
+from .schema import (
+    LLM_SCHEMA_FIELD_NAMES,
+    MULTI_COLUMN_SCHEMA_FIELD_NAMES,
+    SINGLE_COLUMN_SCHEMA_FIELD_NAMES,
+    EmbeddingColumnNames,
+    EmbeddingFeatures,
+    RetrievalEmbeddingColumnNames,
+    Schema,
+    SchemaFieldName,
+    SchemaFieldValue,
+)
+from .validation import validate_dataset_inputs
+
+logger = logging.getLogger(__name__)
+
+# A schema like object. Not recommended to use this directly
+SchemaLike: TypeAlias = Any
+
+
+class Inferences:
+    """
+    A dataset to use for analysis using phoenix.
+    Used to construct a phoenix session via px.launch_app
+
+    Parameters
+    ----------
+    dataframe : pandas.DataFrame
+        The pandas dataframe containing the data to analyze
+    schema : phoenix.Schema
+        the schema of the dataset. Maps dataframe columns to the appropriate
+        model inference dimensions (features, predictions, actuals).
+    name : str, optional
+        The name of the dataset. If not provided, a random name will be generated.
+        Is helpful for identifying the dataset in the application.
+
+    Returns
+    -------
+    dataset : Dataset
+        The dataset object that can be used in a phoenix session
+
+    Examples
+    --------
+    >>> primary_dataset = px.Inferences(
+    >>>    dataframe=production_dataframe, schema=schema, name="primary"
+    >>> )
+    """
+
+    _data_file_name: str = "data.parquet"
+    _schema_file_name: str = "schema.json"
+    _is_persisted: bool = False
+    _is_empty: bool = False
+
+    def __init__(
+        self,
+        dataframe: DataFrame,
+        schema: Union[Schema, SchemaLike],
+        name: Optional[str] = None,
+    ):
+        # allow for schema like objects
+        if not isinstance(schema, Schema):
+            schema = _get_schema_from_unknown_schema_param(schema)
+        errors = validate_dataset_inputs(
+            dataframe=dataframe,
+            schema=schema,
+        )
+        if errors:
+            raise err.DatasetError(errors)
+        dataframe, schema = _parse_dataframe_and_schema(dataframe, schema)
+        dataframe, schema = _normalize_timestamps(
+            dataframe, schema, default_timestamp=Timestamp.utcnow()
+        )
+        dataframe = _sort_dataframe_rows_by_timestamp(dataframe, schema)
+        self.__dataframe: DataFrame = dataframe
+        self.__schema: Schema = schema
+        self.__name: str = (
+            name if name is not None else f"{GENERATED_DATASET_NAME_PREFIX}{str(uuid.uuid4())}"
+        )
+        self._is_empty = self.dataframe.empty
+        logger.info(f"""Dataset: {self.__name} initialized""")
+
+    def __repr__(self) -> str:
+        return f'<Dataset "{self.name}">'
+
+    @property
+    def dataframe(self) -> DataFrame:
+        return self.__dataframe
+
+    @property
+    def schema(self) -> "Schema":
+        return self.__schema
+
+    @property
+    def name(self) -> str:
+        return self.__name
+
+    @classmethod
+    def from_name(cls, name: str) -> "Inferences":
+        """Retrieves a dataset by name from the file system"""
+        directory = DATASET_DIR / name
+        df = read_parquet(directory / cls._data_file_name)
+        with open(directory / cls._schema_file_name) as schema_file:
+            schema_json = schema_file.read()
+        schema = Schema.from_json(schema_json)
+        return cls(df, schema, name)
+
+    def to_disc(self) -> None:
+        """writes the data and schema to disc"""
+        directory = DATASET_DIR / self.name
+        directory.mkdir(parents=True, exist_ok=True)
+        self.dataframe.to_parquet(
+            directory / self._data_file_name,
+            allow_truncated_timestamps=True,
+            coerce_timestamps="ms",
+        )
+        schema_json_data = self.schema.to_json()
+        with open(directory / self._schema_file_name, "w+") as schema_file:
+            schema_file.write(schema_json_data)
+
+    @classmethod
+    @deprecated("Inferences.from_open_inference is deprecated and will be removed.")
+    def from_open_inference(cls, dataframe: DataFrame) -> "Inferences":
+        schema = Schema()
+        column_renaming: Dict[str, str] = {}
+        for group_name, group in groupby(
+            sorted(
+                map(_parse_open_inference_column_name, dataframe.columns),
+                key=lambda column: column.name,
+            ),
+            key=lambda column: column.name,
+        ):
+            open_inference_columns = list(group)
+            if group_name == "":
+                column_names_by_category = {
+                    column.category: column.full_name for column in open_inference_columns
+                }
+                schema = replace(
+                    schema,
+                    prediction_id_column_name=column_names_by_category.get(
+                        OpenInferenceCategory.id
+                    ),
+                    timestamp_column_name=column_names_by_category.get(
+                        OpenInferenceCategory.timestamp
+                    ),
+                )
+                continue
+            column_names_by_specifier = {
+                column.specifier: column.full_name for column in open_inference_columns
+            }
+            if group_name == "response":
+                response_vector_column_name = column_names_by_specifier.get(
+                    OpenInferenceSpecifier.embedding
+                )
+                if response_vector_column_name is not None:
+                    column_renaming[response_vector_column_name] = "response"
+                    schema = replace(
+                        schema,
+                        response_column_names=EmbeddingColumnNames(
+                            vector_column_name=column_renaming[response_vector_column_name],
+                            raw_data_column_name=column_names_by_specifier.get(
+                                OpenInferenceSpecifier.default
+                            ),
+                        ),
+                    )
+                else:
+                    response_text_column_name = column_names_by_specifier.get(
+                        OpenInferenceSpecifier.default
+                    )
+                    if response_text_column_name is None:
+                        raise ValueError(
+                            "invalid OpenInference format: missing text column for response"
+                        )
+                    column_renaming[response_text_column_name] = "response"
+                    schema = replace(
+                        schema,
+                        response_column_names=column_renaming[response_text_column_name],
+                    )
+            elif group_name == "prompt":
+                prompt_vector_column_name = column_names_by_specifier.get(
+                    OpenInferenceSpecifier.embedding
+                )
+                if prompt_vector_column_name is None:
+                    raise ValueError(
+                        "invalid OpenInference format: missing embedding vector column for prompt"
+                    )
+                column_renaming[prompt_vector_column_name] = "prompt"
+                schema = replace(
+                    schema,
+                    prompt_column_names=RetrievalEmbeddingColumnNames(
+                        vector_column_name=column_renaming[prompt_vector_column_name],
+                        raw_data_column_name=column_names_by_specifier.get(
+                            OpenInferenceSpecifier.default
+                        ),
+                        context_retrieval_ids_column_name=column_names_by_specifier.get(
+                            OpenInferenceSpecifier.retrieved_document_ids
+                        ),
+                        context_retrieval_scores_column_name=column_names_by_specifier.get(
+                            OpenInferenceSpecifier.retrieved_document_scores
+                        ),
+                    ),
+                )
+            elif OpenInferenceSpecifier.embedding in column_names_by_specifier:
+                vector_column_name = column_names_by_specifier[OpenInferenceSpecifier.embedding]
+                column_renaming[vector_column_name] = group_name
+                embedding_feature_column_names = schema.embedding_feature_column_names or {}
+                embedding_feature_column_names.update(
+                    {
+                        group_name: EmbeddingColumnNames(
+                            vector_column_name=column_renaming[vector_column_name],
+                            raw_data_column_name=column_names_by_specifier.get(
+                                OpenInferenceSpecifier.raw_data
+                            ),
+                            link_to_data_column_name=column_names_by_specifier.get(
+                                OpenInferenceSpecifier.link_to_data
+                            ),
+                        )
+                    }
+                )
+                schema = replace(
+                    schema,
+                    embedding_feature_column_names=embedding_feature_column_names,
+                )
+            elif len(open_inference_columns) == 1:
+                open_inference_column = open_inference_columns[0]
+                raw_column_name = open_inference_column.full_name
+                column_renaming[raw_column_name] = open_inference_column.name
+                if open_inference_column.category is OpenInferenceCategory.feature:
+                    schema = replace(
+                        schema,
+                        feature_column_names=(
+                            (schema.feature_column_names or []) + [column_renaming[raw_column_name]]
+                        ),
+                    )
+                elif open_inference_column.category is OpenInferenceCategory.tag:
+                    schema = replace(
+                        schema,
+                        tag_column_names=(
+                            (schema.tag_column_names or []) + [column_renaming[raw_column_name]]
+                        ),
+                    )
+                elif open_inference_column.category is OpenInferenceCategory.prediction:
+                    if open_inference_column.specifier is OpenInferenceSpecifier.score:
+                        schema = replace(
+                            schema,
+                            prediction_score_column_name=column_renaming[raw_column_name],
+                        )
+                    if open_inference_column.specifier is OpenInferenceSpecifier.label:
+                        schema = replace(
+                            schema,
+                            prediction_label_column_name=column_renaming[raw_column_name],
+                        )
+                elif open_inference_column.category is OpenInferenceCategory.actual:
+                    if open_inference_column.specifier is OpenInferenceSpecifier.score:
+                        schema = replace(
+                            schema,
+                            actual_score_column_name=column_renaming[raw_column_name],
+                        )
+                    if open_inference_column.specifier is OpenInferenceSpecifier.label:
+                        schema = replace(
+                            schema,
+                            actual_label_column_name=column_renaming[raw_column_name],
+                        )
+            else:
+                raise ValueError(f"invalid OpenInference format: duplicated name `{group_name}`")
+
+        return cls(
+            dataframe.rename(
+                column_renaming,
+                axis=1,
+                copy=False,
+            ),
+            schema,
+        )
+
+
+class OpenInferenceCategory(Enum):
+    id = "id"
+    timestamp = "timestamp"
+    feature = "feature"
+    tag = "tag"
+    prediction = "prediction"
+    actual = "actual"
+
+
+class OpenInferenceSpecifier(Enum):
+    default = ""
+    score = "score"
+    label = "label"
+    embedding = "embedding"
+    raw_data = "raw_data"
+    link_to_data = "link_to_data"
+    retrieved_document_ids = "retrieved_document_ids"
+    retrieved_document_scores = "retrieved_document_scores"
+
+
+@dataclass(frozen=True)
+class _OpenInferenceColumnName:
+    full_name: str
+    category: OpenInferenceCategory
+    data_type: str
+    specifier: OpenInferenceSpecifier = OpenInferenceSpecifier.default
+    name: str = ""
+
+
+def _parse_open_inference_column_name(column_name: str) -> _OpenInferenceColumnName:
+    pattern = (
+        r"^:(?P<category>\w+)\.(?P<data_type>\[\w+\]|\w+)(\.(?P<specifier>\w+))?:(?P<name>.*)?$"
+    )
+    if match := re.match(pattern, column_name):
+        extract = match.groupdict(default="")
+        return _OpenInferenceColumnName(
+            full_name=column_name,
+            category=OpenInferenceCategory(extract.get("category", "").lower()),
+            data_type=extract.get("data_type", "").lower(),
+            specifier=OpenInferenceSpecifier(extract.get("specifier", "").lower()),
+            name=extract.get("name", ""),
+        )
+    raise ValueError(f"Invalid format for column name: {column_name}")
+
+
+def _parse_dataframe_and_schema(dataframe: DataFrame, schema: Schema) -> Tuple[DataFrame, Schema]:
+    """
+    Parses a dataframe according to a schema, infers feature columns names when
+    they are not explicitly provided, and removes excluded column names from
+    both dataframe and schema.
+
+    Removes column names in `schema.excluded_column_names` from the input dataframe and schema. To
+    remove an embedding feature and all associated columns, add the name of the embedding feature to
+    `schema.excluded_column_names` rather than the associated column names. If
+    `schema.feature_column_names` is `None`, automatically discovers features by adding all column
+    names present in the dataframe but not included in any other schema fields.
+    """
+
+    unseen_excluded_column_names: Set[str] = (
+        set(schema.excluded_column_names) if schema.excluded_column_names is not None else set()
+    )
+    unseen_column_names: Set[str] = set(dataframe.columns.to_list())
+    column_name_to_include: Dict[str, bool] = {}
+    schema_patch: Dict[SchemaFieldName, SchemaFieldValue] = {}
+
+    for schema_field_name in SINGLE_COLUMN_SCHEMA_FIELD_NAMES:
+        _check_single_column_schema_field_for_excluded_columns(
+            schema,
+            schema_field_name,
+            unseen_excluded_column_names,
+            schema_patch,
+            column_name_to_include,
+            unseen_column_names,
+        )
+
+    for schema_field_name in MULTI_COLUMN_SCHEMA_FIELD_NAMES:
+        _check_multi_column_schema_field_for_excluded_columns(
+            schema,
+            schema_field_name,
+            unseen_excluded_column_names,
+            schema_patch,
+            column_name_to_include,
+            unseen_column_names,
+        )
+
+    if schema.embedding_feature_column_names:
+        _check_embedding_features_schema_field_for_excluded_columns(
+            schema.embedding_feature_column_names,
+            unseen_excluded_column_names,
+            schema_patch,
+            column_name_to_include,
+            unseen_column_names,
+        )
+
+    for llm_schema_field_name in LLM_SCHEMA_FIELD_NAMES:
+        embedding_column_name_mapping = getattr(schema, llm_schema_field_name)
+        if isinstance(embedding_column_name_mapping, EmbeddingColumnNames):
+            _check_embedding_column_names_for_excluded_columns(
+                embedding_column_name_mapping,
+                column_name_to_include,
+                unseen_column_names,
+            )
+
+    if not schema.feature_column_names and unseen_column_names:
+        _discover_feature_columns(
+            dataframe,
+            unseen_excluded_column_names,
+            schema_patch,
+            column_name_to_include,
+            unseen_column_names,
+        )
+
+    if unseen_excluded_column_names:
+        logger.warning(
+            "The following columns and embedding features were excluded in the schema but were "
+            "not found in the dataframe: {}".format(", ".join(unseen_excluded_column_names))
+        )
+
+    parsed_dataframe, parsed_schema = _create_and_normalize_dataframe_and_schema(
+        dataframe, schema, schema_patch, column_name_to_include
+    )
+
+    return parsed_dataframe, parsed_schema
+
+
+def _check_single_column_schema_field_for_excluded_columns(
+    schema: Schema,
+    schema_field_name: str,
+    unseen_excluded_column_names: Set[str],
+    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: Dict[str, bool],
+    unseen_column_names: Set[str],
+) -> None:
+    """
+    Checks single-column schema fields for excluded column names.
+    """
+    column_name: str = getattr(schema, schema_field_name)
+    include_column: bool = column_name not in unseen_excluded_column_names
+    column_name_to_include[column_name] = include_column
+    if not include_column:
+        schema_patch[schema_field_name] = None
+        unseen_excluded_column_names.discard(column_name)
+        logger.debug(f"excluded {schema_field_name}: {column_name}")
+    unseen_column_names.discard(column_name)
+
+
+def _check_multi_column_schema_field_for_excluded_columns(
+    schema: Schema,
+    schema_field_name: str,
+    unseen_excluded_column_names: Set[str],
+    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: Dict[str, bool],
+    unseen_column_names: Set[str],
+) -> None:
+    """
+    Checks multi-column schema fields for excluded columns names.
+    """
+    column_names: Optional[List[str]] = getattr(schema, schema_field_name)
+    if column_names:
+        included_column_names: List[str] = []
+        excluded_column_names: List[str] = []
+        for column_name in column_names:
+            is_included_column = column_name not in unseen_excluded_column_names
+            column_name_to_include[column_name] = is_included_column
+            if is_included_column:
+                included_column_names.append(column_name)
+            else:
+                excluded_column_names.append(column_name)
+                unseen_excluded_column_names.discard(column_name)
+                logger.debug(f"excluded {schema_field_name}: {column_name}")
+            unseen_column_names.discard(column_name)
+        schema_patch[schema_field_name] = included_column_names if included_column_names else None
+
+
+def _check_embedding_features_schema_field_for_excluded_columns(
+    embedding_features: EmbeddingFeatures,
+    unseen_excluded_column_names: Set[str],
+    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: Dict[str, bool],
+    unseen_column_names: Set[str],
+) -> None:
+    """
+    Check embedding features for excluded column names.
+    """
+    included_embedding_features: EmbeddingFeatures = {}
+    for (
+        embedding_feature_name,
+        embedding_column_name_mapping,
+    ) in embedding_features.items():
+        include_embedding_feature = embedding_feature_name not in unseen_excluded_column_names
+        if include_embedding_feature:
+            included_embedding_features[embedding_feature_name] = deepcopy(
+                embedding_column_name_mapping
+            )
+        else:
+            unseen_excluded_column_names.discard(embedding_feature_name)
+
+        for embedding_field in fields(embedding_column_name_mapping):
+            column_name: Optional[str] = getattr(
+                embedding_column_name_mapping, embedding_field.name
+            )
+            if column_name is not None:
+                column_name_to_include[column_name] = include_embedding_feature
+                if (
+                    column_name != embedding_feature_name
+                    and column_name in unseen_excluded_column_names
+                ):
+                    logger.warning(
+                        f"Excluding embedding feature columns such as "
+                        f'"{column_name}" has no effect; instead exclude the '
+                        f'corresponding embedding feature name "{embedding_feature_name}".'
+                    )
+                    unseen_excluded_column_names.discard(column_name)
+                unseen_column_names.discard(column_name)
+    schema_patch["embedding_feature_column_names"] = (
+        included_embedding_features if included_embedding_features else None
+    )
+
+
+def _check_embedding_column_names_for_excluded_columns(
+    embedding_column_name_mapping: EmbeddingColumnNames,
+    column_name_to_include: Dict[str, bool],
+    unseen_column_names: Set[str],
+) -> None:
+    """
+    Check embedding column names for excluded column names.
+    """
+    for embedding_field in fields(embedding_column_name_mapping):
+        column_name: Optional[str] = getattr(embedding_column_name_mapping, embedding_field.name)
+        if column_name is not None:
+            column_name_to_include[column_name] = True
+            unseen_column_names.discard(column_name)
+
+
+def _discover_feature_columns(
+    dataframe: DataFrame,
+    unseen_excluded_column_names: Set[str],
+    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: Dict[str, bool],
+    unseen_column_names: Set[str],
+) -> None:
+    """
+    Adds unseen and un-excluded columns as features, with the exception of "prediction_id"
+    which is reserved
+    """
+    discovered_feature_column_names = []
+    for column_name in unseen_column_names:
+        if column_name not in unseen_excluded_column_names and column_name != "prediction_id":
+            discovered_feature_column_names.append(column_name)
+            column_name_to_include[column_name] = True
+        else:
+            unseen_excluded_column_names.discard(column_name)
+            logger.debug(f"excluded feature: {column_name}")
+    original_column_positions: List[int] = dataframe.columns.get_indexer(
+        discovered_feature_column_names
+    )  # type: ignore
+    feature_column_name_to_position: Dict[str, int] = dict(
+        zip(discovered_feature_column_names, original_column_positions)
+    )
+    discovered_feature_column_names.sort(key=lambda col: feature_column_name_to_position[col])
+    schema_patch["feature_column_names"] = discovered_feature_column_names
+    logger.debug(
+        "Discovered feature column names: {}".format(", ".join(discovered_feature_column_names))
+    )
+
+
+def _create_and_normalize_dataframe_and_schema(
+    dataframe: DataFrame,
+    schema: Schema,
+    schema_patch: Dict[SchemaFieldName, SchemaFieldValue],
+    column_name_to_include: Dict[str, bool],
+) -> Tuple[DataFrame, Schema]:
+    """
+    Creates new dataframe and schema objects to reflect excluded column names
+    and discovered features. This also normalizes dataframe columns to ensure a
+    standard set of columns (i.e. timestamp and prediction_id) and datatypes for
+    those columns.
+    """
+    included_column_names: List[str] = []
+    for column_name in dataframe.columns:
+        if column_name_to_include.get(str(column_name), False):
+            included_column_names.append(str(column_name))
+    parsed_dataframe = dataframe[included_column_names].copy()
+    parsed_schema = replace(schema, excluded_column_names=None, **schema_patch)  # type: ignore
+    pred_id_col_name = parsed_schema.prediction_id_column_name
+    if pred_id_col_name is None:
+        parsed_schema = replace(parsed_schema, prediction_id_column_name="prediction_id")
+        parsed_dataframe["prediction_id"] = _add_prediction_id(len(parsed_dataframe))
+    elif is_numeric_dtype(parsed_dataframe.dtypes[pred_id_col_name]):
+        parsed_dataframe[pred_id_col_name] = parsed_dataframe[pred_id_col_name].astype(str)
+    for embedding in (
+        [parsed_schema.prompt_column_names, parsed_schema.response_column_names]
+        + list(parsed_schema.embedding_feature_column_names.values())
+        if parsed_schema.embedding_feature_column_names is not None
+        else []
+    ):
+        if not isinstance(embedding, EmbeddingColumnNames):
+            continue
+        vector_column_name = embedding.vector_column_name
+        if vector_column_name not in parsed_dataframe.columns:
+            continue
+        parsed_dataframe.loc[:, vector_column_name] = _coerce_vectors_as_arrays_if_necessary(
+            parsed_dataframe.loc[:, vector_column_name],
+            vector_column_name,
+        )
+    return parsed_dataframe, parsed_schema
+
+
+def _coerce_vectors_as_arrays_if_necessary(
+    series: "pd.Series[Any]",
+    column_name: str,
+) -> "pd.Series[Any]":
+    not_na = ~series.isna()
+    if not_na.sum() == 0:
+        return series
+    if invalid_types := set(map(type, series.loc[not_na])) - {np.ndarray}:
+        logger.warning(
+            f"converting items in column `{column_name}` to numpy.ndarray, "
+            f"because they have the following "
+            f"type{'s' if len(invalid_types) > 1 else ''}: "
+            f"{', '.join(map(lambda t: t.__name__, invalid_types))}"
+        )
+        return series.mask(not_na, series.loc[not_na].apply(np.array))
+    return series
+
+
+def _sort_dataframe_rows_by_timestamp(dataframe: DataFrame, schema: Schema) -> DataFrame:
+    """
+    Sorts dataframe rows by timestamp.
+    """
+    timestamp_column_name = schema.timestamp_column_name
+    if timestamp_column_name is None:
+        raise ValueError("Schema must specify a timestamp column name.")
+    dataframe.set_index(timestamp_column_name, drop=False, inplace=True)
+    dataframe.sort_index(inplace=True)
+    return dataframe
+
+
+def _normalize_timestamps(
+    dataframe: DataFrame,
+    schema: Schema,
+    default_timestamp: Timestamp,
+) -> Tuple[DataFrame, Schema]:
+    """
+    Ensures that the dataframe has a timestamp column and the schema has a timestamp field. If the
+    input dataframe contains a Unix or datetime timestamp or ISO8601 timestamp strings column, it
+    is converted to UTC timezone-aware timestamp. If the input dataframe and schema do not contain
+    timestamps, the default timestamp is used. If a timestamp is timezone-naive, it is localized
+    as per local timezone and then converted to UTC
+    """
+    timestamp_column: Series[Timestamp]
+    if (timestamp_column_name := schema.timestamp_column_name) is None:
+        timestamp_column_name = "timestamp"
+        schema = replace(schema, timestamp_column_name=timestamp_column_name)
+        timestamp_column = (
+            Series([default_timestamp] * len(dataframe), index=dataframe.index)
+            if len(dataframe)
+            else Series([default_timestamp]).iloc[:0].set_axis(dataframe.index, axis=0)
+        )
+    else:
+        timestamp_column = normalize_timestamps(
+            dataframe[timestamp_column_name],
+        )
+    dataframe[timestamp_column_name] = timestamp_column
+    return dataframe, schema
+
+
+def _get_schema_from_unknown_schema_param(schemaLike: SchemaLike) -> Schema:
+    """
+    Compatibility function for converting from arize.utils.types.Schema to phoenix.inferences.Schema
+    """
+    try:
+        from arize.utils.types import (
+            EmbeddingColumnNames as ArizeEmbeddingColumnNames,  # fmt: off type: ignore
+        )
+        from arize.utils.types import Schema as ArizeSchema
+
+        if not isinstance(schemaLike, ArizeSchema):
+            raise ValueError("Unknown schema passed to Dataset. Please pass a phoenix Schema")
+
+        embedding_feature_column_names: Dict[str, EmbeddingColumnNames] = {}
+        if schemaLike.embedding_feature_column_names is not None:
+            for (
+                embedding_name,
+                arize_embedding_feature_column_names,
+            ) in schemaLike.embedding_feature_column_names.items():
+                if isinstance(arize_embedding_feature_column_names, ArizeEmbeddingColumnNames):
+                    embedding_feature_column_names[embedding_name] = EmbeddingColumnNames(
+                        vector_column_name=arize_embedding_feature_column_names.vector_column_name,
+                        link_to_data_column_name=arize_embedding_feature_column_names.link_to_data_column_name,
+                        raw_data_column_name=arize_embedding_feature_column_names.data_column_name,
+                    )
+        prompt_column_names: Optional[EmbeddingColumnNames] = None
+        if schemaLike.prompt_column_names is not None and isinstance(
+            schemaLike.prompt_column_names, ArizeEmbeddingColumnNames
+        ):
+            prompt_column_names = EmbeddingColumnNames(
+                vector_column_name=schemaLike.prompt_column_names.vector_column_name,
+                raw_data_column_name=schemaLike.prompt_column_names.data_column_name,
+                link_to_data_column_name=schemaLike.prompt_column_names.link_to_data_column_name,
+            )
+        response_column_names: Optional[EmbeddingColumnNames] = None
+        if schemaLike.response_column_names is not None and isinstance(
+            schemaLike.response_column_names, ArizeEmbeddingColumnNames
+        ):
+            response_column_names = EmbeddingColumnNames(
+                vector_column_name=schemaLike.response_column_names.vector_column_name,
+                raw_data_column_name=schemaLike.response_column_names.data_column_name,
+                link_to_data_column_name=schemaLike.response_column_names.link_to_data_column_name,
+            )
+        return Schema(
+            feature_column_names=schemaLike.feature_column_names,
+            tag_column_names=schemaLike.tag_column_names,
+            prediction_label_column_name=schemaLike.prediction_label_column_name,
+            actual_label_column_name=schemaLike.actual_label_column_name,
+            prediction_id_column_name=schemaLike.prediction_id_column_name,
+            timestamp_column_name=schemaLike.timestamp_column_name,
+            embedding_feature_column_names=embedding_feature_column_names,
+            prompt_column_names=prompt_column_names,
+            response_column_names=response_column_names,
+        )
+    except Exception:
+        raise ValueError(
+            """Unsupported Arize Schema. Please pass a phoenix Schema or update
+            to the latest version of Arize."""
+        )
+
+
+def _add_prediction_id(num_rows: int) -> List[str]:
+    return [str(uuid.uuid4()) for _ in range(num_rows)]
+
+
+# A dataset with no data. Useful for stubs
+EMPTY_INFERENCES = Inferences(pd.DataFrame(), schema=Schema())