pulled sorting out of loader

Author: czaloom

src/valor_lite/cache/__init__.py

@@ -0,0 +1,14 @@
+from .compute import heapsort
+from .datatype import DataType, convert_type_mapping_to_fields
+from .ephemeral import MemoryCacheReader, MemoryCacheWriter
+from .persistent import FileCacheReader, FileCacheWriter
+
+__all__ = [
+    "DataType",
+    "convert_type_mapping_to_fields",
+    "FileCacheReader",
+    "FileCacheWriter",
+    "MemoryCacheReader",
+    "MemoryCacheWriter",
+    "heapsort",
+]

src/valor_lite/cache/compute.py

@@ -0,0 +1,85 @@
+import heapq
+
+import pyarrow as pa
+
+from valor_lite.cache.ephemeral import MemoryCacheReader, MemoryCacheWriter
+from valor_lite.cache.persistent import FileCacheReader, FileCacheWriter
+
+
+def heapsort(
+    source: MemoryCacheReader | FileCacheReader,
+    sink: MemoryCacheWriter | FileCacheWriter,
+    batch_size: int,
+    sorting: list[tuple[str, str]],
+):
+    """
+    Perform heapsort on a cache object.
+
+    Parameters
+    ----------
+    source : MemoryCacheReader | FileCacheReader
+        The read-only source cache.
+    sink : MemoryCacheWriter | FileCacheWriter
+        The cache where sorted data will be written.
+    batch_size : int
+        Maximum number of rows allowed to be read into memory per cache file.
+    sorting : list[tuple[str, str]]
+        Sorting arguments in PyArrow format (e.g. [('a', 'ascending'), ('b', 'descending')]).
+    """
+    if source.count_tables() == 1 or isinstance(source, MemoryCacheReader):
+        for tbl in source.iterate_tables():
+            sorted_tbl = tbl.sort_by(sorting)
+            sink.write_table(sorted_tbl)
+    else:
+
+        def create_sort_key(
+            batches: list[pa.RecordBatch],
+            batch_idx: int,
+            row_idx: int,
+        ):
+            args = [
+                -batches[batch_idx][name][row_idx].as_py()
+                if direction == "descending"
+                else batches[batch_idx][name][row_idx].as_py()
+                for name, direction in sorting
+            ]
+            return (
+                *args,
+                batch_idx,
+                row_idx,
+            )
+
+        # merge sorted rows
+        heap = []
+        batch_iterators = []
+        batches = []
+        for batch_idx, batch_fragment in enumerate(source.iterate_fragments()):
+            batch_iter = batch_fragment.to_batches(batch_size=batch_size)
+            batch_iterators.append(batch_iter)
+            batches.append(next(batch_iterators[batch_idx], None))
+            if batches[batch_idx] is not None and len(batches[batch_idx]) > 0:
+                heapq.heappush(heap, create_sort_key(batches, batch_idx, 0))
+
+        while heap:
+            _, _, batch_idx, row_idx = heapq.heappop(heap)
+            row_table = batches[batch_idx].slice(row_idx, 1)
+            sink.write_batch(row_table)
+            row_idx += 1
+            if row_idx < len(batches[batch_idx]):
+                heapq.heappush(
+                    heap,
+                    create_sort_key(batches, batch_idx, row_idx),
+                )
+            else:
+                batches[batch_idx] = next(batch_iterators[batch_idx], None)
+                if (
+                    batches[batch_idx] is not None
+                    and len(batches[batch_idx]) > 0
+                ):
+                    heapq.heappush(
+                        heap,
+                        create_sort_key(batches, batch_idx, 0),
+                    )
+
+    # flush any buffers
+    sink.flush()

src/valor_lite/common/datatype.py src/valor_lite/cache/datatype.pysrc/valor_lite/common/datatype.py renamed to src/valor_lite/cache/datatype.py

@@ -36,7 +36,7 @@ def to_arrow(self):
                 return pa.timestamp("us")
 
 
