feat: keep column groups through compaction and add rewrite_columns

docs/src/guide/read_and_write.md

@@ -85,6 +85,26 @@ the target version, actual version, and file path. A persistent compaction targe
 can be set through `lance.compaction.data_storage_version` in the table config;
 an explicit operation target takes precedence.
 
+Compaction re-encodes every column of the fragments it rewrites. To migrate
+only some columns, `rewrite_columns` (Python: `LanceDataset.rewrite_columns`,
+Rust: `Dataset::rewrite_columns`) reads just the named top-level columns of each
+fragment, writes them to one new data file per fragment in the requested V2
+version, and tombstones them in the files they came from. The files holding the
+other columns are neither read nor written, and rows, fragment ids, row
+addresses and indices are unchanged. Fragments already in the requested layout
+are skipped, so an interrupted rewrite can be rerun; `LanceFragment.rewrite_columns`
+does the same for one fragment so the work can be spread over workers and
+committed as a single `Update` in `rewrite_columns` mode.
+
+A compaction folds those files back into one per fragment unless it is told to
+keep them apart: the `column_groups` compaction option (config key
+`lance.compaction.column_groups`, groups separated by `;` and columns by `,`)
+writes each listed group of columns to its own data file per fragment and the
+remaining columns to one shared file. Setting the config key once makes every
+later compaction preserve the layout that `rewrite_columns` produced. Binary
+copy is disabled when groups are set, and `max_bytes_per_file` is ignored so
+all groups split at the same rows.
+
 ### Upgrading clients before mixed-version writes
 
 Before writing files that differ from the dataset default, upgrade every reader

python/python/lance/dataset.py

@@ -2868,6 +2868,52 @@ def drop_columns(self, columns: List[str]):
         # Indices might have changed
         self._list_indices_res = None
 
