implement get_partial_values

earth-mover · Sep 13, 2024 · c6d8782 · c6d8782
1 parent 95ac8dd
commit c6d8782
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git a/icechunk-python/python/icechunk/__init__.py b/icechunk-python/python/icechunk/__init__.py
@@ -69,6 +69,7 @@ async def get(
                 return None
         except ValueError as _e:
             # Zarr python expects None to be returned if the key does not exist
+            # but an IcechunkStore returns an error if the key does not exist
             return None
 
         return prototype.buffer.from_bytes(result)

diff --git a/icechunk-python/tests/test_zarr/test_store/test_icechunk_store.py b/icechunk-python/tests/test_zarr/test_store/test_icechunk_store.py
@@ -26,6 +26,7 @@ async def get(self, store: IcechunkStore, key: str) -> Buffer:
                 return None
         except ValueError as _e:
             # Zarr python expects None to be returned if the key does not exist
+            # but an IcechunkStore returns an error if the key does not exist
             return None
 
         return self.buffer_cls.from_bytes(result)
@@ -129,13 +130,25 @@ async def test_delete(self, store: IcechunkStore) -> None:
         await store.delete("foo/zarr.json")
         assert not await store.exists("foo/zarr.json")
 
-    @pytest.mark.xfail(reason="Invalid usage pattern with Icechunk")
     async def test_get_partial_values(
         self,
         store: IcechunkStore,
-        key_ranges: list[tuple[str, tuple[int | None, int | None]]],
     ) -> None:
-        await super().test_get_partial_values(store, key_ranges)
+        await self.set(
+            store, "zarr.json", self.buffer_cls.from_bytes(DEFAULT_GROUP_METADATA)
+        )
+        # read back just part of it
+        values = await store.get_partial_values(
+            default_buffer_prototype(),
+            [
+                ("zarr.json", (0, 5)),
+            ],
+        )
+
+        assert len(values) == 1
+        data = values[0].to_bytes()
+        assert len(data) == 5
+        assert data == DEFAULT_GROUP_METADATA[:5]
 
     async def test_set(self, store: IcechunkStore) -> None:
         await store.set("zarr.json", self.buffer_cls.from_bytes(DEFAULT_GROUP_METADATA))
@@ -144,7 +157,9 @@ async def test_set(self, store: IcechunkStore) -> None:
         assert result.to_bytes() == DEFAULT_GROUP_METADATA
 
     async def test_get(self, store: IcechunkStore) -> None:
-        await self.set(store, "zarr.json", self.buffer_cls.from_bytes(DEFAULT_GROUP_METADATA))
+        await self.set(
+            store, "zarr.json", self.buffer_cls.from_bytes(DEFAULT_GROUP_METADATA)
+        )
         assert await store.exists("zarr.json")
         result = await store.get("zarr.json", default_buffer_prototype())
         assert result.to_bytes() == DEFAULT_GROUP_METADATA
diff --git a/icechunk/src/zarr.rs b/icechunk/src/zarr.rs
@@ -174,13 +174,23 @@ impl Store {
     }
 
     // TODO: prototype argument
-    pub async fn get(&self, key: &str, _byte_range: &ByteRange) -> StoreResult<Bytes> {
-        match Key::parse(key)? {
+    pub async fn get(&self, key: &str, byte_range: &ByteRange) -> StoreResult<Bytes> {
+        let bytes = match Key::parse(key)? {
             Key::Metadata { node_path } => self.get_metadata(key, &node_path).await,
             Key::Chunk { node_path, coords } => {
                 self.get_chunk(key, node_path, coords).await
             }
-        }
+        }?;
+
+        let (start, end) = byte_range;
+        let bytes = match (start, end) {
+            (Some(start), Some(end)) => bytes.slice(*start as usize..*end as usize),
+            (Some(start), None) => bytes.slice(*start as usize..),
+            (None, Some(end)) => bytes.slice(..*end as usize),
+            (None, None) => bytes,
+        };
+
+        Ok(bytes)
     }
 
     /// Get all the requested keys concurrently.