fix _set_vector_unsafe and add tests on _nb_fill

Author: ggmarshall

.gitignore

@@ -1,5 +1,8 @@
 /src/lgdo/_version.py
 
+#uv
+uv.lock
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

.pre-commit-config.yaml

@@ -50,6 +50,8 @@ repos:
     rev: "v2.4.1"
     hooks:
       - id: codespell
+        additional_dependencies:
+          - tomli
 
   - repo: https://github.com/shellcheck-py/shellcheck-py
     rev: "v0.10.0.1"

src/lgdo/types/vectorofvectors.py

@@ -476,7 +476,10 @@ def _set_vector_unsafe(
             else:
                 nan_val = np.nan
             vovutils._nb_fill(
-                vec, lens, nan_val, self.flattened_data.nda[start : cum_lens[-1]]
+                vec,
+                lens,
+                np.array([nan_val]).astype(self.flattened_data.nda.dtype),
+                self.flattened_data.nda[start : cum_lens[-1]],
             )
 
             # add new vector(s) length to cumulative_length

src/lgdo/types/vovutils.py

@@ -131,7 +131,7 @@ def _nb_fill(
     for i, ll in enumerate(len_in):
         stop = start + ll
         if ll > max_len:
-            flattened_array_out[start : start + max_len] = aoa_in[i, :]
+            flattened_array_out[start : start + max_len] = aoa_in[i, :max_len]
             flattened_array_out[start + max_len : stop] = nan_val
         else:
             flattened_array_out[start:stop] = aoa_in[i, :ll]

tests/types/test_vovutils.py

@@ -27,6 +27,58 @@ def testvov():
     return VovColl(v2d, v3d, v4d)
 
 
+def test_nb_fill():
+    # test 1d array of int
+    aoa_in = np.arange(5, dtype="int32").reshape(1, 5)
+    len_in = np.array([5])
+    nan_val = np.array([0], dtype=aoa_in.dtype)
+    flattened_array_out = np.empty(5, dtype=aoa_in.dtype)
+
+    vovutils._nb_fill(aoa_in, len_in, nan_val, flattened_array_out)
+    assert np.array_equal(
+        flattened_array_out, np.array([0, 1, 2, 3, 4], dtype=aoa_in.dtype)
+    )
+    # test 1d array of uint
+    aoa_in = np.arange(5, dtype="uint16").reshape(1, 5)
+    len_in = np.array([5])
+    nan_val = np.array([0], dtype=aoa_in.dtype)
+    flattened_array_out = np.empty(5, dtype=aoa_in.dtype)
+
+    vovutils._nb_fill(aoa_in, len_in, nan_val, flattened_array_out)
+    assert np.array_equal(
+        flattened_array_out, np.array([0, 1, 2, 3, 4], dtype=aoa_in.dtype)
+    )
+    # test 1d array of float
+    aoa_in = np.arange(5, dtype="float32").reshape(1, 5)
+    len_in = np.array([5])
+    nan_val = np.array([0], dtype=aoa_in.dtype)
+    flattened_array_out = np.empty(5, dtype=aoa_in.dtype)
+
+    vovutils._nb_fill(aoa_in, len_in, nan_val, flattened_array_out)
+    assert np.array_equal(
+        flattened_array_out, np.array([0, 1, 2, 3, 4], dtype=aoa_in.dtype)
+    )
+    # test 2d array of int
+    aoa_in = np.array([[1, 2, 3], [4, 5, 6]], dtype="int32")
+    len_in = np.array([3, 3])
+    nan_val = np.array([0], dtype=aoa_in.dtype)
+    flattened_array_out = np.empty(6, dtype=aoa_in.dtype)
+    vovutils._nb_fill(aoa_in, len_in, nan_val, flattened_array_out)
+    assert np.array_equal(
+        flattened_array_out,
+        np.array([1, 2, 3, 4, 5, 6], dtype=aoa_in.dtype),
+    )
+    # test nan value addition
+    aoa_in = np.array([[1, 2, 3], [4, 5, 6]], dtype="int32")
+    len_in = np.array([4, 3])
+    flattened_array_out = np.empty(7, dtype=aoa_in[0].dtype)
+    vovutils._nb_fill(aoa_in, len_in, nan_val, flattened_array_out)
+    assert np.array_equal(
+        flattened_array_out,
+        np.array([1, 2, 3, 0, 4, 5, 6], dtype=aoa_in[0].dtype),
+    )
+
+
 def test_ak_input_validity(testvov):
     for v in testvov:
         assert vovutils._ak_is_jagged(v) is True