+    def rewrite_columns(
+        self,
+        columns: List[str],
+        *,
+        data_storage_version: Optional[str] = None,
+    ):
+        """Rewrite columns into their own data files, leaving other files alone.
+
+        Every fragment gets one new data file holding exactly ``columns``, and
+        those columns are tombstoned in the files they came from. The files
+        holding the other columns are not read or written, so this migrates
+        narrow columns to a newer data file version without re-encoding the
+        wide columns that dominate a table's size. Rows, fragment ids, row
+        addresses and indices are unchanged.
+
+        Fragments whose ``columns`` already sit alone in a file of the requested
+        version are skipped, so an interrupted rewrite can be rerun. Fragments
+        are rewritten one after another; to spread the work over many workers,
+        call :meth:`lance.fragment.LanceFragment.rewrite_columns` per fragment
+        and commit the returned metadata in one
+        :class:`LanceOperation.Update` with ``update_mode="rewrite_columns"``.
+
+        A later ``compact_files`` folds the new files back into one file per
+        fragment unless its ``column_groups`` lists the same columns.
+
+        Parameters
+        ----------
+        columns : list of str
+            Top-level column names to rewrite.
+        data_storage_version : str, optional
+            Data file version for the new files, such as ``"2.2"`` or
+            ``"stable"``. Defaults to the dataset's default write version and
+            never changes it. Must be a V2 version.
+
+        Examples
+        --------
+        >>> import lance
+        >>> import pyarrow as pa
+        >>> table = pa.table({"a": [1, 2, 3], "b": ["x", "y", "z"]})
+        >>> dataset = lance.write_dataset(table, "example", data_storage_version="2.0")
+        >>> dataset.rewrite_columns(["b"], data_storage_version="2.2")
+        >>> [f.fields for f in dataset.get_fragments()[0].data_files()]
+        [[0, -2], [1]]
+        """
+        self._ds.rewrite_columns(columns, data_storage_version)
+
     def delete(
         self,
         predicate: Union[str, Expression],
@@ -7511,6 +7557,7 @@ def compact_files(
         max_source_bytes: Optional[int] = None,
         excluded_fragment_ids: Optional[list[int]] = None,
         data_storage_version: Optional[str] = None,
+        column_groups: Optional[list[list[str]]] = None,
     ) -> CompactionMetrics:
         """Compacts small files in the dataset, reducing total number of files.
 
@@ -7545,7 +7592,9 @@ def compact_files(
         ``lance.compaction.max_source_fragments``,
         ``lance.compaction.max_source_rows``,
         ``lance.compaction.max_source_bytes``,
-        ``lance.compaction.data_storage_version``.
+        ``lance.compaction.data_storage_version``,
+        ``lance.compaction.column_groups`` (groups separated by ``;``,
+        columns by ``,``, e.g. ``"embedding;caption,tags"``).
 
         Parameters
         ----------
@@ -7626,6 +7675,16 @@ def compact_files(
             Uses the compaction config target when set, otherwise the dataset's
             default write version. Does not change that default or the versions
             of unselected files. V1/V2 cross-family targets are rejected.
+        column_groups: list[list[str]], optional
+            Top-level columns to keep in their own data files. Each inner list
+            becomes one data file per compacted fragment holding exactly those
+            columns; every column not listed goes to one shared file. This is
+            the layout :meth:`LanceDataset.rewrite_columns` produces, so a
+            compaction configured with the same groups preserves it instead of
+            folding wide columns back next to narrow ones. Binary copy is
+            disabled when set, and ``max_bytes_per_file`` is ignored so every
+            group splits at the same rows. Uses the manifest config value when
+            not specified.
 
         Returns
         -------
@@ -7654,6 +7713,7 @@ def compact_files(
                 max_source_bytes=max_source_bytes,
                 excluded_fragment_ids=excluded_fragment_ids,
                 data_storage_version=data_storage_version,
+                column_groups=column_groups,
             ).items()
             if v is not None
         }

python/python/lance/fragment.py

@@ -1015,6 +1015,30 @@ def update_columns(
             return metadata, fields_modified, matched_offsets
         return metadata, fields_modified
 
+    def rewrite_columns(
+        self,
+        columns: List[str],
+        *,
+        data_storage_version: Optional[str] = None,
+    ) -> Optional[FragmentMetadata]:
+        """Rewrite columns of this fragment into one new data file.
+
+        .. warning::
+
+            Internal API. This method is not intended to be used by end users.
+
+        The per-fragment half of
+        :meth:`lance.dataset.LanceDataset.rewrite_columns`, for spreading a
+        rewrite over many workers. The new file is written but nothing is
+        committed: collect the returned metadata from every fragment and commit
+        it in one :class:`lance.dataset.LanceOperation.Update` with
+        ``update_mode="rewrite_columns"`` and no ``fields_modified``.
+
+        Returns ``None`` when ``columns`` already sit alone in a file of the
+        requested version.
+        """
+        return self._fragment.rewrite_columns(columns, data_storage_version)
+
     def merge_columns(
         self,
         value_func: (

python/python/lance/lance/__init__.pyi

@@ -636,6 +636,9 @@ class _Dataset:
     def validate(self): ...
     def migrate_manifest_paths_v2(self): ...
     def drop_columns(self, columns: List[str]): ...
+    def rewrite_columns(
+        self, columns: List[str], data_storage_version: Optional[str] = None
+    ): ...
     def add_columns_from_reader(
         self, reader: pa.RecordBatchReader, batch_size: Optional[int] = None
     ): ...
@@ -759,6 +762,9 @@ class _Fragment:
         read_columns: Optional[List[str]],
         batch_size: Optional[int],
     ) -> Tuple[FragmentMetadata, LanceSchema]: ...
+    def rewrite_columns(
+        self, columns: List[str], data_storage_version: Optional[str] = None
+    ) -> Optional[FragmentMetadata]: ...
     def delete(self, predicate: str) -> Optional[_Fragment]: ...
     def delete_rows(self, offsets: List[int]) -> Optional[_Fragment]: ...
     def schema(self) -> pa.Schema: ...

python/python/lance/optimize.py

@@ -131,3 +131,13 @@ class CompactionOptions(TypedDict, total=False):
     input versions and no overlays; TryBinaryCopy reencodes ineligible inputs,
     while ForceBinaryCopy reports an error.
     """
+    column_groups: Optional[list[list[str]]]
+    """
+    Top-level columns to keep in their own data files. Each inner list becomes
+    one data file per compacted fragment holding exactly those columns; every
+    column not listed goes to one shared file. This is the layout
+    ``LanceDataset.rewrite_columns`` produces, so a compaction configured with
+    the same groups preserves it. Binary copy is disabled when set, and
+    ``max_bytes_per_file`` is ignored so every group splits at the same rows.
+    (default: None, one file per fragment)
+    """

python/python/tests/test_optimize.py

@@ -909,3 +909,43 @@ def test_remap_row_addrs(tmp_path: Path):
         pa.array([old[i] for i in sample], pa.uint64())
     ).to_pylist()
     assert remapped == [new[i] for i in sample]
