From e0499de959da158b8de91eb0cca98c104e4c2244 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:01:42 +0000 Subject: [PATCH] fix(python): expose optimize compaction options --- docs/src/python/python.md | 2 + python/python/lancedb/__init__.py | 3 +- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/remote/table.py | 12 +++- python/python/lancedb/table.py | 81 +++++++++++++++++++++++++++ python/python/tests/test_table.py | 23 ++++++++ python/src/table.rs | 67 +++++++++++++++++++--- rust/lancedb/src/table.rs | 4 +- rust/lancedb/src/table/optimize.rs | 2 +- 9 files changed, 184 insertions(+), 11 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 36044d35d..3b9b50cec 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -42,6 +42,8 @@ listing a storage directory. ::: lancedb.table.Table +::: lancedb.table.CompactionOptions + ::: lancedb.table.FragmentStatistics ::: lancedb.table.FragmentSummaryStats diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 235049f97..460d0695b 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -21,7 +21,7 @@ from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector, BlobType from .job import AsyncJob, Job -from .table import AsyncTable, Table +from .table import AsyncTable, CompactionOptions, Table from .types import BaseTokenizerType from ._lancedb import Session from .namespace import ( @@ -504,6 +504,7 @@ __all__ = [ "AsyncJob", "AsyncLanceNamespaceDBConnection", "AsyncTable", + "CompactionOptions", "FtsToken", "col", "Expr", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 47e727f99..983199acf 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -347,6 +347,7 @@ class Table: *, cleanup_since_ms: Optional[int] = None, delete_unverified: Optional[bool] = None, + compaction_options: Optional[Dict[str, Any]] = None, ) -> OptimizeStats: ... async def uri(self) -> str: ... async def initial_storage_options(self) -> Optional[Dict[str, str]]: ... diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index acc2f4c9d..71bebc78f 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -64,7 +64,16 @@ from ..query import ( LanceTakeQueryBuilder, LanceVectorQueryBuilder, ) -from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags +from ..table import ( + AsyncTable, + BlobMode, + Branches, + CompactionOptions, + IndexStatistics, + Query, + Table, + Tags, +) from ..types import BaseTokenizerType @@ -942,6 +951,7 @@ class RemoteTable(Table): *, cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, + compaction_options: Optional[CompactionOptions] = None, ): """ optimize() is a no-op on LanceDB Cloud. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 31a70c298..2b8fa0a24 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -22,6 +22,7 @@ from typing import ( Optional, Sequence, Tuple, + TypedDict, Union, overload, ) @@ -211,6 +212,69 @@ IndexConfigType = Union[ FTS, ] + +class CompactionOptions(TypedDict, total=False): + """Options that control file compaction during table optimization. + + Unspecified options use Lance's defaults. + + Examples + -------- + Limit fragment size and compaction concurrency for tables with large rows: + + >>> options: CompactionOptions = { + ... "target_rows_per_fragment": 500, + ... "batch_size": 1024, + ... "num_threads": 1, + ... } + >>> await table.optimize(compaction_options=options) # doctest: +SKIP + """ + + target_rows_per_fragment: int + """Target number of rows per fragment (default: 1,048,576).""" + + max_rows_per_group: int + """Maximum number of rows per row group (default: 1,024).""" + + max_bytes_per_file: Optional[int] + """Maximum number of bytes per data file.""" + + materialize_deletions: bool + """Whether to rewrite fragments containing deleted rows (default: True).""" + + materialize_deletions_threshold: float + """Minimum deleted-row fraction that makes a fragment eligible (default: 0.1).""" + + num_threads: Optional[int] + """Number of compaction tasks to run in parallel.""" + + batch_size: Optional[int] + """Number of rows per input scan batch.""" + + io_buffer_size: Optional[int] + """Maximum number of bytes queued in the input scan I/O buffer.""" + + defer_index_remap: bool + """Whether to defer index remapping during compaction (default: False).""" + + index_remap_mode: Literal["direct", "compact"] + """How to construct the old-to-new row-address mapping.""" + + compaction_mode: Optional[ + Literal["reencode", "try_binary_copy", "force_binary_copy"] + ] + """Whether compaction re-encodes data or uses binary copying.""" + + binary_copy_read_batch_bytes: Optional[int] + """Number of bytes read per batch during binary-copy compaction.""" + + max_source_fragments: Optional[int] + """Maximum number of source fragments compacted in one run.""" + + max_overlays_per_fragment: Optional[int] + """Maximum overlays before a fragment is fully compacted.""" + + # Known distance metrics for legacy API detection KNOWN_METRICS = {"l2", "cosine", "dot", "hamming"} @@ -1820,6 +1884,7 @@ class Table(ABC): cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, retrain: bool = False, + compaction_options: Optional[CompactionOptions] = None, ): """ Optimize the on-disk data and indices for better performance. @@ -1852,6 +1917,10 @@ class Table(ABC): retrain: bool, default False This parameter is no longer used and is deprecated. + compaction_options: CompactionOptions, optional + Options that control file compaction. This can be used to bound + memory usage by reducing ``target_rows_per_fragment``, ``batch_size``, + or ``num_threads``. Notes ----- @@ -3825,6 +3894,7 @@ class LanceTable(Table): cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, retrain: bool = False, + compaction_options: Optional[CompactionOptions] = None, ): """ Optimize the on-disk data and indices for better performance. @@ -3857,6 +3927,10 @@ class LanceTable(Table): retrain: bool, default False This parameter is no longer used and is deprecated. + compaction_options: CompactionOptions, optional + Options that control file compaction. This can be used to bound + memory usage by reducing ``target_rows_per_fragment``, ``batch_size``, + or ``num_threads``. Notes ----- @@ -3871,6 +3945,7 @@ class LanceTable(Table): cleanup_older_than=cleanup_older_than, delete_unverified=delete_unverified, retrain=retrain, + compaction_options=compaction_options, ) ) @@ -6019,6 +6094,7 @@ class AsyncTable: cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, retrain=False, + compaction_options: Optional[CompactionOptions] = None, ) -> OptimizeStats: """ Optimize the on-disk data and indices for better performance. @@ -6051,6 +6127,10 @@ class AsyncTable: retrain: bool, default False This parameter is no longer used and is deprecated. + compaction_options: CompactionOptions, optional + Options that control file compaction. This can be used to bound + memory usage by reducing ``target_rows_per_fragment``, ``batch_size``, + or ``num_threads``. Notes ----- @@ -6077,6 +6157,7 @@ class AsyncTable: return await self._inner.optimize( cleanup_since_ms=cleanup_since_ms, delete_unverified=delete_unverified, + compaction_options=compaction_options, ) async def list_indices(self) -> Iterable[IndexConfig]: diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 542c944c8..5fb5ce9c9 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3381,6 +3381,29 @@ async def test_optimize(mem_db_async: AsyncConnection): assert await table.query().to_arrow() == pa.table({"x": [[1], [2]]}) +@pytest.mark.asyncio +async def test_optimize_compaction_options(mem_db_async: AsyncConnection): + table = await mem_db_async.create_table("test", data=[{"x": 1}]) + await table.add([{"x": 2}]) + + stats = await table.optimize( + compaction_options={ + "target_rows_per_fragment": 1, + "batch_size": 1, + "num_threads": 1, + } + ) + assert stats.compaction.fragments_removed == 0 + assert stats.compaction.fragments_added == 0 + + stats = await table.optimize(compaction_options={"target_rows_per_fragment": 3}) + assert stats.compaction.fragments_removed == 2 + assert stats.compaction.fragments_added == 1 + + with pytest.raises(ValueError, match="Invalid compaction option: unknown"): + await table.optimize(compaction_options={"unknown": 1}) + + @pytest.mark.asyncio async def test_optimize_delete_unverified(tmp_db_async: AsyncConnection, tmp_path): table = await tmp_db_async.create_table( diff --git a/python/src/table.rs b/python/src/table.rs index f20bcf2bf..963183dea 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -20,8 +20,9 @@ use arrow::{ use lancedb::blob::{BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ - AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, - NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, + AddDataMode, ColumnAlteration, CompactionMode, CompactionOptions, Duration, + FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform, + OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, }; use lancedb::tokenize as lancedb_tokenize; use pyo3::{ @@ -39,6 +40,56 @@ enum PredicateArg { Sql(String), } +fn parse_compaction_options(options: Option<&Bound<'_, PyDict>>) -> PyResult { + let mut parsed = CompactionOptions::default(); + let Some(options) = options else { + return Ok(parsed); + }; + + for (key, value) in options.iter() { + let key: String = key.extract()?; + match key.as_str() { + "target_rows_per_fragment" => parsed.target_rows_per_fragment = value.extract()?, + "max_rows_per_group" => parsed.max_rows_per_group = value.extract()?, + "max_bytes_per_file" => parsed.max_bytes_per_file = value.extract()?, + "materialize_deletions" => parsed.materialize_deletions = value.extract()?, + "materialize_deletions_threshold" => { + parsed.materialize_deletions_threshold = value.extract()? + } + "num_threads" => parsed.num_threads = value.extract()?, + "batch_size" => parsed.batch_size = value.extract()?, + "io_buffer_size" => parsed.io_buffer_size = value.extract()?, + "defer_index_remap" => parsed.defer_index_remap = value.extract()?, + "index_remap_mode" => { + let mode: String = value.extract()?; + parsed.index_remap_mode = IndexRemapMode::try_from(mode.as_str()) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + } + "compaction_mode" => { + let mode: Option = value.extract()?; + parsed.compaction_mode = mode + .map(|mode| { + CompactionMode::try_from(mode.as_str()) + .map_err(|err| PyValueError::new_err(err.to_string())) + }) + .transpose()?; + } + "binary_copy_read_batch_bytes" => { + parsed.binary_copy_read_batch_bytes = value.extract()? + } + "max_source_fragments" => parsed.max_source_fragments = value.extract()?, + "max_overlays_per_fragment" => parsed.max_overlays_per_fragment = value.extract()?, + _ => { + return Err(PyValueError::new_err(format!( + "Invalid compaction option: {key}" + ))); + } + } + } + + Ok(parsed) +} + /// Statistics about a compaction operation. #[pyclass(get_all, from_py_object)] #[derive(Clone, Debug)] @@ -1195,13 +1246,15 @@ impl Table { } /// Optimize the on-disk data by compacting and pruning old data, for better performance. - #[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))] - pub fn optimize( - self_: PyRef<'_, Self>, + #[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None, compaction_options=None))] + pub fn optimize<'py>( + self_: PyRef<'py, Self>, cleanup_since_ms: Option, delete_unverified: Option, - ) -> PyResult> { + compaction_options: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { let inner = self_.inner_ref()?.clone(); + let compaction_options = parse_compaction_options(compaction_options)?; let older_than = if let Some(ms) = cleanup_since_ms { if ms > i64::MAX as u64 { return Err(PyValueError::new_err(format!( @@ -1217,7 +1270,7 @@ impl Table { future_into_py(self_.py(), async move { let compaction_stats = inner .optimize(OptimizeAction::Compact { - options: lancedb::table::CompactionOptions::default(), + options: compaction_options, remap_options: None, }) .await diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 0d8a8e8b9..65677f506 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -95,7 +95,9 @@ pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTa pub use lance::dataset::scanner::DatasetRecordBatchStream; use lance::dataset::statistics::DatasetStatisticsExt; pub use lance_index::optimize::OptimizeOptions; -pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats}; +pub use optimize::{ + CompactionMode, CompactionOptions, IndexRemapMode, OptimizeAction, OptimizeStats, +}; pub use schema_evolution::{ AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate, UpdateFieldMetadataResult, diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index e29445b2d..9599164b4 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -15,7 +15,7 @@ use lance_index::optimize::OptimizeOptions; use log::info; pub use chrono::Duration; -pub use lance::dataset::optimize::CompactionOptions; +pub use lance::dataset::optimize::{CompactionMode, CompactionOptions, IndexRemapMode}; use super::NativeTable; use crate::error::Result;