[AL-1443] Tiling + Slicing + Updating

hub/api/tests/test_api.py

@@ -117,10 +117,25 @@ def test_populate_dataset(ds):
 
 def test_larger_data_memory(memory_ds):
     memory_ds.create_tensor("image")
-    memory_ds.image.extend(np.ones((4, 4096, 4096)))
+    x = np.ones((4, 4096, 4096))
+    memory_ds.image.extend(x)
     assert len(memory_ds) == 4
-    assert memory_ds.image.shape == (4, 4096, 4096)
-    np.testing.assert_array_equal(memory_ds.image.numpy(), np.ones((4, 4096, 4096)))
+    assert memory_ds.image.shape == x.shape
+    np.testing.assert_array_equal(memory_ds.image.numpy(), x)
+    idxs = [
+        0,
+        1,
+        3,
+        -1,
+        slice(0, 3),
+        slice(2, 4),
+        slice(2, None),
+        (0, slice(5, None), slice(None, 714)),
+        (2, 100, 1007),
+        (slice(1, 3), [20, 1000, 2, 400], [-2, 3, 577, 4095]),
+    ]
+    for idx in idxs:
+        np.testing.assert_array_equal(memory_ds.image[idx].numpy(), x[idx])
 
 
 def test_stringify(memory_ds):

hub/api/tests/test_api_tiling.py

@@ -25,7 +25,18 @@ def test_mixed_small_large(local_ds_generator, compression):
     ds = local_ds_generator()
     arr1 = np.random.randint(0, 255, (3003, 2001, 3)).astype(np.uint8)
     arr2 = np.random.randint(0, 255, (500, 500, 3)).astype(np.uint8)
-    arr3 = np.random.randint(0, 255, (2503, 2501, 3)).astype(np.uint8)
+    arr3 = np.random.randint(0, 255, (250, 250, 3)).astype(np.uint8)
+
+    idxs = [
+        (slice(73, 117), slice(9, 17)),
+        4,
+        -1,
+        slice(
+            10,
+        ),
+        slice(20, 37),
+    ]
+
     with ds:
         ds.create_tensor("abc", **compression)
         for i in range(10):
@@ -37,15 +48,24 @@ def test_mixed_small_large(local_ds_generator, compression):
     for i in range(10):
         if i % 5 == 0:
             np.testing.assert_array_equal(ds.abc[i].numpy(), arr1)
+            for idx in idxs:
+                np.testing.assert_array_equal(ds.abc[i][idx].numpy(), arr1[idx])
         else:
             np.testing.assert_array_equal(ds.abc[i].numpy(), arr2)
+            for idx in idxs:
+                np.testing.assert_array_equal(ds.abc[i][idx].numpy(), arr2[idx])
 
     ds = local_ds_generator()
+
     for i in range(10):
         if i % 5 == 0:
             np.testing.assert_array_equal(ds.abc[i].numpy(), arr1)
+            for idx in idxs:
+                np.testing.assert_array_equal(ds.abc[i][idx].numpy(), arr1[idx])
         else:
             np.testing.assert_array_equal(ds.abc[i].numpy(), arr2)
+            for idx in idxs:
+                np.testing.assert_array_equal(ds.abc[i][idx].numpy(), arr2[idx])
 
     with ds:
         ds.abc.extend([arr3] * 3)
@@ -74,27 +94,42 @@ def test_updates(memory_ds, compression):
     arr2 = np.random.randint(0, 255, (500, 500, 3)).astype(np.uint8)
     arr3 = np.random.randint(0, 255, (2503, 2501, 3)).astype(np.uint8)
     arr4 = np.random.randint(0, 255, (250, 250, 3)).astype(np.uint8)
+
+    update_idx = (slice(73, 117), slice(1765, 1901))
+
+    arr5 = arr1 * 2
+    arr6 = arr5[update_idx]
+    arr6 += 1
+
     with memory_ds:
-        memory_ds.create_tensor("abc", **compression)
+        memory_ds.create_tensor("abc", max_chunk_size=2 ** 21, **compression)
         for i in range(10):
             if i % 5 == 0:
                 memory_ds.abc.append(arr1)
             else:
                 memory_ds.abc.append(arr2)
-
+            len(memory_ds)
     with memory_ds:
         for i in range(10):
