feat: time-windowed euclidean distance (#247)

* feat: gql interface for timeseries

* format fix

* correct docstrings

* tmp

* tmp

* tmp

* added tests and got them passing

* update graphql schema

* hard-code time range in driftMetric query

* run prettier

* readability improvements

* readability improvements

* update graphql schema

* simplify code to not round to nearest hour

---------

Co-authored-by: Mikyo King <mikyo@arize.com>

app/schema.graphql

@@ -58,10 +58,13 @@ type DriftTimeSeries implements TimeSeries {
 type EmbeddingDimension implements Node {
   id: GlobalID!
   name: String!
+  driftMetric(metric: DriftMetric!, timeRange: TimeRange!): Float
   driftTimeSeries(
+    metric: DriftMetric!
+
     """The time range of the primary dataset"""
     timeRange: TimeRange!
-  ): DriftTimeSeries!
+  ): DriftTimeSeries
   UMAPPoints(
     """The time range of the primary dataset to generate the UMAP points for"""
     timeRange: TimeRange!
@@ -78,7 +81,6 @@ type EmbeddingDimension implements Node {
     """UMAP N samples"""
     nSamples: Int = 500
   ): UMAPPoints!
-  driftMetric(metric: DriftMetric!): Float
 }
 
 type EmbeddingDimensionConnection {

app/src/components/model/ModelEmbeddingsTable.tsx

@@ -25,7 +25,13 @@ export function ModelEmbeddingsTable(props: ModelEmbeddingsTable) {
               embedding: node {
                 id
                 name
-                euclideanDistance: driftMetric(metric: euclideanDistance)
+                euclideanDistance: driftMetric(
+                  metric: euclideanDistance
+                  timeRange: {
+                    start: "1970-01-20 02:00:00"
+                    end: "1970-01-20 04:00:00"
+                  }
+                )
               }
             }
           }

src/phoenix/datasets/dataset.py

@@ -160,6 +160,12 @@ def _get_embedding_feature_column_names(
             raise err.SchemaError(err.MissingEmbeddingFeatureColumnNames(embedding_feature_name))
         return embedding_feature_column_names[embedding_feature_name]
 
+    def get_timestamp_column(self) -> "Series[Any]":
+        timestamp_column_name = self.schema.timestamp_column_name
+        if timestamp_column_name is None:
+            raise err.SchemaError(err.MissingTimestampColumnName())
+        return self.dataframe[timestamp_column_name]
+
     # TODO(mikeldking): add strong vector type
     def get_embedding_vector_column(self, embedding_feature_name: str) -> "Series[Any]":
         column_names = self._get_embedding_feature_column_names(embedding_feature_name)

src/phoenix/datasets/errors.py

@@ -134,6 +134,16 @@ def error_message(self) -> str:
         )
 
 
+class MissingTimestampColumnName(ValidationError):
+    """
+    An error raised when trying to access a timestamp column that is absent from
+    the schema.
+    """
+
+    def error_message(self) -> str:
+        return "Schema is missing timestamp_column_name."
+
+
 class SchemaError(Exception):
     """An error raised when the Schema is invalid or incomplete"""
 

src/phoenix/metrics/embeddings.py

@@ -6,10 +6,10 @@
 
 
 def euclidean_distance(
-    array0: npt.NDArray[np.float64],
-    array1: npt.NDArray[np.float64],
+    pt0: npt.NDArray[np.float64],
+    pt1: npt.NDArray[np.float64],
 ) -> Optional[float]:
     """
-    Computes Euclidean distance between the centroids of two arrays.
+    Computes Euclidean distance between two points.
     """
-    return cast(float, euclidean(np.mean(array0, axis=0), np.mean(array1, axis=0)))
+    return cast(float, euclidean(pt0, pt1))

src/phoenix/server/api/input_types/TimeRange.py

@@ -7,3 +7,6 @@
 class TimeRange:
     start: datetime.datetime
     end: datetime.datetime
+
+    def is_valid(self) -> bool:
+        return self.start < self.end

src/phoenix/server/api/types/DriftTimeSeries.py

@@ -7,4 +7,4 @@
 class DriftTimeSeries(TimeSeries):
     """A time series of drift metrics"""
 
-    pass
+    ...