refactor: don't use schema dictionary unless using legacy storage (#4023)

Closes #4010

Author: westonpace

python/python/tests/test_dataset.py

@@ -2294,6 +2294,80 @@ def test_scan_with_batch_size(tmp_path: Path):
         assert batch.num_rows != 12
 
 
+def test_dictionaries(tmp_path: Path):
+    data = pa.table(
+        {
+            "id": pa.array([1, 2, 3]),
+            "dict": pa.array(
+                ["foo", "bar", "baz"], pa.dictionary(pa.int32(), pa.string())
+            ),
+        }
+    )
+    ds = lance.write_dataset(data, tmp_path)
+    assert ds.schema == pa.schema(
+        {"id": pa.int64(), "dict": pa.dictionary(pa.int32(), pa.string())}
+    )
+    assert ds.to_table() == data
+
+    # Can insert data with new values
+    new_data = pa.table(
+        {
+            "id": [4, 5, 6],
+            "dict": pa.array(
+                ["qux", "quux", "corge"], pa.dictionary(pa.int32(), pa.string())
+            ),
+        }
+    )
+    ds.insert(new_data)
+    table = ds.to_table().combine_chunks()
+    assert table == pa.table(
+        {
+            "id": [1, 2, 3, 4, 5, 6],
+            "dict": pa.array(
+                ["foo", "bar", "baz", "qux", "quux", "corge"],
+                pa.dictionary(pa.int32(), pa.string()),
+            ),
+        }
+    )
+
+    dict_arr = table.column("dict").chunk(0)
+    assert dict_arr.type == pa.dictionary(pa.int32(), pa.string())
+    assert dict_arr.to_pylist() == ["foo", "bar", "baz", "qux", "quux", "corge"]
+
+    assert dict_arr.dictionary.to_pylist() == [
+        "foo",
+        "bar",
+        "baz",
+        "qux",
+        "quux",
+        "corge",
+    ]
+
+    # Can merge insert data that has even more values
+    new_data = pa.table(
+        {
+            "id": [1, 7],
+            "dict": pa.array(
+                ["grault", "garply"], pa.dictionary(pa.int32(), pa.string())
+            ),
+        }
+    )
+    ds.merge_insert(
+        "id"
+    ).when_matched_update_all().when_not_matched_insert_all().execute(new_data)
+    table = ds.to_table().combine_chunks().sort_by("id")
+    assert table.column("id").to_pylist() == [1, 2, 3, 4, 5, 6, 7]
+    assert table.column("dict").to_pylist() == [
+        "grault",
+        "bar",
+        "baz",
+        "qux",
+        "quux",
+        "corge",
+        "garply",
+    ]
+
+
 @pytest.mark.slow
 def test_io_buffer_size(tmp_path: Path):
     # These cases regress deadlock issues that happen when the

rust/lance-table/src/format/manifest.rs

@@ -632,7 +632,9 @@ impl SelfDescribingFileReader for FileReader {
             location: location!(),
         })?;
         let mut manifest: Manifest = read_struct(reader.as_ref(), manifest_position).await?;