-            if i % 5 != 0:
+            if i % 5 == 0:
+                memory_ds.abc[i] = arr1 * 2
+                memory_ds.abc[i][update_idx] = arr6
+            else:
                 memory_ds.abc[i] = arr3 if i % 2 == 0 else arr4
 
     for i in range(10):
         if i % 5 == 0:
-            np.testing.assert_array_equal(memory_ds.abc[i].numpy(), arr1)
+            np.testing.assert_array_equal(memory_ds.abc[i].numpy(), arr5)
         elif i % 2 == 0:
             np.testing.assert_array_equal(memory_ds.abc[i].numpy(), arr3)
         else:
             np.testing.assert_array_equal(memory_ds.abc[i].numpy(), arr4)
 
+    # update tiled sample with small sample
+    arr7 = np.random.randint(0, 255, (3, 2, 3)).astype(np.uint8)
+    memory_ds.abc[0] = arr7
+    np.testing.assert_array_equal(memory_ds.abc[0].numpy(), arr7)
+
 
 def test_cachable_overflow(memory_ds):
     ds = memory_ds

hub/core/chunk/base_chunk.py

@@ -77,6 +77,9 @@ def __init__(
         self._decompressed_samples: Optional[List[np.ndarray]] = None
         self._decompressed_bytes: Optional[bytes] = None
 
+        # Whether tensor meta is updated by chunk. Used by chunk engine.
+        self._update_meta: bool = True
+
     @property
     def num_data_bytes(self) -> int:
         return len(self.data_bytes)
@@ -127,7 +130,7 @@ def frombuffer(cls, buffer: bytes, chunk_args: list, copy=True):  # type: ignore
         return chunk
 
     @abstractmethod
-    def extend_if_has_space(self, incoming_samples) -> float:
+    def extend_if_has_space(self, incoming_samples, update_meta: bool = True) -> float:
         """Extends the chunk with the incoming samples."""
 
     @abstractmethod
@@ -242,8 +245,9 @@ def copy(self, chunk_args=None):
     def register_in_meta_and_headers(self, sample_nbytes: Optional[int], shape):
         """Registers a new sample in meta and headers"""
         self.register_sample_to_headers(sample_nbytes, shape)
-        self.tensor_meta.length += 1
-        self.tensor_meta.update_shape_interval(shape)
+        if self._update_meta:
+            self.tensor_meta.length += 1
+            self.tensor_meta.update_shape_interval(shape)
 
     def update_in_meta_and_headers(
         self, local_index: int, sample_nbytes: Optional[int], shape
@@ -286,7 +290,7 @@ def write_tile(self, sample: SampleTiles, skip_bytes=False):
         data, tile_shape = sample.yield_tile()
         sample_nbytes = None if skip_bytes else len(data)
         self.data_bytes = data
-        update_meta = sample.is_first_write
+        update_meta = self._update_meta and sample.is_first_write
         self.register_sample_to_headers(sample_nbytes, tile_shape)
         if update_meta:
             self.tensor_meta.length += 1

hub/core/chunk/chunk_compressed_chunk.py

@@ -16,7 +16,7 @@
 
 
 class ChunkCompressedChunk(BaseChunk):
-    def extend_if_has_space(self, incoming_samples: List[InputSample]) -> float:
+    def extend_if_has_space(self, incoming_samples: List[InputSample]) -> float:  # type: ignore
         self.prepare_for_write()
         if self.is_byte_compression:
             return self.extend_if_has_space_byte_compression(incoming_samples)

hub/core/chunk/sample_compressed_chunk.py

@@ -10,7 +10,7 @@
 
 
 class SampleCompressedChunk(BaseChunk):
-    def extend_if_has_space(self, incoming_samples: List[InputSample]) -> float:
+    def extend_if_has_space(self, incoming_samples: List[InputSample]) -> float:  # type: ignore
         self.prepare_for_write()
         num_samples: float = 0
         dtype = self.dtype if self.is_byte_compression else None

hub/core/chunk/uncompressed_chunk.py

@@ -7,7 +7,7 @@
 
 
 class UncompressedChunk(BaseChunk):
-    def extend_if_has_space(
+    def extend_if_has_space(  # type: ignore
         self, incoming_samples: Union[List[InputSample], np.ndarray]
     ) -> float:
         self.prepare_for_write()