Add virtual ref support EAR-1183

.github/workflows/rust-ci.yaml

@@ -81,6 +81,8 @@ jobs:
         env:
           AWS_ACCESS_KEY_ID: minio123
           AWS_SECRET_ACCESS_KEY: minio123
+          AWS_ALLOW_HTTP: 1
+          AWS_ENDPOINT_URL: http://localhost:9000
           AWS_DEFAULT_REGION: "us-east-1"
         run: |
           just pre-commit

compose.yaml

@@ -3,7 +3,7 @@ volumes:
 
 services:
   minio:
-    container_name: minio
+    container_name: icechunk_minio
     image: quay.io/minio/minio
     entrypoint: |
         /bin/sh -c '

icechunk/Cargo.toml

@@ -25,6 +25,7 @@ base32 = "0.5.1"
 chrono = { version = "0.4.38", features = ["serde"] }
 async-recursion = "1.1.1"
 rmp-serde = "1.3.0"
+url = "2.5.2"
 
 [dev-dependencies]
 pretty_assertions = "1.4.1"

icechunk/src/dataset.rs

@@ -6,14 +6,18 @@ use std::{
     sync::Arc,
 };
 
+use crate::storage::virtual_ref::VirtualChunkResolver;
 pub use crate::{
-    format::{manifest::ChunkPayload, snapshot::ZarrArrayMetadata, ChunkIndices, Path},
+    format::{
+        manifest::{ChunkPayload, VirtualChunkLocation},
+        snapshot::ZarrArrayMetadata,
+        ChunkIndices, Path,
+    },
     metadata::{
         ArrayShape, ChunkKeyEncoding, ChunkShape, Codec, DataType, DimensionName,
         DimensionNames, FillValue, StorageTransformer, UserAttributes,
     },
 };
-
 use bytes::Bytes;
 use chrono::Utc;
 use futures::{future::ready, Stream, StreamExt, TryStreamExt};
@@ -23,7 +27,9 @@ use tokio::task;
 
 use crate::{
     format::{
-        manifest::{ChunkInfo, ChunkRef, Manifest, ManifestExtents, ManifestRef},
+        manifest::{
+            ChunkInfo, ChunkRef, Manifest, ManifestExtents, ManifestRef, VirtualChunkRef,
+        },
         snapshot::{
             NodeData, NodeSnapshot, NodeType, Snapshot, SnapshotMetadata,
             SnapshotProperties, UserAttributesSnapshot,
@@ -34,6 +40,7 @@ use crate::{
         create_tag, fetch_branch_tip, fetch_tag, update_branch, BranchVersion, Ref,
         RefError,
     },
+    storage::virtual_ref::ObjectStoreVirtualChunkResolver,
     MemCachingStorage, Storage, StorageError,
 };
 
@@ -54,13 +61,14 @@ impl Default for DatasetConfig {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub struct Dataset {
     config: DatasetConfig,
     storage: Arc<dyn Storage + Send + Sync>,
     snapshot_id: ObjectId,
     last_node_id: Option<NodeId>,
     change_set: ChangeSet,
+    virtual_resolver: Arc<dyn VirtualChunkResolver + Send + Sync>,
 }
 
 #[derive(Clone, Debug, PartialEq, Default)]
@@ -342,6 +350,7 @@ impl Dataset {
             storage,
             last_node_id: None,
             change_set: ChangeSet::default(),
+            virtual_resolver: Arc::new(ObjectStoreVirtualChunkResolver::default()),
         }
     }
 
@@ -645,8 +654,18 @@ impl Dataset {
                 Ok(self.storage.fetch_chunk(&id, byte_range).await.map(Some)?)
             }
             Some(ChunkPayload::Inline(bytes)) => Ok(Some(byte_range.slice(bytes))),
-            //FIXME: implement virtual fetch
-            Some(ChunkPayload::Virtual(_)) => todo!(),
+            Some(ChunkPayload::Virtual(VirtualChunkRef { location, offset, length })) => {
+                Ok(self
+                    .virtual_resolver
+                    .fetch_chunk(
+                        &location,
+                        &byte_range.replace_unbounded_with(
+                            &ByteRange::from_offset_to_length(offset, length),
+                        ),
+                    )
+                    .await
+                    .map(Some)?)
+            }
             None => Ok(None),
         }
     }

icechunk/src/format/manifest.rs

@@ -1,9 +1,14 @@
+use itertools::Itertools;
 use std::{collections::BTreeMap, ops::Bound, sync::Arc};
+use thiserror::Error;
 
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 
-use super::{ChunkIndices, Flags, IcechunkFormatError, IcechunkResult, NodeId, ObjectId};
+use super::{
+    ChunkIndices, ChunkLength, ChunkOffset, Flags, IcechunkFormatError, IcechunkResult,
+    NodeId, ObjectId,
+};
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ManifestExtents(pub Vec<ChunkIndices>);
@@ -15,18 +20,64 @@ pub struct ManifestRef {
     pub extents: ManifestExtents,
 }
 
