fix: updating schema/field metadata now retains fragments (#3384)

fixes: #3383 --------- Co-authored-by: Will Jones <[email protected]>
lancedb · Jan 15, 2025 · 9f7e012 · 9f7e012
1 parent 7cc43a7
commit 9f7e012
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 2 deletions.
diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs
@@ -3563,6 +3563,68 @@ mod tests {
         assert_eq!(dataset.manifest.config, desired_config);
     }
 
+    #[rstest]
+    #[tokio::test]
+    async fn test_replace_schema_metadata_preserves_fragments() {
+        let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "i",
+            DataType::UInt32,
+            false,
+        )]));
+
+        let data = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(UInt32Array::from_iter_values(0..100))],
+        );
+
+        let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema);
+        let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap();
+
+        let manifest_before = dataset.manifest.clone();
+
+        let mut new_schema_meta = HashMap::new();
+        new_schema_meta.insert("new_key".to_string(), "new_value".to_string());
+        dataset
+            .replace_schema_metadata(new_schema_meta.clone())
+            .await
+            .unwrap();
+
+        let manifest_after = dataset.manifest.clone();
+
+        assert_eq!(manifest_before.fragments, manifest_after.fragments);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn test_replace_fragment_metadata_preserves_fragments() {
+        let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "i",
+            DataType::UInt32,
+            false,
+        )]));
+
+        let data = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(UInt32Array::from_iter_values(0..100))],
+        );
+
+        let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema);
+        let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap();
+
+        let manifest_before = dataset.manifest.clone();
+
+        let mut new_field_meta = HashMap::new();
+        new_field_meta.insert("new_key".to_string(), "new_value".to_string());
+        dataset
+            .replace_field_metadata(vec![(0, new_field_meta.clone())])
+            .await
+            .unwrap();
+
+        let manifest_after = dataset.manifest.clone();
+
+        assert_eq!(manifest_before.fragments, manifest_after.fragments);
+    }
+
     #[rstest]
     #[tokio::test]
     async fn test_tag(

diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs
@@ -710,7 +710,7 @@ impl Transaction {
                 });
                 final_indices.extend(new_indices.clone());
             }
-            Operation::ReserveFragments { .. } => {
+            Operation::ReserveFragments { .. } | Operation::UpdateConfig { .. } => {
                 final_fragments.extend(maybe_existing_fragments?.clone());
             }
             Operation::Merge { ref fragments, .. } => {
@@ -744,7 +744,6 @@ impl Transaction {
             Operation::Restore { .. } => {
                 unreachable!()
             }
-            Operation::UpdateConfig { .. } => {}
         };
 
         // If a fragment was reserved then it may not belong at the end of the fragments list.