Skip to content

Commit

Permalink
fix(sessions): exclude null input/output values in dataloader query f…
Browse files Browse the repository at this point in the history
…or sessions table (#5781)
  • Loading branch information
RogerHYang authored Dec 18, 2024
1 parent 26e83b1 commit fb1330f
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions src/phoenix/server/api/dataloaders/session_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ def _subq(self) -> Select[tuple[Optional[int], str, str, int]]:
)
if self._kind == "first_input":
stmt = stmt.add_columns(
models.Span.attributes[INPUT_VALUE].label("value"),
models.Span.attributes[INPUT_MIME_TYPE].label("mime_type"),
models.Span.attributes[INPUT_VALUE].as_string().label("value"),
models.Span.attributes[INPUT_MIME_TYPE].as_string().label("mime_type"),
func.row_number()
.over(
partition_by=models.Trace.project_session_rowid,
Expand All @@ -42,8 +42,8 @@ def _subq(self) -> Select[tuple[Optional[int], str, str, int]]:
)
elif self._kind == "last_output":
stmt = stmt.add_columns(
models.Span.attributes[OUTPUT_VALUE].label("value"),
models.Span.attributes[OUTPUT_MIME_TYPE].label("mime_type"),
models.Span.attributes[OUTPUT_VALUE].as_string().label("value"),
models.Span.attributes[OUTPUT_MIME_TYPE].as_string().label("mime_type"),
func.row_number()
.over(
partition_by=models.Trace.project_session_rowid,
Expand All @@ -57,7 +57,11 @@ def _subq(self) -> Select[tuple[Optional[int], str, str, int]]:

def _stmt(self, *keys: Key) -> Select[tuple[int, str, str]]:
subq = self._subq.where(models.Trace.project_session_rowid.in_(keys)).subquery()
return select(subq.c.id_, subq.c.value, subq.c.mime_type).filter_by(rank=1)
return (
select(subq.c.id_, subq.c.value, subq.c.mime_type)
.filter_by(rank=1)
.where(subq.c.value.isnot(None))
)

async def _load_fn(self, keys: list[Key]) -> list[Result]:
async with self._db() as session:
Expand Down

0 comments on commit fb1330f

Please sign in to comment.