add min max count (#327)

Arize-ai · Mar 3, 2023 · b01b182 · b01b182
1 parent a802866
commit b01b182
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 6 deletions.
diff --git a/app/schema.graphql b/app/schema.graphql
@@ -8,6 +8,9 @@ enum DataQualityMetric {
   percentEmpty
   mean
   sum
+  min
+  max
+  count
 }
 
 type DataQualityTimeSeries implements TimeSeries {

diff --git a/src/phoenix/metrics/metrics.py b/src/phoenix/metrics/metrics.py
@@ -51,6 +51,16 @@ def calc(self, df: pd.DataFrame) -> Union[float, npt.NDArray[np.float64]]:
             )
 
 
+class Min(UnaryOperator, BaseMetric):
+    def calc(self, df: pd.DataFrame) -> float:
+        return cast(float, df.loc[:, self.operand].min())
+
+
+class Max(UnaryOperator, BaseMetric):
+    def calc(self, df: pd.DataFrame) -> float:
+        return cast(float, df.loc[:, self.operand].max())
+
+
 class Cardinality(UnaryOperator, BaseMetric):
     def calc(self, df: pd.DataFrame) -> int:
         return df.loc[:, self.operand].nunique()

diff --git a/src/phoenix/server/api/types/DataQualityMetric.py b/src/phoenix/server/api/types/DataQualityMetric.py
@@ -9,3 +9,6 @@ class DataQualityMetric(Enum):
     percentEmpty = "PercentEmpty"
     mean = "Mean"
     sum = "Sum"
+    min = "Min"
+    max = "Max"
+    count = "Count"
diff --git a/src/phoenix/server/api/types/Dimension.py b/src/phoenix/server/api/types/Dimension.py
@@ -1,3 +1,4 @@
+import math
 from datetime import timedelta
 from typing import List, Optional
 
@@ -34,12 +35,12 @@ class Dimension(Node):
     async def dataQualityMetric(
         self, metric: DataQualityMetric, info: Info[Context, None]
     ) -> Optional[float]:
-        dimension_name = self.name
-        if metric is DataQualityMetric.cardinality:
-            return await info.context.loaders.cardinality.load(dimension_name)
-        elif metric is DataQualityMetric.percentEmpty:
-            return await info.context.loaders.percent_empty.load(dimension_name)
-        raise NotImplementedError(f"Metric {metric} is not implemented.")
+        metric_cls = METRICS.get(metric.value, None)
+        if not metric_cls or not issubclass(metric_cls, UnaryOperator):
+            raise NotImplementedError(f"Metric {metric} is not implemented.")
+        df = info.context.model.primary_dataset.dataframe
+        _, ans = metric_cls(self.name)(df)
+        return None if math.isnan(ans) else ans
 
     @strawberry.field(
         description=(