+
+
+def test_compact_files_column_groups(tmp_path: Path):
+    data = pa.table({"a": range(8), "b": [str(i) for i in range(8)], "c": range(8)})
+    dataset = lance.write_dataset(
+        data, tmp_path / "dataset", max_rows_per_file=2, data_storage_version="2.0"
+    )
+    dataset.delete("a = 5")
+    expected = dataset.to_table()
+
+    with pytest.raises(OSError, match="more than once"):
+        dataset.optimize.compact_files(column_groups=[["c"], ["c"]])
+    with pytest.raises(OSError, match="not a top-level column"):
+        dataset.optimize.compact_files(column_groups=[["missing"]])
+    with pytest.raises(OSError, match="binary copy is not supported"):
+        dataset.optimize.compact_files(
+            column_groups=[["c"]], compaction_mode="force_binary_copy"
+        )
+
+    metrics = dataset.optimize.compact_files(
+        target_rows_per_fragment=100,
+        column_groups=[["c"]],
+        data_storage_version="2.2",
+    )
+    assert metrics.fragments_added == 1
+    assert metrics.files_added == 2
+    (fragment,) = dataset.get_fragments()
+    files = fragment.data_files()
+    assert [f.fields for f in files] == [[0, 1], [2]]
+    assert all((f.file_major_version, f.file_minor_version) == (2, 2) for f in files)
+    assert dataset.to_table() == expected
+    assert dataset.data_storage_version == "2.0"
+
+    # The persisted config makes later compactions keep the layout.
+    dataset.update_config({"lance.compaction.column_groups": "c"})
+    dataset.insert(pa.table({"a": [8], "b": ["8"], "c": [8]}))
+    dataset.optimize.compact_files(target_rows_per_fragment=100)
+    (fragment,) = dataset.get_fragments()
+    assert [f.fields for f in fragment.data_files()] == [[0, 1], [2]]
+    assert dataset.count_rows() == 8

python/python/tests/test_schema_evolution.py

@@ -627,3 +627,78 @@ def test_project_nullability_assertion_round_trips(tmp_path: Path):
     )
     with pytest.raises(Exception, match="preempted"):
         lance.LanceDataset.commit(tmp_path, relax, read_version=written_at)
+
+
+def test_rewrite_columns(tmp_path: Path):
+    table = pa.table({"a": range(8), "b": [str(i) for i in range(8)], "c": range(8)})
+    dataset = lance.write_dataset(
+        table, tmp_path, max_rows_per_file=4, data_storage_version="2.0"
+    )
+    dataset.delete("a = 3")
+    dataset.create_scalar_index("a", "BTREE")
+    expected = dataset.to_table()
+    version = dataset.version
+
+    dataset.rewrite_columns(["c"], data_storage_version="2.2")
+
+    assert dataset.version == version + 1
+    for fragment in dataset.get_fragments():
+        files = fragment.data_files()
+        assert [f.fields for f in files] == [[0, 1, -2], [2]]
+        assert (files[0].file_major_version, files[0].file_minor_version) == (2, 0)
+        assert (files[1].file_major_version, files[1].file_minor_version) == (2, 2)
+    assert dataset.to_table() == expected
+    assert dataset.data_storage_version == "2.0"
+    # The values did not change, so the index still covers every fragment.
+    (index,) = dataset.describe_indices()
+    (segment,) = index.segments
+    assert set(segment.fragment_ids) == {0, 1}
+    assert dataset.to_table(filter="a = 5").num_rows == 1
+
+    # Already in the requested layout: no new version.
+    dataset.rewrite_columns(["c"], data_storage_version="2.2")
+    assert dataset.version == version + 1
+
+    with pytest.raises(ValueError, match="not a top-level column"):
+        dataset.rewrite_columns(["missing"])
+
+    # A compaction that knows the group keeps `c` apart while moving the rest.
+    dataset.optimize.compact_files(
+        target_rows_per_fragment=100,
+        column_groups=[["c"]],
+        data_storage_version="2.2",
+    )
+    (fragment,) = dataset.get_fragments()
+    files = fragment.data_files()
+    assert [f.fields for f in files] == [[0, 1], [2]]
+    assert all((f.file_major_version, f.file_minor_version) == (2, 2) for f in files)
+    assert dataset.to_table() == expected
+
+
+def test_rewrite_columns_per_fragment_commit(tmp_path: Path):
+    table = pa.table({"a": range(6), "b": [str(i) for i in range(6)]})
+    dataset = lance.write_dataset(
+        table, tmp_path, max_rows_per_file=3, data_storage_version="2.0"
+    )
+    expected = dataset.to_table()
+
+    # The distributed shape: rewrite each fragment on its own, commit once.
+    updated = [
+        fragment.rewrite_columns(["b"], data_storage_version="2.1")
+        for fragment in dataset.get_fragments()
+    ]
+    assert all(metadata is not None for metadata in updated)
+    operation = lance.LanceOperation.Update(
+        updated_fragments=updated, update_mode="rewrite_columns"
+    )
+    dataset = lance.LanceDataset.commit(
+        dataset, operation, read_version=dataset.version
+    )
+
+    for fragment in dataset.get_fragments():
+        files = fragment.data_files()
+        assert [f.fields for f in files] == [[0, -2], [1]]
+        assert (files[1].file_major_version, files[1].file_minor_version) == (2, 1)
+        # Nothing left to do for this fragment.
+        assert fragment.rewrite_columns(["b"], data_storage_version="2.1") is None
+    assert dataset.to_table() == expected