-        populate_schema_dictionary(&mut manifest.schema, reader.as_ref()).await?;
+        if manifest.should_use_legacy_format() {
+            populate_schema_dictionary(&mut manifest.schema, reader.as_ref()).await?;
+        }
         let schema = manifest.schema;
         let max_field_id = schema.max_field_id().unwrap_or_default();
         Self::try_new_from_reader(

rust/lance-table/src/io/manifest.rs

@@ -162,9 +162,10 @@ pub async fn write_manifest(
 ) -> Result<usize> {
     // Write dictionary values.
     let max_field_id = manifest.schema.max_field_id().unwrap_or(-1);
+    let is_legacy_storage = manifest.should_use_legacy_format();
     for field_id in 0..max_field_id + 1 {
         if let Some(field) = manifest.schema.mut_field_by_id(field_id) {
-            if field.data_type().is_dictionary() {
+            if field.data_type().is_dictionary() && is_legacy_storage {
                 let dict_info = field.dictionary.as_mut().ok_or_else(|| {
                     Error::io(
                         format!("Lance field {} misses dictionary info", field.name),

rust/lance/src/dataset.rs

@@ -483,7 +483,9 @@ impl Dataset {
             }
         }
 
-        populate_schema_dictionary(&mut manifest.schema, object_reader.as_ref()).await?;
+        if manifest.should_use_legacy_format() {
+            populate_schema_dictionary(&mut manifest.schema, object_reader.as_ref()).await?;
+        }
 
         Ok(manifest)
     }
@@ -629,7 +631,7 @@ impl Dataset {
             return Ok((self.manifest.clone(), self.manifest_location.clone()));
         }
         let mut manifest = read_manifest(&self.object_store, &location.path, location.size).await?;
-        if manifest.schema.has_dictionary_types() {
+        if manifest.schema.has_dictionary_types() && manifest.should_use_legacy_format() {
             let reader = if let Some(size) = location.size {
                 self.object_store
                     .open_with_size(&location.path, size as usize)
@@ -2651,8 +2653,12 @@ mod tests {
         assert!(matches!(result, Err(Error::SchemaMismatch { .. })))
     }
 
+    #[rstest]
     #[tokio::test]
-    async fn append_dictionary() {
+    async fn append_dictionary(
+        #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)]
+        data_storage_version: LanceFileVersion,
+    ) {
         // We store the dictionary as part of the schema, so we check that the
         // dictionary is consistent between appends.
 
@@ -2676,6 +2682,7 @@ mod tests {
         let mut write_params = WriteParams {
             max_rows_per_file: 40,
             max_rows_per_group: 10,
+            data_storage_version: Some(data_storage_version),
             ..Default::default()
         };
         let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone());
@@ -2711,10 +2718,14 @@ mod tests {
         )
         .unwrap()];
 
-        // Try write to dataset (fail)
+        // Try write to dataset (fails with legacy format)
         let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone());
         let result = Dataset::write(batches, test_uri, Some(write_params)).await;
-        assert!(result.is_err());
+        if data_storage_version == LanceFileVersion::Legacy {
+            assert!(result.is_err());
+        } else {
+            assert!(result.is_ok());
+        }
     }
 
     #[rstest]

rust/lance/src/dataset/builder.rs

@@ -308,7 +308,7 @@ impl DatasetBuilder {
             let location = commit_handler
                 .resolve_version_location(&base_path, manifest.version, &object_store.inner)
                 .await?;
-            if manifest.schema.has_dictionary_types() {
+            if manifest.schema.has_dictionary_types() && manifest.should_use_legacy_format() {
                 let reader = object_store.open(&location.path).await?;
                 populate_schema_dictionary(&mut manifest.schema, reader.as_ref()).await?;
             }

rust/lance/src/dataset/write.rs

@@ -369,7 +369,7 @@ pub async fn write_fragments_internal(
                         compare_nullability: NullabilityComparison::Ignore,
                         allow_missing_if_nullable: true,
                         ignore_field_order: true,
-                        compare_dictionary: true,
+                        compare_dictionary: dataset.is_legacy_storage(),
                         ..Default::default()
                     },
                 )?;

rust/lance/src/dataset/write/insert.rs

@@ -307,7 +307,12 @@ impl<'a> InsertBuilder<'a> {
                 }
                 let m = dataset.manifest.as_ref();
                 let mut schema_cmp_opts = SchemaCompareOptions {
-                    compare_dictionary: true,
+                    // In the legacy format we stored the dictionary in the manifest and
+                    // all files must have identical dictionaries.
+                    //
+                    // In 2.0+ the dictionary is stored in the files and dictionaries may
+                    // fluctuate between files.
+                    compare_dictionary: m.should_use_legacy_format(),
                     // array nullability is checked later, using actual data instead
                     // of the schema
                     compare_nullability: NullabilityComparison::Ignore,

rust/lance/src/dataset/write/merge_insert.rs

@@ -403,7 +403,7 @@ impl MergeInsertJob {
         let is_compatible = lance_schema.check_compatible(
             self.dataset.schema(),
             &SchemaCompareOptions {
-                compare_dictionary: true,
+                compare_dictionary: self.dataset.is_legacy_storage(),
                 ..Default::default()
             },
         );