+#[derive(Debug, Error)]
+pub enum VirtualReferenceError {
+    #[error("error parsing virtual ref URL {0}")]
+    CannotParseUrl(#[from] url::ParseError),
+    #[error("virtual reference has no path segments {0}")]
+    NoPathSegments(String),
+    #[error("unsupported scheme for virtual chunk refs: {0}")]
+    UnsupportedScheme(String),
+    #[error("error parsing bucket name from virtual ref URL {0}")]
+    CannotParseBucketName(String),
+    #[error("error parsing virtual reference {0}")]
+    OtherError(#[from] Box<dyn std::error::Error + Send + Sync>),
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+pub enum VirtualChunkLocation {
+    Absolute(String),
+    // Relative(prefix_id, String)
+}
+
+impl VirtualChunkLocation {
+    pub fn from_absolute_path(
+        path: &str,
+    ) -> Result<VirtualChunkLocation, VirtualReferenceError> {
+        // make sure we can parse the provided URL before creating the enum
+        // TODO: consider other validation here.
+        let url = url::Url::parse(path)?;
+        let new_path: String = url
+            .path_segments()
+            .ok_or(VirtualReferenceError::NoPathSegments(path.into()))?
+            // strip empty segments here, object_store cannot handle them.
+            .filter(|x| !x.is_empty())
+            .join("/");
+
+        let host = url
+            .host()
+            .ok_or_else(|| VirtualReferenceError::CannotParseBucketName(path.into()))?;
+        Ok(VirtualChunkLocation::Absolute(format!(
+            "{}://{}/{}",
+            url.scheme(),
+            host,
+            new_path,
+        )))
+    }
+}
+
 #[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct VirtualChunkRef {
-    location: String, // FIXME: better type
-    offset: u64,
-    length: u64,
+    pub location: VirtualChunkLocation,
+    pub offset: ChunkOffset,
+    pub length: ChunkLength,
 }
 
 #[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct ChunkRef {
     pub id: ObjectId,
-    pub offset: u64,
-    pub length: u64,
+    pub offset: ChunkOffset,
+    pub length: ChunkLength,
 }
 
 #[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]

icechunk/src/format/mod.rs

@@ -105,6 +105,10 @@ impl ByteRange {
         Self(Bound::Included(offset), Bound::Unbounded)
     }
 
+    pub fn from_offset_to_length(offset: ChunkOffset, length: ChunkOffset) -> Self {
+        Self(Bound::Included(offset), Bound::Excluded(offset + length))
+    }
+
     pub fn to_offset(offset: ChunkOffset) -> Self {
         Self(Bound::Unbounded, Bound::Excluded(offset))
     }
@@ -115,6 +119,17 @@ impl ByteRange {
 
     pub const ALL: Self = Self(Bound::Unbounded, Bound::Unbounded);
 
+    pub fn replace_unbounded_with(&self, other: &ByteRange) -> Self {
+        let first = match self.0 {
+            Bound::Unbounded => other.0,
+            _ => self.0,
+        };
+        let second = match self.1 {
+            Bound::Unbounded => other.1,
+            _ => self.1,
+        };
+        ByteRange(first, second)
+    }
     pub fn slice(&self, bytes: Bytes) -> Bytes {
         match (self.0, self.1) {
             (Bound::Included(start), Bound::Excluded(end)) => {

icechunk/src/storage/mod.rs

@@ -12,13 +12,16 @@ pub mod caching;
 pub mod logging;
 
 pub mod object_store;
+pub mod virtual_ref;
 
 pub use caching::MemCachingStorage;
 pub use object_store::ObjectStorage;
 
 use crate::format::{
-    attributes::AttributesTable, manifest::Manifest, snapshot::Snapshot, ByteRange,
-    ObjectId, Path,
+    attributes::AttributesTable,
+    manifest::{Manifest, VirtualReferenceError},
+    snapshot::Snapshot,
+    ByteRange, ObjectId, Path,
 };
 
 #[derive(Debug, Error)]
@@ -37,6 +40,8 @@ pub enum StorageError {
     RefAlreadyExists(String),
     #[error("ref not found: {0}")]
     RefNotFound(String),
+    #[error("error parsing virtual reference {0}")]
+    VirtualReferenceError(#[from] VirtualReferenceError),
     #[error("generic storage error: {0}")]
     OtherError(#[from] Arc<dyn std::error::Error + Sync + Send>),
     #[error("unknown storage error: {0}")]

icechunk/src/storage/object_store.rs

@@ -1,19 +1,17 @@
-use core::fmt;
-use std::{
-    fs::create_dir_all, future::ready, ops::Bound, path::Path as StdPath, sync::Arc,
+use crate::format::{
+    attributes::AttributesTable, manifest::Manifest, snapshot::Snapshot, ByteRange,
+    ObjectId,
 };
-
 use async_trait::async_trait;
 use bytes::Bytes;
+use core::fmt;
 use futures::{stream::BoxStream, StreamExt, TryStreamExt};
 use object_store::{
     local::LocalFileSystem, memory::InMemory, path::Path as ObjectPath, GetOptions,
     GetRange, ObjectStore, PutMode, PutOptions, PutPayload,
 };
-
-use crate::format::{
-    attributes::AttributesTable, manifest::Manifest, snapshot::Snapshot, ByteRange,
-    ObjectId,
+use std::{
+    fs::create_dir_all, future::ready, ops::Bound, path::Path as StdPath, sync::Arc,
 };
 
 use super::{Storage, StorageError, StorageResult};

icechunk/src/storage/virtual_ref.rs

@@ -0,0 +1,82 @@
+use crate::format::manifest::{VirtualChunkLocation, VirtualReferenceError};
+use crate::format::ByteRange;
+use crate::storage::StorageResult;
+use async_trait::async_trait;
+use bytes::Bytes;
+use object_store::{
+    aws::AmazonS3Builder, path::Path as ObjectPath, GetOptions, GetRange, ObjectStore,
+};
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::sync::Arc;
+use tokio::sync::RwLock;
+use url;
+
+#[async_trait]
+pub trait VirtualChunkResolver: Debug {
+    async fn fetch_chunk(
+        &self,
+        location: &VirtualChunkLocation,
+        range: &ByteRange,
+    ) -> StorageResult<Bytes>;
+}
+
+#[derive(PartialEq, Eq, Hash, Clone, Debug, Default)]
+struct StoreCacheKey(String, String);
+
+#[derive(Debug, Default)]
+pub struct ObjectStoreVirtualChunkResolver {
+    stores: Arc<RwLock<HashMap<StoreCacheKey, Arc<dyn ObjectStore>>>>,
+}
+
+#[async_trait]
+impl VirtualChunkResolver for ObjectStoreVirtualChunkResolver {
+    async fn fetch_chunk(
+        &self,
+        location: &VirtualChunkLocation,
+        range: &ByteRange,
+    ) -> StorageResult<Bytes> {
+        let VirtualChunkLocation::Absolute(location) = location;
+        let parsed =
+            url::Url::parse(location).map_err(VirtualReferenceError::CannotParseUrl)?;
+        let bucket_name = parsed
+            .host_str()
+            .ok_or(VirtualReferenceError::CannotParseBucketName(
+                "error parsing bucket name".into(),
+            ))?
+            .to_string();
+        let path = ObjectPath::parse(parsed.path())
+            .map_err(|e| VirtualReferenceError::OtherError(Box::new(e)))?;
+        let scheme = parsed.scheme();
+        let cache_key = StoreCacheKey(scheme.into(), bucket_name);
+
+        let options =
+            GetOptions { range: Option::<GetRange>::from(range), ..Default::default() };
+        let store = {
+            let stores = self.stores.read().await;
+            #[allow(clippy::expect_used)]
+            stores.get(&cache_key).map(Arc::clone)
+        };
+        let store = match store {
+            Some(store) => store,
+            None => {
+                let builder = match scheme {
+                    "s3" => AmazonS3Builder::from_env(),
+                    _ => {
+                        Err(VirtualReferenceError::UnsupportedScheme(scheme.to_string()))?
+                    }
+                };
+                let new_store: Arc<dyn ObjectStore> =
+                    Arc::new(builder.with_bucket_name(&cache_key.1).build()?);
+                {
+                    self.stores
+                        .write()
+                        .await
+                        .insert(cache_key.clone(), Arc::clone(&new_store));
+                }
+                new_store
+            }
+        };
+        Ok(store.get_opts(&path, options).await?.bytes().await?)
+    }
+}

icechunk/src/strategies.rs

@@ -18,19 +18,23 @@ pub fn node_paths() -> impl Strategy<Value = Path> {
     any::<PathBuf>()
 }
 
-#[allow(clippy::expect_used)]
-pub fn empty_datasets() -> impl Strategy<Value = Dataset> {
+prop_compose! {
+    #[allow(clippy::expect_used)]
+    pub fn empty_datasets()(_id in any::<u32>()) -> Dataset {
+    // _id is used as a hack to avoid using prop_oneof![Just(dataset)]
+    // Using Just requires Dataset impl Clone, which we do not want
+
     // FIXME: add storages strategy
     let storage = ObjectStorage::new_in_memory_store(None);
     let runtime = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
 
-    let dataset = runtime.block_on(async {
+    runtime.block_on(async {
         Dataset::init(Arc::new(storage), false)
             .await
             .expect("Failed to initialize dataset")
             .build()
-    });
-    prop_oneof![Just(dataset)]
+    })
+}
 }
 
 pub fn codecs() -> impl Strategy<Value = Vec<Codec>> {