python/src/dataset.rs

@@ -3272,6 +3272,28 @@ impl Dataset {
         Ok(())
     }
 
+    #[pyo3(signature = (columns, data_storage_version = None))]
+    fn rewrite_columns(
+        &mut self,
+        columns: Vec<String>,
+        data_storage_version: Option<&str>,
+    ) -> PyResult<()> {
+        let version = data_storage_version
+            .map(str::parse)
+            .transpose()
+            .infer_error()?;
+        let mut new_self = self.ds.as_ref().clone();
+        let new_self = rt()
+            .spawn(None, async move {
+                let columns: Vec<&str> = columns.iter().map(String::as_str).collect();
+                new_self.rewrite_columns(&columns, version).await?;
+                Ok(new_self)
+            })?
+            .infer_error()?;
+        self.ds = Arc::new(new_self);
+        Ok(())
+    }
+
     #[pyo3(signature = (reader, batch_size = None))]
     fn add_columns_from_reader(
         &mut self,

python/src/dataset/optimize.rs

@@ -92,6 +92,11 @@ fn parse_compaction_options(
                     opts.data_storage_version = Some(version.parse().infer_error()?);
                 }
             }
+            "column_groups" => {
+                opts.column_groups = value
+                    .extract::<Option<Vec<Vec<String>>>>()?
+                    .unwrap_or_default();
+            }
             _ => {
                 return Err(PyValueError::new_err(format!(
                     "Invalid compaction option: {}",
@@ -129,8 +134,8 @@ pub struct PyCompactionMetrics {
     /// int : The number of files that have been removed, including deletion files.
     #[pyo3(get)]
     pub files_removed: usize,
-    /// int : The number of files that have been added, which is always equal to the
-    /// number of fragments.
+    /// int : The number of files that have been added: one per new fragment, or
+    /// one per column group and fragment when ``column_groups`` is set.
     #[pyo3(get)]
     pub files_added: usize,
 }

python/src/fragment.rs

@@ -380,6 +380,26 @@ impl FileFragment {
         Ok((PyLance(fragment), LanceSchema(schema)))
     }
 
+    #[pyo3(signature = (columns, data_storage_version = None))]
+    fn rewrite_columns(
+        &self,
+        columns: Vec<String>,
+        data_storage_version: Option<&str>,
+    ) -> PyResult<Option<PyLance<Fragment>>> {
+        let version = data_storage_version
+            .map(str::parse)
+            .transpose()
+            .infer_error()?;
+        let fragment = self.fragment.clone();
+        let updated = rt()
+            .spawn(None, async move {
+                let columns: Vec<&str> = columns.iter().map(String::as_str).collect();
+                fragment.rewrite_columns(&columns, version).await
+            })?
+            .infer_error()?;
+        Ok(updated.map(PyLance))
+    }
+
     fn merge(
         &mut self,
         reader: PyArrowType<ArrowArrayStreamReader>,

rust/lance/src/dataset.rs

@@ -82,6 +82,7 @@ pub mod optimize;
 pub(crate) mod overlay;
 pub mod progress;
 pub mod refs;
+pub mod rewrite_columns;
 pub mod rowids;
 pub mod scanner;
 mod schema_evolution;