[WIP] Optimize split manifest writes.

Author: dcherian

icechunk/src/change_set.rs

@@ -9,9 +9,10 @@ use itertools::{Either, Itertools as _};
 use serde::{Deserialize, Serialize};
 
 use crate::{
+    config::ManifestSplittingConfig,
     format::{
         ChunkIndices, NodeId, Path,
-        manifest::{ChunkInfo, ChunkPayload},
+        manifest::{ChunkInfo, ChunkPayload, ManifestExtents, ManifestSplits},
         snapshot::{ArrayShape, DimensionName, NodeData, NodeSnapshot},
     },
     session::SessionResult,
@@ -24,20 +25,81 @@ pub struct ArrayData {
     pub user_data: Bytes,
 }
 
+impl ManifestSplits {
+    pub fn which_extent(&self, coord: &ChunkIndices) -> SessionResult<&ManifestExtents> {
+        Ok(self.0.get(self.which(coord)?).expect("logic bug"))
+    }
+}
+
+#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
+pub struct SplitManifest {
+    from: Vec<u32>,
+    to: Vec<u32>,
+    // It's important we keep these sorted, we use this fact in TransactionLog creation
+    chunks: BTreeMap<ChunkIndices, Option<ChunkPayload>>,
+}
+
+impl SplitManifest {
+    pub fn update(&mut self, coord: ChunkIndices, data: Option<ChunkPayload>) {
+        if self.from.is_empty() {
+            debug_assert!(self.to.is_empty());
+            debug_assert!(self.chunks.is_empty());
+            // important to remember that `to` is not inclusive, so we need +1
+            let mut coord0 = coord.0.clone();
+            self.to.extend(coord0.iter().map(|n| *n + 1));
+            self.from.append(&mut coord0);
+        } else {
+            for (existing, coord0) in self.from.iter_mut().zip(coord.0.iter()) {
+                if coord0 < existing {
+                    *existing = *coord0
+                }
+            }
+            for (existing, coord0) in self.to.iter_mut().zip(coord.0.iter()) {
+                // important to remember that `to` is not inclusive, so we need +1
+                let range_value = coord0 + 1;
+                if range_value > *existing {
+                    *existing = range_value
+                }
+            }
+        }
+        self.chunks.insert(coord, data);
+    }
+
+    pub fn retain(&mut self, predicate: impl Fn(&ChunkIndices) -> bool) {
+        self.chunks.retain(|coord, _| {
+            if !predicate(coord) {
+                // FIXME: handle from, to updating
+                todo!();
+            } else {
+                false
+            }
+        })
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
 pub struct ChangeSet {
+    splitting: ManifestSplittingConfig,
+    // This is an optimization so that we needn't figure out the split sizes on every set.
+    // FIXME: this is needed to catch the case where the splitting config changes in a session
+    //        after chunks are written
+    splits: HashMap<NodeId, ManifestSplits>,
     new_groups: HashMap<Path, (NodeId, Bytes)>,
     new_arrays: HashMap<Path, (NodeId, ArrayData)>,
     updated_arrays: HashMap<NodeId, ArrayData>,
-    updated_groups: HashMap<NodeId, Bytes>,
-    // It's important we keep these sorted, we use this fact in TransactionLog creation
-    // TODO: could track ManifestExtents
-    set_chunks: BTreeMap<NodeId, BTreeMap<ChunkIndices, Option<ChunkPayload>>>,
+    updated3_groups: HashMap<NodeId, Bytes>,
+    // FIXME: It's important we keep these sorted, we use this fact in TransactionLog creation
+    //        Change HashMap -> BTreeMap, need to check Ord on ManifestExtents
+    set_chunks: BTreeMap<NodeId, HashMap<ManifestExtents, SplitManifest>>,
     deleted_groups: HashSet<(Path, NodeId)>,
     deleted_arrays: HashSet<(Path, NodeId)>,
 }
 
 impl ChangeSet {
+    pub fn new(splitting: ManifestSplittingConfig) -> Self {
+        Self { splitting, ..Default::default() }
+    }
+
     pub fn deleted_arrays(&self) -> impl Iterator<Item = &(Path, NodeId)> {
         self.deleted_arrays.iter()
     }
@@ -58,10 +120,13 @@ impl ChangeSet {
         self.deleted_arrays.contains(path_and_id)
     }
 
+    pub fn splits(&self, id: &NodeId) -> Option<&ManifestSplits> {
+        self.splits.get(id)
+    }
+
     pub fn chunk_changes(
         &self,
-    ) -> impl Iterator<Item = (&NodeId, &BTreeMap<ChunkIndices, Option<ChunkPayload>>)>
-    {
+    ) -> impl Iterator<Item = (&NodeId, &HashMap<ManifestExtents, SplitManifest>)> {
         self.set_chunks.iter()
     }
 
@@ -100,11 +165,51 @@ impl ChangeSet {
         }
     }
 
-    pub fn add_array(&mut self, path: Path, node_id: NodeId, array_data: ArrayData) {
+    fn update_cached_splits(
+        &mut self,
+        node_id: &NodeId,
+        path: &Path,
+        shape: &ArrayShape,
+        dimension_names: &Option<Vec<DimensionName>>,
+    ) -> SessionResult<()> {
+        // FIXME: What happens if we set a chunk, then change a dimension name, so
+        //        that the split changes.
+        let splits = self.splitting.get_split_sizes(path, shape, dimension_names)?;
+        // FIXME: Check that splits matches any existing split used? This is for the case when
+        //        we write a chunk, populate `self.splits`, then say a dimension name was changed,
+        //        which changes the splits, so things are now inconsistent.
+        self.splits.insert(node_id.clone(), splits);
+        Ok(())
+    }
+
+    pub fn add_array(
+        &mut self,
+        path: Path,
+        node_id: NodeId,
+        array_data: ArrayData,
+    ) -> SessionResult<()> {
+        self.update_cached_splits(
+            &node_id,
+            &path,
+            &array_data.shape,
+            &array_data.dimension_names,
+        );
         self.new_arrays.insert(path, (node_id, array_data));
+        Ok(())
     }
 
-    pub fn update_array(&mut self, node_id: &NodeId, path: &Path, array_data: ArrayData) {
+    pub fn update_array(
+        &mut self,
+        node_id: &NodeId,
+        path: &Path,
+        array_data: ArrayData,
+    ) -> SessionResult<()> {
+        self.update_cached_splits(
+            &node_id,
+            &path,
+            &array_data.shape,
+            &array_data.dimension_names,
+        );
         match self.new_arrays.get(path) {
             Some((id, _)) => {
                 debug_assert!(!self.updated_arrays.contains_key(id));
@@ -114,6 +219,7 @@ impl ChangeSet {
                 self.updated_arrays.insert(node_id.clone(), array_data);
             }
         }
+        Ok(())
     }
 
     pub fn update_group(&mut self, node_id: &NodeId, path: &Path, definition: Bytes) {
@@ -139,6 +245,7 @@ impl ChangeSet {
 
         self.updated_arrays.remove(node_id);
         self.set_chunks.remove(node_id);
+        self.splits.remove(node_id);
         if !is_new_array {
             self.deleted_arrays.insert((path, node_id.clone()));
         }
@@ -166,34 +273,60 @@ impl ChangeSet {
         node_id: NodeId,
         coord: ChunkIndices,
         data: Option<ChunkPayload>,
-    ) {
+    ) -> SessionResult<()> {
+        let cached_splits = self.splits.get(&node_id).expect("logic bug");
+
+        let extent = cached_splits.which_extent(&coord)?;
         // this implementation makes delete idempotent
         // it allows deleting a deleted chunk by repeatedly setting None.
         self.set_chunks
             .entry(node_id)
             .and_modify(|h| {
-                h.insert(coord.clone(), data.clone());
+                h.entry(extent.clone()).or_default().update(coord.clone(), data.clone());
             })
-            .or_insert(BTreeMap::from([(coord, data)]));
+            .or_insert_with(|| {
+                let mut h = HashMap::<ManifestExtents, SplitManifest>::with_capacity(
+                    cached_splits.len(),
+                );
+                h.entry(extent.clone())
+                    // TODO: this is duplicative. I can't use `or_default` because it's
+                    // nice to create the HashMap using `with_capacity`
+                    .or_default()
+                    .update(coord, data);
+                h
+            });
+        Ok(())
     }
 
     pub fn get_chunk_ref(
         &self,
         node_id: &NodeId,
         coords: &ChunkIndices,
     ) -> Option<&Option<ChunkPayload>> {
-        self.set_chunks.get(node_id).and_then(|h| h.get(coords))
+        self.splits
+            .get(node_id)
+            .and_then(|splits| {
+                splits.which_extent(coords).ok().map(|extent| {
+                    self.set_chunks
+                        .get(node_id)
+                        .and_then(|h| h.get(&extent))
+                        .and_then(|s| s.chunks.get(coords))
+                })
+            })
+            .flatten()
     }
 
     /// Drop the updated chunk references for the node.
     /// This will only drop the references for which `predicate` returns true
     pub fn drop_chunk_changes(
         &mut self,
         node_id: &NodeId,
-        predicate: impl Fn(&ChunkIndices) -> bool,
+        predicate: impl Fn(&ChunkIndices) -> bool + Clone,
     ) {
         if let Some(changes) = self.set_chunks.get_mut(node_id) {
-            changes.retain(|coord, _| !predicate(coord));
+            for split in changes.values_mut() {
+                split.retain(predicate.clone());
+            }
         }
     }
 
@@ -207,7 +340,7 @@ impl ChangeSet {
         }
         match self.set_chunks.get(node_id) {
             None => Either::Left(iter::empty()),
-            Some(h) => Either::Right(h.iter()),
+            Some(h) => Either::Right(h.values().flat_map(|x| x.chunks.iter())),
         }
     }
 
@@ -235,6 +368,44 @@ impl ChangeSet {
         )
     }
 
+    pub fn array_manifests_iterator(
+        &self,
+        node_id: &NodeId,
+        node_path: &Path,
+    ) -> impl Iterator<Item = (&ManifestExtents, &SplitManifest)> + use<'_> {
+        if self.is_deleted(node_path, node_id) {
+            return Either::Left(iter::empty());
+        }
+        match self.set_chunks.get(node_id) {
+            None => Either::Left(iter::empty()),
+            Some(h) => Either::Right(h.iter()),
+        }
+    }
+
+    pub fn array_manifest(
+        &self,
+        node_id: &NodeId,
+        extent: &ManifestExtents,
+    ) -> Option<&SplitManifest> {
+        self.set_chunks.get(node_id).and_then(|x| x.get(extent))
+    }
+
+    // pub fn new_array_manifest_iterator<'a>(
+    //     &'a self,
+    //     node_id: &'a NodeId,
+    //     node_path: &Path,
+    // ) -> impl Iterator<Item = ChunkInfo> + use<'a> {
+    //     self.array_manifests_iterator(node_id, node_path).filter_map(
+    //         move |(coords, payload)| {
+    //             payload.as_ref().map(|p| ChunkInfo {
+    //                 node: node_id.clone(),
+    //                 coord: coords.clone(),
+    //                 payload: p.clone(),
+    //             })
+    //         },
+    //     )
+    // }
+
     pub fn new_nodes(&self) -> impl Iterator<Item = (&Path, &NodeId)> {
         self.new_groups().chain(self.new_arrays())
     }
@@ -274,6 +445,7 @@ impl ChangeSet {
         self.updated_arrays.extend(other.updated_arrays);
         self.deleted_groups.extend(other.deleted_groups);
         self.deleted_arrays.extend(other.deleted_arrays);
+        // FIXME: handle splits
 
         for (node, other_chunks) in other.set_chunks.into_iter() {
             match self.set_chunks.remove(&node) {

icechunk/src/format/manifest.rs

@@ -18,7 +18,7 @@ use super::{
     ChunkId, ChunkIndices, ChunkLength, ChunkOffset, IcechunkResult, ManifestId, NodeId,
 };
 
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub struct ManifestExtents(Vec<Range<u32>>);
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -54,8 +54,8 @@ impl ManifestExtents {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct ManifestSplits(Vec<ManifestExtents>);
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ManifestSplits(pub Vec<ManifestExtents>);
 
 impl ManifestSplits {
     /// Used at read-time