wrapped DataType

Author: czaloom

src/valor_lite/classification/evaluator.py

@@ -50,6 +50,7 @@
     unpack_rocauc,
 )
 from valor_lite.exceptions import EmptyCacheError
+from valor_lite.filtering import DataType, Expression
 
 
 class Builder:
@@ -58,7 +59,9 @@ def __init__(
         writer: MemoryCacheWriter | FileCacheWriter,
         roc_curve_writer: MemoryCacheWriter | FileCacheWriter,
         intermediate_writer: MemoryCacheWriter | FileCacheWriter,
-        metadata_fields: list[tuple[str, str | pa.DataType]] | None = None,
+        metadata_fields: list[tuple[str, DataType]]
+        | list[tuple[str, str]]
+        | None = None,
     ):
         self._writer = writer
         self._roc_curve_writer = roc_curve_writer
@@ -69,7 +72,9 @@ def __init__(
     def in_memory(
         cls,
         batch_size: int = 10_000,
-        metadata_fields: list[tuple[str, str | pa.DataType]] | None = None,
+        metadata_fields: list[tuple[str, DataType]]
+        | list[tuple[str, str]]
+        | None = None,
     ):
         """
         Create an in-memory evaluator cache.
@@ -78,7 +83,7 @@ def in_memory(
         ----------
         batch_size : int, default=10_000
             The target number of rows to buffer before writing to the cache. Defaults to 10_000.
-        metadata_fields : list[tuple[str, str | pa.DataType]], optional
+        metadata_fields : list[tuple[str, str | DataType]], optional
             Optional metadata field definitions.
         """
         writer = MemoryCacheWriter.create(
@@ -107,7 +112,9 @@ def persistent(
         batch_size: int = 10_000,
         rows_per_file: int = 100_000,
         compression: str = "snappy",
-        metadata_fields: list[tuple[str, str | pa.DataType]] | None = None,
+        metadata_fields: list[tuple[str, DataType]]
+        | list[tuple[str, str]]
+        | None = None,
     ):
         """
         Create a persistent file-based evaluator cache.
@@ -122,7 +129,7 @@ def persistent(
             Sets the maximum number of rows per file. This may be exceeded as files are datum aligned.
         compression : str, default="snappy"
             Sets the pyarrow compression method.
-        metadata_fields : list[tuple[str, str | pa.DataType]], optional
+        metadata_fields : list[tuple[str, str | DataType]], optional
             Optionally sets metadata description for use in filtering.
         """
         path = Path(path)
@@ -324,20 +331,26 @@ def __init__(
         reader: MemoryCacheReader | FileCacheReader,
         roc_curve_reader: MemoryCacheReader | FileCacheReader,
         index_to_label: dict[int, str],
-        metadata_fields: list[tuple[str, str | pa.DataType]] | None = None,
+        metadata_fields: list[tuple[str, str]]
+        | list[tuple[str, DataType]]
+        | None = None,
     ):
         self._reader = reader
         self._roc_curve_reader = roc_curve_reader
         self._index_to_label = index_to_label
-        self._metadata_fields = metadata_fields
+        self._metadata_fields = (
+            [(name, str(dtype)) for name, dtype in metadata_fields]
+            if metadata_fields
+            else None
+        )
 
     @property
     def info(self) -> EvaluatorInfo:
         return self.get_info()
 
     def get_info(
         self,
-        datums: pc.Expression | None = None,
+        datums: Expression | None = None,
     ) -> EvaluatorInfo:
         info = EvaluatorInfo()
         info.metadata_fields = self._metadata_fields
@@ -403,21 +416,21 @@ def load(
 
     def filter(
         self,
-        datums: pc.Expression | None = None,
-        groundtruths: pc.Expression | None = None,
-        predictions: pc.Expression | None = None,
+        datums: Expression | None = None,
+        groundtruths: Expression | None = None,
+        predictions: Expression | None = None,
         path: str | Path | None = None,
     ) -> Evaluator:
         """
         Filter evaluator cache.
 
         Parameters
         ----------
-        datums : pc.Expression | None = None
+        datums : Expression | None = None
             A filter expression used to filter datums.
-        groundtruths : pc.Expression | None = None
+        groundtruths : Expression | None = None
             A filter expression used to filter ground truth annotations.
-        predictions : pc.Expression | None = None
+        predictions : Expression | None = None
             A filter expression used to filter predictions.
         path : str | Path, optional
             Where to store the filtered cache if storing on disk.
@@ -447,7 +460,8 @@ def filter(
                 metadata_fields=self.info.metadata_fields,
             )
 
-        for tbl in self._reader.iterate_tables(filter=datums):
+        datum_filter = datums.to_arrow() if datums is not None else None
+        for tbl in self._reader.iterate_tables(filter=datum_filter):
             columns = (
                 "datum_id",
                 "gt_label_id",
@@ -461,7 +475,7 @@ def filter(
 
             if groundtruths is not None:
                 mask_valid_gt = np.zeros(n_pairs, dtype=np.bool_)
-                gt_tbl = tbl.filter(groundtruths)
+                gt_tbl = tbl.filter(groundtruths.to_arrow())
                 gt_pairs = np.column_stack(
                     [
                         gt_tbl[col].to_numpy()
@@ -475,7 +489,7 @@ def filter(
 
             if predictions is not None:
                 mask_valid_pd = np.zeros(n_pairs, dtype=np.bool_)
-                pd_tbl = tbl.filter(predictions)
+                pd_tbl = tbl.filter(predictions.to_arrow())
                 pd_pairs = np.column_stack(
                     [
                         pd_tbl[col].to_numpy()
@@ -503,7 +517,7 @@ def filter(
 
         return loader.finalize(index_to_label_override=self._index_to_label)
 
-    def iterate_values(self, datums: pc.Expression | None = None):
+    def _iterate_values(self, datum_filter: pc.Expression | None = None):
         columns = [
             "datum_id",
             "gt_label_id",
@@ -512,7 +526,9 @@ def iterate_values(self, datums: pc.Expression | None = None):
             "pd_winner",
             "match",
         ]
-        for tbl in self._reader.iterate_tables(columns=columns, filter=datums):
+        for tbl in self._reader.iterate_tables(
+            columns=columns, filter=datum_filter
+        ):
             ids = np.column_stack(
                 [
                     tbl[col].to_numpy()
@@ -528,8 +544,10 @@ def iterate_values(self, datums: pc.Expression | None = None):
             matches = tbl["match"].to_numpy()
             yield ids, scores, winners, matches
 
-    def iterate_values_with_tables(self, datums: pc.Expression | None = None):
-        for tbl in self._reader.iterate_tables(filter=datums):
+    def _iterate_values_with_tables(
+        self, datum_filter: pc.Expression | None = None
+    ):
+        for tbl in self._reader.iterate_tables(filter=datum_filter):
             ids = np.column_stack(
                 [
                     tbl[col].to_numpy()
@@ -594,7 +612,7 @@ def compute_precision_recall(
         self,
         score_thresholds: list[float] = [0.0],
         hardmax: bool = True,
-        datums: pc.Expression | None = None,
+        datums: Expression | None = None,
     ) -> dict[MetricType, list]:
         """
         Performs an evaluation and returns metrics.
@@ -613,6 +631,7 @@ def compute_precision_recall(
         dict[MetricType, list]
             A dictionary mapping MetricType enumerations to lists of computed metrics.
         """
+        datum_filter = datums.to_arrow() if datums is not None else None
         if not score_thresholds:
             raise ValueError("At least one score threshold must be passed.")
 
@@ -623,7 +642,9 @@ def compute_precision_recall(
         # intermediates
         counts = np.zeros((n_scores, n_labels, 4), dtype=np.uint64)
 
-        for ids, scores, winners, _ in self.iterate_values(datums=datums):
+        for ids, scores, winners, _ in self._iterate_values(
+            datum_filter=datum_filter
+        ):
             batch_counts = compute_counts(
                 ids=ids,
                 scores=scores,
@@ -654,7 +675,7 @@ def compute_confusion_matrix(
         self,
         score_thresholds: list[float] = [0.0],
         hardmax: bool = True,
-        datums: pc.Expression | None = None,
+        datums: Expression | None = None,
     ) -> list[Metric]:
         """
         Compute a confusion matrix.
@@ -673,6 +694,7 @@ def compute_confusion_matrix(
         list[Metric]
             A list of confusion matrices.
         """
+        datum_filter = datums.to_arrow() if datums is not None else None
         if not score_thresholds:
             raise ValueError("At least one score threshold must be passed.")
 
@@ -684,8 +706,8 @@ def compute_confusion_matrix(
         unmatched_groundtruths = np.zeros(
             (n_scores, n_labels), dtype=np.uint64
         )
-        for ids, scores, winners, matches in self.iterate_values(
-            datums=datums
+        for ids, scores, winners, matches in self._iterate_values(
+            datum_filter=datum_filter
         ):
             (
                 mask_tp,
@@ -722,7 +744,7 @@ def compute_examples(
         self,
         score_thresholds: list[float] = [0.0],
         hardmax: bool = True,
-        datums: pc.Expression | None = None,
+        datums: Expression | None = None,
     ) -> list[Metric]:
         """
         Compute examples per datum.
@@ -743,6 +765,7 @@ def compute_examples(
         list[Metric]
             A list of confusion matrices.
         """
+        datum_filter = datums.to_arrow() if datums is not None else None
         if not score_thresholds:
             raise ValueError("At least one score threshold must be passed.")
 
@@ -753,7 +776,7 @@ def compute_examples(
             winners,
             _,
             tbl,
-        ) in self.iterate_values_with_tables(datums=datums):
+        ) in self._iterate_values_with_tables(datum_filter=datum_filter):
             if ids.size == 0:
                 continue
 
@@ -795,7 +818,7 @@ def compute_confusion_matrix_with_examples(
         self,
         score_thresholds: list[float] = [0.0],
         hardmax: bool = True,
-        datums: pc.Expression | None = None,
+        datums: Expression | None = None,
     ) -> list[Metric]:
         """
         Compute confusion matrix with examples.
@@ -818,6 +841,7 @@ def compute_confusion_matrix_with_examples(
         list[Metric]
             A list of confusion matrices.
         """
+        datum_filter = datums.to_arrow() if datums is not None else None
         if not score_thresholds:
             raise ValueError("At least one score threshold must be passed.")
 
@@ -835,7 +859,7 @@ def compute_confusion_matrix_with_examples(
             winners,
             _,
             tbl,
-        ) in self.iterate_values_with_tables(datums=datums):
+        ) in self._iterate_values_with_tables(datum_filter=datum_filter):
             if ids.size == 0:
                 continue
 

src/valor_lite/classification/shared.py

@@ -7,6 +7,7 @@
 from numpy.typing import NDArray
 
 from valor_lite.cache import FileCacheReader, MemoryCacheReader
+from valor_lite.filtering import DataType, Expression
 
 
 @dataclass
@@ -34,7 +35,9 @@ def generate_metadata_path(path: str | Path) -> Path:
 
 
 def generate_schema(
-    metadata_fields: list[tuple[str, str | pa.DataType]] | None
+    metadata_fields: list[tuple[str, DataType]]
+    | list[tuple[str, str]]
+    | None = None
 ) -> pa.Schema:
     metadata_fields = metadata_fields if metadata_fields else []
     reserved_fields = [
@@ -59,7 +62,15 @@ def generate_schema(
         raise ValueError(
             f"metadata fields {conflicting} conflict with reserved fields"
         )
-    return pa.schema(reserved_fields + metadata_fields)
+    return pa.schema(
+        reserved_fields
+        + [
+            (name, dtype.to_arrow())
+            if isinstance(dtype, DataType)
+            else (name, dtype)
+            for name, dtype in metadata_fields
+        ]
+    )
 
 
 def generate_intermediate_schema() -> pa.Schema:
@@ -83,7 +94,9 @@ def generate_roc_curve_schema() -> pa.Schema:
 
 
 def encode_metadata_fields(
-    metadata_fields: list[tuple[str, str | pa.DataType]] | None
+    metadata_fields: list[tuple[str, DataType]]
+    | list[tuple[str, str]]
+    | None = None,
 ) -> dict[str, str]:
     metadata_fields = metadata_fields if metadata_fields else []
     return {k: str(v) for k, v in metadata_fields}
@@ -133,10 +146,11 @@ def extract_labels(
 
 def extract_counts(
     reader: MemoryCacheReader | FileCacheReader,
-    datums: pc.Expression | None = None,
+    datums: Expression | None = None,
 ):
     n_dts = 0
-    for tbl in reader.iterate_tables(filter=datums):
+    datum_filter = datums.to_arrow() if datums is not None else None
+    for tbl in reader.iterate_tables(filter=datum_filter):
         n_dts += int(np.unique(tbl["datum_id"].to_numpy()).shape[0])
     return n_dts
 

src/valor_lite/filtering.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import StrEnum
-from typing import Any
+from typing import Any, Iterable
 from zoneinfo import ZoneInfo
 
 import pyarrow as pa
@@ -167,7 +167,7 @@ def __le__(self, other: Any) -> Expression:
         other = other._symbol if isinstance(other, _Symbol) else other
         return Expression(self._symbol <= other)
 
-    def isin(self, values: set[Any]) -> Expression:
+    def isin(self, values: Iterable[Any]) -> Expression:
         values = {v._symbol if isinstance(v, _Symbol) else v for v in values}
         return Expression(self._symbol.isin(values))