data type tests

Author: czaloom

src/valor_lite/cache.py

@@ -17,7 +17,7 @@ class DataType(StrEnum):
     STRING = "string"
     TIMESTAMP = "timestamp"
 
-    def to_type(self):
+    def to_py(self):
         match self:
             case DataType.INTEGER:
                 return int
@@ -37,7 +37,7 @@ def to_arrow(self):
             case DataType.STRING:
                 return pa.string()
             case DataType.TIMESTAMP:
-                return pa.timestamp()
+                return pa.timestamp("us")
 
 
 class CacheReader:

src/valor_lite/object_detection/loader.py

@@ -27,10 +27,8 @@ def __init__(
         self,
         directory: str | Path = ".valor",
         name: str = "default",
-        batch_size: int = 1,
-        rows_per_file: int = 1,
-        # batch_size: int = 10_000,
-        # rows_per_file: int = 1_000_000,
+        batch_size: int = 10_000,
+        rows_per_file: int = 100_000,
         compression: str = "snappy",
         datum_metadata_types: dict[str, DataType] | None = None,
         groundtruth_metadata_types: dict[str, DataType] | None = None,

tests/common/test_cache.py

@@ -1,10 +1,25 @@
 import tempfile
+from datetime import datetime
 from pathlib import Path
 
 import numpy as np
 import pyarrow as pa
 
-from valor_lite.cache import CacheWriter
+from valor_lite.cache import CacheWriter, DataType
+
+
+def test_datatype_casting_to_arrow():
+    assert DataType.FLOAT.to_arrow() == pa.float64()
+    assert DataType.INTEGER.to_arrow() == pa.int64()
+    assert DataType.STRING.to_arrow() == pa.string()
+    assert DataType.TIMESTAMP.to_arrow() == pa.timestamp("us")
+
+
+def test_datatype_casting_to_python():
+    assert DataType.FLOAT.to_py() is float
+    assert DataType.INTEGER.to_py() is int
+    assert DataType.STRING.to_py() is str
+    assert DataType.TIMESTAMP.to_py() is datetime
 
 
 def test_cache_write_batch():