merge main

Author: czaloom

.gitignore

@@ -30,4 +30,9 @@ site/*
 *.png
 *.jpg
 *.parquet
-*.valor
+<<<<<<< HEAD
+*.valor
+=======
+
+.valor/*
+>>>>>>> origin/main

src/valor_lite/cache/compute.py

@@ -16,7 +16,7 @@ def _merge(
     batch_size: int,
     sorting: list[tuple[str, str]],
     columns: list[str] | None = None,
-    table_sort_override: Callable[[pa.table], pa.table] | None = None,
+    table_sort_override: Callable[[pa.Table], pa.Table] | None = None,
 ):
     """Merge locally sorted cache fragments."""
     for tbl in source.iterate_tables(columns=columns):
@@ -75,7 +75,7 @@ def create_sort_key(
                     create_sort_key(batches, batch_idx, 0),
                 )
 
-        sink.flush()
+    sink.flush()
 
 
 def sort(
@@ -84,7 +84,7 @@ def sort(
     batch_size: int,
     sorting: list[tuple[str, str]],
     columns: list[str] | None = None,
-    table_sort_override: Callable[[pa.table], pa.table] | None = None,
+    table_sort_override: Callable[[pa.Table], pa.Table] | None = None,
 ):
     """
     Sort data into new cache.
@@ -100,9 +100,9 @@ def sort(
     sorting : list[tuple[str, str]]
         Sorting arguments in PyArrow format (e.g. [('a', 'ascending'), ('b', 'descending')]).
         Note that only numeric fields are supported.
-    columns : list[tuple[str, str]], optional
+    columns : list[str], optional
         Option to only read a subset of columns.
-    table_sort_override : Callable[[pa.table], pa.table], optional
+    table_sort_override : Callable[[pa.Table], pa.Table], optional
         Option to override sort function for singular cache fragments.
     """
 

src/valor_lite/cache/ephemeral.py

@@ -106,25 +106,30 @@ def iterate_tables_with_arrays(
         Iterator[tuple[pa.Table, np.ndarray]]
 
         """
+        _columns = set(columns) if columns else set()
+        _numeric_columns = set(numeric_columns) if numeric_columns else set()
+        columns = list(_columns.union(_numeric_columns))
         for tbl in self.iterate_tables(
             columns=columns,
             filter=filter,
         ):
-            columns = numeric_columns if numeric_columns else tbl.columns
+            table_columns = numeric_columns if numeric_columns else tbl.columns
             yield tbl, np.column_stack(
-                [tbl[col].to_numpy() for col in columns]
+                [tbl[col].to_numpy() for col in table_columns]
             )
 
     def iterate_fragments(
         self, batch_size: int
     ) -> Iterator[Iterator[pa.RecordBatch]]:
         """
-        Iterate over fragment batch iterators within the file-based cache.
+        Yield a table batch iterator.
+
+        This is intended to emulate file-based access patterns.
 
         Parameters
         ----------
         batch_size : int
-            Maximum number of rows allowed to be read into memory per cache file.
+            Maximum number of rows allowed to be read per batch.
 
         Yields
         ------

src/valor_lite/cache/persistent.py

@@ -231,13 +231,16 @@ def iterate_tables_with_arrays(
         Iterator[tuple[pa.Table, np.ndarray]]
 
         """
+        _columns = set(columns) if columns else set()
+        _numeric_columns = set(numeric_columns) if numeric_columns else set()
+        columns = list(_columns.union(_numeric_columns))
         for tbl in self.iterate_tables(
             columns=columns,
             filter=filter,
         ):
-            columns = numeric_columns if numeric_columns else tbl.columns
+            table_columns = numeric_columns if numeric_columns else tbl.columns
             yield tbl, np.column_stack(
-                [tbl[col].to_numpy() for col in columns]
+                [tbl[col].to_numpy() for col in table_columns]
             )
 
     def iterate_fragments(