-def convert_type_mapping_to_schema(
+def convert_type_mapping_to_fields(
     type_mapping: dict[str, DataType] | None
 ) -> list[tuple[str, pl.DataType]]:
     """

src/valor_lite/common/ephemeral.py src/valor_lite/cache/ephemeral.pysrc/valor_lite/common/ephemeral.py renamed to src/valor_lite/cache/ephemeral.py

@@ -8,9 +8,7 @@
 import pyarrow.compute as pc
 from numpy.typing import NDArray
 
-from valor_lite.common.datatype import DataType
-from valor_lite.common.ephemeral import MemoryCacheReader
-from valor_lite.common.persistent import FileCacheReader
+from valor_lite.cache import DataType, FileCacheReader, MemoryCacheReader
 from valor_lite.object_detection.computation import (
     compute_average_precision,
     compute_average_recall,

src/valor_lite/common/persistent.py src/valor_lite/cache/persistent.pysrc/valor_lite/common/persistent.py renamed to src/valor_lite/cache/persistent.py

@@ -1,4 +1,3 @@
-import heapq
 import json
 from pathlib import Path
 
@@ -7,9 +6,13 @@
 from numpy.typing import NDArray
 from tqdm import tqdm
 
-from valor_lite.common.datatype import DataType, convert_type_mapping_to_schema
-from valor_lite.common.ephemeral import MemoryCacheReader, MemoryCacheWriter
-from valor_lite.common.persistent import FileCacheWriter
+from valor_lite.cache import (
+    DataType,
+    FileCacheWriter,
+    MemoryCacheWriter,
+    convert_type_mapping_to_fields,
+    heapsort,
+)
 from valor_lite.exceptions import EmptyCacheError
 from valor_lite.object_detection.annotation import (
     Bitmask,
@@ -58,13 +61,13 @@ def in_memory(
         groundtruth_metadata_types: dict[str, DataType] | None = None,
         prediction_metadata_types: dict[str, DataType] | None = None,
     ):
-        datum_metadata_fields = convert_type_mapping_to_schema(
+        datum_metadata_fields = convert_type_mapping_to_fields(
             datum_metadata_types
         )
-        groundtruth_metadata_fields = convert_type_mapping_to_schema(
+        groundtruth_metadata_fields = convert_type_mapping_to_fields(
             groundtruth_metadata_types
         )
-        prediction_metadata_fields = convert_type_mapping_to_schema(
+        prediction_metadata_fields = convert_type_mapping_to_fields(
             prediction_metadata_types
         )
 
@@ -108,13 +111,13 @@ def persistent(
         if delete_if_exists and path.exists():
             cls.delete(path)
 
-        datum_metadata_fields = convert_type_mapping_to_schema(
+        datum_metadata_fields = convert_type_mapping_to_fields(
             datum_metadata_types
         )
-        groundtruth_metadata_fields = convert_type_mapping_to_schema(
+        groundtruth_metadata_fields = convert_type_mapping_to_fields(
             groundtruth_metadata_types
         )
-        prediction_metadata_fields = convert_type_mapping_to_schema(
+        prediction_metadata_fields = convert_type_mapping_to_fields(
             prediction_metadata_types
         )
 
@@ -488,88 +491,47 @@ def rank(
             for field in self._ranked_writer.schema
             if field.name not in {"high_score", "iou_prev"}
         ]
-        if (
-            isinstance(detailed_reader, MemoryCacheReader)
-            or detailed_reader.count_tables() == 1
-        ):
-            for tbl in detailed_reader.iterate_tables(columns=subset_columns):
-                ranked_tbl = rank_table(tbl, n_labels)
-                self._ranked_writer.write_table(ranked_tbl)
-        elif isinstance(self._ranked_writer, FileCacheWriter):
+        if isinstance(self._ranked_writer, FileCacheWriter):
             if not self._path:
                 raise ValueError(
                     "missing path definition in file-based loader"
                 )
             path = self._generate_temporary_cache_path(self._path)
-            with FileCacheWriter.create(
+            tmp_sink = FileCacheWriter.create(
                 path=path,
-                schema=self._ranked_writer._schema,
+                schema=self._ranked_writer.schema,
                 batch_size=self._ranked_writer._batch_size,
                 rows_per_file=self._ranked_writer._rows_per_file,
                 compression=self._ranked_writer._compression,
                 delete_if_exists=True,
-            ) as tmp_writer:
-
-                # rank individual files
-                for tbl in detailed_reader.iterate_tables(
-                    columns=subset_columns
-                ):
-                    ranked_tbl = rank_table(tbl, n_labels)
-                    tmp_writer.write_table(ranked_tbl)
-
-            tmp_reader = tmp_writer.to_reader()
-
-            def generate_heap_item(batches, batch_idx, row_idx) -> tuple:
-                score = batches[batch_idx]["score"][row_idx].as_py()
-                iou = batches[batch_idx]["iou"][row_idx].as_py()
-                return (
-                    -score,
-                    -iou,
-                    batch_idx,
-                    row_idx,
-                )
-
-            # merge sorted rows
-            heap = []
-            batch_iterators = []
-            batches = []
-            for batch_idx, batch_fragment in enumerate(
-                tmp_reader.iterate_fragments()
-            ):
-                batch_iter = batch_fragment.to_batches(batch_size=batch_size)
-                batch_iterators.append(batch_iter)
-                batches.append(next(batch_iterators[batch_idx], None))
-                if (
-                    batches[batch_idx] is not None
-                    and len(batches[batch_idx]) > 0
-                ):
-                    heapq.heappush(
-                        heap, generate_heap_item(batches, batch_idx, 0)
-                    )
-
-            while heap:
-                _, _, batch_idx, row_idx = heapq.heappop(heap)
-                row_table = batches[batch_idx].slice(row_idx, 1)
-                self._ranked_writer.write_batch(row_table)
-                row_idx += 1
-                if row_idx < len(batches[batch_idx]):
-                    heapq.heappush(
-                        heap,
-                        generate_heap_item(batches, batch_idx, row_idx),
-                    )
-                else:
-                    batches[batch_idx] = next(batch_iterators[batch_idx], None)
-                    if (
-                        batches[batch_idx] is not None
-                        and len(batches[batch_idx]) > 0
-                    ):
-                        heapq.heappush(
-                            heap,
-                            generate_heap_item(batches, batch_idx, 0),
-                        )
+            )
+        else:
+            tmp_sink = MemoryCacheWriter.create(
+                schema=self._ranked_writer.schema,
+                batch_size=self._ranked_writer._batch_size,
+            )
+
+        # rank individual files
+        for tbl in detailed_reader.iterate_tables(columns=subset_columns):
+            ranked_tbl = rank_table(tbl, n_labels)
+            tmp_sink.write_table(ranked_tbl)
+        tmp_source = tmp_sink.to_reader()
+
+        # sort ranked pairs across all chunks
+        heapsort(
+            source=tmp_source,
+            sink=self._ranked_writer,
+            batch_size=batch_size,
+            sorting=[
+                ("score", "descending"),
+                ("iou", "descending"),
+            ],
+        )
 
-            FileCacheWriter.delete(path)
+        # clean up
         self._ranked_writer.flush()
+        if isinstance(tmp_sink, FileCacheWriter):
+            FileCacheWriter.delete(tmp_sink.path)
 
     def finalize(
         self,

src/valor_lite/object_detection/evaluator.py

@@ -2,7 +2,7 @@
 
 import pyarrow as pa
 
-from valor_lite.common.datatype import DataType, convert_type_mapping_to_schema
+from valor_lite.cache.datatype import DataType, convert_type_mapping_to_fields
 
 
 def test_datatype_casting_to_arrow():
@@ -19,8 +19,8 @@ def test_datatype_casting_to_python():
     assert DataType.TIMESTAMP.to_py() is datetime
 
 
-def test_convert_type_mapping_to_schema():
-    x = convert_type_mapping_to_schema(
+def test_convert_type_mapping_to_fields():
+    x = convert_type_mapping_to_fields(
         {
             "a": DataType.FLOAT,
             "b": DataType.STRING,
@@ -31,5 +31,5 @@ def test_convert_type_mapping_to_schema():
         ("b", pa.string()),
     ]
 
-    assert convert_type_mapping_to_schema({}) == []
-    assert convert_type_mapping_to_schema(None) == []
+    assert convert_type_mapping_to_fields({}) == []
+    assert convert_type_mapping_to_fields(None) == []

src/valor_lite/object_detection/loader.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pyarrow as pa
 
-from valor_lite.common.ephemeral import MemoryCacheWriter
+from valor_lite.cache.ephemeral import MemoryCacheWriter
 
 
 def test_cache_reader():

tests/common/test_datatype.py

@@ -5,7 +5,7 @@
 import pyarrow as pa
 import pytest
 
-from valor_lite.common.persistent import (
+from valor_lite.cache.persistent import (
     FileCache,
     FileCacheReader,
     FileCacheWriter,