arize-phoenix 4.4.3__py3-none-any.whl → 4.4.4rc0__py3-none-any.whl

phoenix/inferences/fixtures.py

@@ -9,7 +9,7 @@ from urllib.parse import quote, urljoin
 
 from pandas import read_parquet
 
-from phoenix.config import DATASET_DIR
+from phoenix.config import INFERENCES_DIR
 from phoenix.inferences.inferences import Inferences
 from phoenix.inferences.schema import (
     EmbeddingColumnNames,
@@ -20,7 +20,7 @@ from phoenix.inferences.schema import (
 logger = logging.getLogger(__name__)
 
 
-class DatasetRole(Enum):
+class InferencesRole(Enum):
     PRIMARY = auto()
     REFERENCE = auto()
     CORPUS = auto()
@@ -39,11 +39,11 @@ class Fixture:
     corpus_file_name: Optional[str] = None
     corpus_schema: Optional[Schema] = None
 
-    def paths(self) -> Iterator[Tuple[DatasetRole, Path]]:
+    def paths(self) -> Iterator[Tuple[InferencesRole, Path]]:
         return (
             (role, Path(self.prefix) / name)
             for role, name in zip(
-                DatasetRole,
+                InferencesRole,
                 (
                     self.primary_file_name,
                     self.reference_file_name,
@@ -413,41 +413,41 @@ FIXTURES: Tuple[Fixture, ...] = (
 NAME_TO_FIXTURE = {fixture.name: fixture for fixture in FIXTURES}
 
 
-def get_datasets(
+def get_inferences(
     fixture_name: str,
     no_internet: bool = False,
 ) -> Tuple[Inferences, Optional[Inferences], Optional[Inferences]]:
     """
-    Downloads primary and reference datasets for a fixture if they are not found
+    Downloads primary and reference inferences for a fixture if they are not found
     locally.
     """
     fixture = _get_fixture_by_name(fixture_name=fixture_name)
     if no_internet:
-        paths = {role: DATASET_DIR / path for role, path in fixture.paths()}
+        paths = {role: INFERENCES_DIR / path for role, path in fixture.paths()}
     else:
-        paths = dict(_download(fixture, DATASET_DIR))
-    primary_dataset = Inferences(
-        read_parquet(paths[DatasetRole.PRIMARY]),
+        paths = dict(_download(fixture, INFERENCES_DIR))
+    primary_inferences = Inferences(
+        read_parquet(paths[InferencesRole.PRIMARY]),
         fixture.primary_schema,
         "production",
     )
-    reference_dataset = None
+    reference_inferences = None
     if fixture.reference_file_name is not None:
-        reference_dataset = Inferences(
-            read_parquet(paths[DatasetRole.REFERENCE]),
+        reference_inferences = Inferences(
+            read_parquet(paths[InferencesRole.REFERENCE]),
             fixture.reference_schema
             if fixture.reference_schema is not None
             else fixture.primary_schema,
             "training",
         )
-    corpus_dataset = None
+    corpus_inferences = None
     if fixture.corpus_file_name is not None:
-        corpus_dataset = Inferences(
-            read_parquet(paths[DatasetRole.CORPUS]),
+        corpus_inferences = Inferences(
+            read_parquet(paths[InferencesRole.CORPUS]),
             fixture.corpus_schema,
             "knowledge_base",
         )
-    return primary_dataset, reference_dataset, corpus_dataset
+    return primary_inferences, reference_inferences, corpus_inferences
 
 
 def _get_fixture_by_name(fixture_name: str) -> Fixture:
@@ -496,14 +496,14 @@ def load_example(use_case: str) -> ExampleInferences:
 
     """
     fixture = _get_fixture_by_name(use_case)
-    primary_dataset, reference_dataset, corpus_dataset = get_datasets(use_case)
+    primary_inferences, reference_inferences, corpus_inferences = get_inferences(use_case)
     print(f"📥 Loaded {use_case} example datasets.")
     print("ℹ️ About this use-case:")
     print(fixture.description)
     return ExampleInferences(
-        primary=primary_dataset,
-        reference=reference_dataset,
-        corpus=corpus_dataset,
+        primary=primary_inferences,
+        reference=reference_inferences,
+        corpus=corpus_inferences,
     )
 
 
@@ -544,7 +544,7 @@ class GCSAssets(NamedTuple):
         )
 
 
-def _download(fixture: Fixture, location: Path) -> Iterator[Tuple[DatasetRole, Path]]:
+def _download(fixture: Fixture, location: Path) -> Iterator[Tuple[InferencesRole, Path]]:
     for role, path in fixture.paths():
         yield role, GCSAssets().metadata(path).save_artifact(location)
 
@@ -556,5 +556,5 @@ if __name__ == "__main__":
     for fixture in FIXTURES:
         start_time = time.time()
         print(f"getting {fixture.name}", end="...")
-        dict(_download(fixture, DATASET_DIR))
+        dict(_download(fixture, INFERENCES_DIR))
         print(f"done ({time.time() - start_time:.2f}s)")

phoenix/inferences/inferences.py

@@ -15,7 +15,7 @@ from pandas.api.types import (
 )
 from typing_extensions import TypeAlias
 
-from phoenix.config import DATASET_DIR, GENERATED_DATASET_NAME_PREFIX
+from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX, INFERENCES_DIR
 from phoenix.datetime_utils import normalize_timestamps
 from phoenix.utilities.deprecation import deprecated
 
@@ -31,7 +31,7 @@ from .schema import (
     SchemaFieldName,
     SchemaFieldValue,
 )
-from .validation import validate_dataset_inputs
+from .validation import validate_inferences_inputs
 
 logger = logging.getLogger(__name__)
 
@@ -62,7 +62,7 @@ class Inferences:
 
     Examples
     --------
-    >>> primary_dataset = px.Inferences(
+    >>> primary_inferences = px.Inferences(
     >>>    dataframe=production_dataframe, schema=schema, name="primary"
     >>> )
     """
@@ -81,7 +81,7 @@ class Inferences:
         # allow for schema like objects
         if not isinstance(schema, Schema):
             schema = _get_schema_from_unknown_schema_param(schema)
-        errors = validate_dataset_inputs(
+        errors = validate_inferences_inputs(
             dataframe=dataframe,
             schema=schema,
         )
@@ -95,7 +95,7 @@ class Inferences:
         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())}"
+            name if name is not None else f"{GENERATED_INFERENCES_NAME_PREFIX}{str(uuid.uuid4())}"
         )
         self._is_empty = self.dataframe.empty
         logger.info(f"""Dataset: {self.__name} initialized""")
@@ -118,7 +118,7 @@ class Inferences:
     @classmethod
     def from_name(cls, name: str) -> "Inferences":
         """Retrieves a dataset by name from the file system"""
-        directory = DATASET_DIR / name
+        directory = INFERENCES_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()
@@ -127,7 +127,7 @@ class Inferences:
 
     def to_disc(self) -> None:
         """writes the data and schema to disc"""
-        directory = DATASET_DIR / self.name
+        directory = INFERENCES_DIR / self.name
         directory.mkdir(parents=True, exist_ok=True)
         self.dataframe.to_parquet(
             directory / self._data_file_name,

phoenix/inferences/validation.py

@@ -34,7 +34,7 @@ def _check_valid_schema(schema: Schema) -> List[err.ValidationError]:
     return []
 
 
-def validate_dataset_inputs(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
+def validate_inferences_inputs(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
     errors = _check_missing_columns(dataframe, schema)
     if errors:
         return errors

phoenix/server/api/context.py

@@ -12,33 +12,49 @@ from typing_extensions import TypeAlias
 from phoenix.core.model_schema import Model
 from phoenix.server.api.dataloaders import (
     CacheForDataLoaders,
+    DatasetExampleRevisionsDataLoader,
+    DatasetExampleSpansDataLoader,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
     EvaluationSummaryDataLoader,
+    ExperimentAnnotationSummaryDataLoader,
+    ExperimentErrorRatesDataLoader,
+    ExperimentSequenceNumberDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
+    ProjectByNameDataLoader,
     RecordCountDataLoader,
     SpanDescendantsDataLoader,
     SpanEvaluationsDataLoader,
+    SpanProjectsDataLoader,
     TokenCountDataLoader,
     TraceEvaluationsDataLoader,
+    TraceRowIdsDataLoader,
 )
 
 
 @dataclass
 class DataLoaders:
+    dataset_example_revisions: DatasetExampleRevisionsDataLoader
+    dataset_example_spans: DatasetExampleSpansDataLoader
     document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
     document_evaluations: DocumentEvaluationsDataLoader
     document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
     evaluation_summaries: EvaluationSummaryDataLoader
+    experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
+    experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_sequence_number: ExperimentSequenceNumberDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
     record_counts: RecordCountDataLoader
     span_descendants: SpanDescendantsDataLoader
     span_evaluations: SpanEvaluationsDataLoader
+    span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
     trace_evaluations: TraceEvaluationsDataLoader
+    trace_row_ids: TraceRowIdsDataLoader
+    project_by_name: ProjectByNameDataLoader
 
 
 ProjectRowId: TypeAlias = int

phoenix/server/api/dataloaders/__init__.py

@@ -8,6 +8,8 @@ from phoenix.db.insertion.evaluation import (
 )
 from phoenix.db.insertion.span import ClearProjectSpansEvent, SpanInsertionEvent
 
+from .dataset_example_revisions import DatasetExampleRevisionsDataLoader
+from .dataset_example_spans import DatasetExampleSpansDataLoader
 from .document_evaluation_summaries import (
     DocumentEvaluationSummaryCache,
     DocumentEvaluationSummaryDataLoader,
@@ -15,27 +17,41 @@ from .document_evaluation_summaries import (
 from .document_evaluations import DocumentEvaluationsDataLoader
 from .document_retrieval_metrics import DocumentRetrievalMetricsDataLoader
 from .evaluation_summaries import EvaluationSummaryCache, EvaluationSummaryDataLoader
+from .experiment_annotation_summaries import ExperimentAnnotationSummaryDataLoader
+from .experiment_error_rates import ExperimentErrorRatesDataLoader
+from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
+from .project_by_name import ProjectByNameDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
 from .span_descendants import SpanDescendantsDataLoader
 from .span_evaluations import SpanEvaluationsDataLoader
+from .span_projects import SpanProjectsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
 from .trace_evaluations import TraceEvaluationsDataLoader
+from .trace_row_ids import TraceRowIdsDataLoader
 
 __all__ = [
     "CacheForDataLoaders",
+    "DatasetExampleRevisionsDataLoader",
+    "DatasetExampleSpansDataLoader",
     "DocumentEvaluationSummaryDataLoader",
     "DocumentEvaluationsDataLoader",
     "DocumentRetrievalMetricsDataLoader",
     "EvaluationSummaryDataLoader",
+    "ExperimentAnnotationSummaryDataLoader",
+    "ExperimentErrorRatesDataLoader",
+    "ExperimentSequenceNumberDataLoader",
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
     "RecordCountDataLoader",
     "SpanDescendantsDataLoader",
     "SpanEvaluationsDataLoader",
+    "SpanProjectsDataLoader",
     "TokenCountDataLoader",
     "TraceEvaluationsDataLoader",
+    "TraceRowIdsDataLoader",
+    "ProjectByNameDataLoader",
 ]
 
 

phoenix/server/api/dataloaders/dataset_example_revisions.py

@@ -0,0 +1,100 @@
+from typing import (
+    AsyncContextManager,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+from sqlalchemy import Integer, case, func, literal, or_, select, union
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
+
+ExampleID: TypeAlias = int
+VersionID: TypeAlias = Optional[int]
+Key: TypeAlias = Tuple[ExampleID, Optional[VersionID]]
+Result: TypeAlias = DatasetExampleRevision
+
+
+class DatasetExampleRevisionsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Union[Result, ValueError]]:
+        # sqlalchemy has limited SQLite support for VALUES, so use UNION ALL instead.
+        # For details, see https://github.com/sqlalchemy/sqlalchemy/issues/7228
+        keys_subquery = union(
+            *(
+                select(
+                    literal(example_id, Integer).label("example_id"),
+                    literal(version_id, Integer).label("version_id"),
+                )
+                for example_id, version_id in keys
+            )
+        ).subquery()
+        revision_ids = (
+            select(
+                keys_subquery.c.example_id,
+                keys_subquery.c.version_id,
+                func.max(models.DatasetExampleRevision.id).label("revision_id"),
+            )
+            .select_from(keys_subquery)
+            .join(
+                models.DatasetExampleRevision,
+                onclause=keys_subquery.c.example_id
+                == models.DatasetExampleRevision.dataset_example_id,
+            )
+            .where(
+                or_(
+                    keys_subquery.c.version_id.is_(None),
+                    models.DatasetExampleRevision.dataset_version_id <= keys_subquery.c.version_id,
+                )
+            )
+            .group_by(keys_subquery.c.example_id, keys_subquery.c.version_id)
+        ).subquery()
+        query = (
+            select(
+                revision_ids.c.example_id,
+                revision_ids.c.version_id,
+                case(
+                    (
+                        or_(
+                            revision_ids.c.version_id.is_(None),
+                            models.DatasetVersion.id.is_not(None),
+                        ),
+                        True,
+                    ),
+                    else_=False,
+                ).label("is_valid_version"),  # check that non-null versions exist
+                models.DatasetExampleRevision,
+            )
+            .select_from(revision_ids)
+            .join(
+                models.DatasetExampleRevision,
+                onclause=revision_ids.c.revision_id == models.DatasetExampleRevision.id,
+            )
+            .join(
+                models.DatasetVersion,
+                onclause=revision_ids.c.version_id == models.DatasetVersion.id,
+                isouter=True,  # keep rows where the version id is null
+            )
+            .where(models.DatasetExampleRevision.revision_kind != "DELETE")
+        )
+        async with self._db() as session:
+            results = {
+                (example_id, version_id): DatasetExampleRevision.from_orm_revision(revision)
+                async for (
+                    example_id,
+                    version_id,
+                    is_valid_version,
+                    revision,
+                ) in await session.stream(query)
+                if is_valid_version
+            }
+        return [results.get(key, ValueError("Could not find revision.")) for key in keys]

phoenix/server/api/dataloaders/dataset_example_spans.py

@@ -0,0 +1,43 @@
+from typing import (
+    AsyncContextManager,
+    Callable,
+    List,
+    Optional,
+)
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+ExampleID: TypeAlias = int
+Key: TypeAlias = ExampleID
+Result: TypeAlias = Optional[models.Span]
+
+
+class DatasetExampleSpansDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        example_ids = keys
+        async with self._db() as session:
+            spans = {
+                example_id: span
+                async for example_id, span in await session.stream(
+                    select(models.DatasetExample.id, models.Span)
+                    .select_from(models.DatasetExample)
+                    .join(models.Span, models.DatasetExample.span_rowid == models.Span.id)
+                    .where(models.DatasetExample.id.in_(example_ids))
+                    .options(
+                        joinedload(models.Span.trace, innerjoin=True).load_only(
+                            models.Trace.trace_id
+                        )
+                    )
+                )
+            }
+        return [spans.get(example_id) for example_id in example_ids]

phoenix/server/api/dataloaders/experiment_annotation_summaries.py

@@ -0,0 +1,85 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import (
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+    Optional,
+)
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+
+@dataclass
+class ExperimentAnnotationSummary:
+    annotation_name: str
+    min_score: float
+    max_score: float
+    mean_score: float
+    count: int
+    error_count: int
+
+
+ExperimentID: TypeAlias = int
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = List[ExperimentAnnotationSummary]
+
+
+class ExperimentAnnotationSummaryDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+        cache_map: Optional[AbstractCache[Key, Result]] = None,
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        experiment_ids = keys
+        summaries: DefaultDict[ExperimentID, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for (
+                experiment_id,
+                annotation_name,
+                min_score,
+                max_score,
+                mean_score,
+                count,
+                error_count,
+            ) in await session.stream(
+                select(
+                    models.ExperimentRun.experiment_id,
+                    models.ExperimentRunAnnotation.name,
+                    func.min(models.ExperimentRunAnnotation.score),
+                    func.max(models.ExperimentRunAnnotation.score),
+                    func.avg(models.ExperimentRunAnnotation.score),
+                    func.count(),
+                    func.count(models.ExperimentRunAnnotation.error),
+                )
+                .join(
+                    models.ExperimentRun,
+                    models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
+                )
+                .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+                .group_by(models.ExperimentRun.experiment_id, models.ExperimentRunAnnotation.name)
+            ):
+                summaries[experiment_id].append(
+                    ExperimentAnnotationSummary(
+                        annotation_name=annotation_name,
+                        min_score=min_score,
+                        max_score=max_score,
+                        mean_score=mean_score,
+                        count=count,
+                        error_count=error_count,
+                    )
+                )
+        return [
+            sorted(summaries[experiment_id], key=lambda summary: summary.annotation_name)
+            for experiment_id in experiment_ids
+        ]

phoenix/server/api/dataloaders/experiment_error_rates.py

@@ -0,0 +1,43 @@
+from typing import (
+    AsyncContextManager,
+    Callable,
+    List,
+    Optional,
+)
+
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+ExperimentID: TypeAlias = int
+ErrorRate: TypeAlias = float
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = Optional[ErrorRate]
+
+
+class ExperimentErrorRatesDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        experiment_ids = keys
+        async with self._db() as session:
+            error_rates = {
+                experiment_id: error_rate
+                async for experiment_id, error_rate in await session.stream(
+                    select(
+                        models.ExperimentRun.experiment_id,
+                        func.count(models.ExperimentRun.error) / func.count(),
+                    )
+                    .group_by(models.ExperimentRun.experiment_id)
+                    .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+                )
+            }
+        return [error_rates.get(experiment_id) for experiment_id in experiment_ids]

phoenix/server/api/dataloaders/experiment_sequence_number.py

@@ -0,0 +1,49 @@
+from typing import (
+    AsyncContextManager,
+    Callable,
+    List,
+    Optional,
+)
+
+from sqlalchemy import distinct, func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+ExperimentId: TypeAlias = int
+Key: TypeAlias = ExperimentId
+Result: TypeAlias = Optional[int]
+
+
+class ExperimentSequenceNumberDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        experiment_ids = keys
+        dataset_ids = (
+            select(distinct(models.Experiment.dataset_id))
+            .where(models.Experiment.id.in_(experiment_ids))
+            .scalar_subquery()
+        )
+        row_number = (
+            func.row_number().over(
+                partition_by=models.Experiment.dataset_id,
+                order_by=models.Experiment.id,
+            )
+        ).label("row_number")
+        subq = (
+            select(models.Experiment.id, row_number)
+            .where(models.Experiment.dataset_id.in_(dataset_ids))
+            .subquery()
+        )
+        stmt = select(subq).where(subq.c.id.in_(experiment_ids))
+        async with self._db() as session:
+            result = {
+                experiment_id: sequence_number
+                async for experiment_id, sequence_number in await session.stream(stmt)
+            }
+        return [result.get(experiment_id) for experiment_id in experiment_ids]

phoenix/server/api/dataloaders/project_by_name.py

@@ -0,0 +1,31 @@
+from collections import defaultdict
+from typing import AsyncContextManager, Callable, DefaultDict, List, Optional
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+ProjectName: TypeAlias = str
+Key: TypeAlias = ProjectName
+Result: TypeAlias = Optional[models.Project]
+
+
+class ProjectByNameDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        project_names = list(set(keys))
+        projects_by_name: DefaultDict[Key, Result] = defaultdict(None)
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(models.Project).where(models.Project.name.in_(project_names))
+            )
+            async for project in data:
+                projects_by_name[project.name] = project
+
+        return [projects_by_name[project_name] for project_name in project_names]

phoenix/server/api/dataloaders/span_descendants.py

@@ -9,7 +9,7 @@ from typing import (
 from aioitertools.itertools import groupby
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import contains_eager
+from sqlalchemy.orm import joinedload
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -52,8 +52,7 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
         stmt = (
             select(descendant_ids.c[root_id_label], models.Span)
             .join(descendant_ids, models.Span.id == descendant_ids.c.id)
-            .join(models.Trace)
-            .options(contains_eager(models.Span.trace))
+            .options(joinedload(models.Span.trace, innerjoin=True).load_only(models.Trace.trace_id))
             .order_by(descendant_ids.c[root_id_label])
         )
         results: Dict[SpanId, Result] = {key: [] for key in keys}