diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f7aeac21a..fa854317d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -106,11 +106,15 @@ at the same path when the published site is assembled. Four things to know before editing: **Code examples are not written into the pages.** They live in real tests under -`docs/web-tests/{py,ts,rs}`, are extracted into `docs/web/snippets/`, and are -imported by the pages. Edit the test, then regenerate from the repository root: +`docs/web-tests/{py,ts,rs}`, `docs/web-tests/blob-v2/{py,ts,rs}`, and +`rust/lancedb/tests`, are extracted into `docs/web/snippets/`, and are imported +by the pages. Edit the test, then regenerate from the repository root: ```bash -uv run docs/web-tests/mdx_snippets_gen.py -s docs/web-tests/py -s docs/web-tests/ts -s docs/web-tests/rs -o docs/web/snippets +uv run docs/web-tests/mdx_snippets_gen.py \ + -s docs/web-tests/py -s docs/web-tests/ts -s docs/web-tests/rs \ + -s docs/web-tests/blob-v2/py -s docs/web-tests/blob-v2/ts -s docs/web-tests/blob-v2/rs \ + -s rust/lancedb/tests -o docs/web/snippets ``` Review the diff before committing: a few modules currently regenerate with diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 356b49ae7..2fefe91f3 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -329,6 +329,8 @@ code that used the former remote-only argument order. ::: lancedb.table.Table +::: lancedb.table.CompactionOptions + ::: lancedb.table.FragmentStatistics ::: lancedb.table.FragmentSummaryStats diff --git a/docs/web-tests/py/test_doc_storage_examples.py b/docs/web-tests/py/test_doc_storage_examples.py index 13bedb731..99ab4a514 100644 --- a/docs/web-tests/py/test_doc_storage_examples.py +++ b/docs/web-tests/py/test_doc_storage_examples.py @@ -7,6 +7,36 @@ import pyarrow as pa import pytest +def test_compaction_tuning(mem_db): + payload = b"x" * 1016 + table = mem_db.create_table( + "compaction_tuning", [{"id": i, "payload": payload} for i in range(32)] + ) + table.add([{"id": i, "payload": payload} for i in range(32, 64)]) + assert table.stats()["fragment_stats"]["num_fragments"] == 2 + + # --8<-- [start:storage_compaction_tuning] + from lancedb import CompactionOptions + + desired_fragment_bytes = 64 * 1024 + average_row_bytes = 1024 # Estimate from representative rows. + options: CompactionOptions = { + "target_rows_per_fragment": max(1, desired_fragment_bytes // average_row_bytes), + "batch_size": 16, + "num_threads": 1, + } + table.optimize(compaction_options=options) + # --8<-- [end:storage_compaction_tuning] + + assert table.stats()["fragment_stats"]["num_fragments"] == 1 + assert table.count_rows() == 64 + + table.add([{"id": 64, "payload": payload}]) + table.optimize(compaction_options=options) + assert table.stats()["fragment_stats"]["num_fragments"] == 2 + assert sorted(table.to_arrow()["id"].to_pylist()) == list(range(65)) + + def test_rebuild_vector_index(tmp_db): rng = np.random.default_rng(7) table = tmp_db.create_table( diff --git a/docs/web/maintenance/compaction.mdx b/docs/web/maintenance/compaction.mdx index 5bee9c65e..5d2a9c42f 100644 --- a/docs/web/maintenance/compaction.mdx +++ b/docs/web/maintenance/compaction.mdx @@ -12,7 +12,7 @@ availability: background, and `optimize()` is not available there. --- -import { PyStorageRebuildVectorIndex } from '/snippets/doc_storage_examples.mdx'; +import { PyStorageCompactionTuning, PyStorageRebuildVectorIndex, RsStorageCompactionTuning } from '/snippets/doc_storage_examples.mdx'; Every write commits a new table version, and versions accumulate: small files from frequent appends, superseded files from updates and deletes, and rows added @@ -104,3 +104,41 @@ table.optimize(cleanup_older_than=timedelta(days=1)) # reclaim space sooner ``` The best practice is to run `optimize()` after large writes or on a schedule. It also bundles incremental index updates. See [Build & Repair](/indexing/reindexing) for the breakdown. Updates also move rows out of the vector index, so they remain searchable but unindexed. Update or rebuild the index after large update batches. + +{/* +Draft for the Python compaction_options release. Keep this section unpublished +until the first supported release is confirmed. + +### Tune compaction for wide rows {#tune-wide-rows} + +For an existing local table, derive a fragment row target from the expected +row size. These examples use rows averaging 1 KiB and a 64 KiB fragment +target; choose sizes for your workload. Smaller scan batches and one compaction +task at a time can also reduce memory use. + + + + {PyStorageCompactionTuning} + + + {RsStorageCompactionTuning} + + + +The Rust example runs compaction only. Run pruning and index optimization +separately when needed; Python's `optimize()` includes those steps. + +TypeScript's `optimize()` does not currently expose the fragment row target, +scan batch size, or compaction thread count. These tuning examples apply to +Python and Rust local tables. + +Planning remains row based. The row target bounds later passes after existing +oversized fragments have been rewritten; the first pass rewrites those +fragments whole, so peak memory can exceed the target. `max_bytes_per_file` +limits output file size rather than compaction memory. + +The `max_source_rows` and `max_source_bytes` budgets admit whole planned tasks. +If the first task exceeds a budget, that run performs no compaction. See the +[CompactionOptions reference](https://lancedb.github.io/lancedb/python/python/#lancedb.table.CompactionOptions) +for all options and their limits. +*/} diff --git a/docs/web/snippets/doc_storage_examples.mdx b/docs/web/snippets/doc_storage_examples.mdx index 47e2da787..f3b6be6a6 100644 --- a/docs/web/snippets/doc_storage_examples.mdx +++ b/docs/web/snippets/doc_storage_examples.mdx @@ -1,4 +1,8 @@ {/* Auto-generated by scripts/mdx_snippets_gen.py. Do not edit manually. */} +export const PyStorageCompactionTuning = "from lancedb import CompactionOptions\n\ndesired_fragment_bytes = 64 * 1024\naverage_row_bytes = 1024 # Estimate from representative rows.\noptions: CompactionOptions = {\n \"target_rows_per_fragment\": max(1, desired_fragment_bytes // average_row_bytes),\n \"batch_size\": 16,\n \"num_threads\": 1,\n}\ntable.optimize(compaction_options=options)\n"; + export const PyStorageRebuildVectorIndex = "from lancedb.index import IvfFlat\n\ntable.create_index(\n \"vector\",\n config=IvfFlat(distance_type=\"l2\", num_partitions=4),\n replace=True,\n)\n"; +export const RsStorageCompactionTuning = "use lancedb::table::{CompactionOptions, OptimizeAction};\n\nlet desired_fragment_bytes: usize = 64 * 1024;\nlet average_row_bytes: usize = 1024; // Estimate from representative rows.\nlet options = CompactionOptions {\n target_rows_per_fragment: (desired_fragment_bytes / average_row_bytes).max(1),\n batch_size: Some(16),\n num_threads: Some(1),\n ..Default::default()\n};\ntable\n .optimize(OptimizeAction::Compact {\n options,\n remap_options: None,\n })\n .await?;\n"; + diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 5e879f61d..6c71c6261 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -49,7 +49,7 @@ from .materialized_view import ( MaterializedViewDefinition, ) from .view import ViewDescription as ViewDescription -from .table import AsyncTable, Table +from .table import AsyncTable, CompactionOptions, Table from .types import BaseTokenizerType from ._lancedb import Session from .namespace import ( @@ -610,6 +610,7 @@ __all__ = [ "AsyncSqlQuery", "AsyncLanceNamespaceDBConnection", "AsyncTable", + "CompactionOptions", "FtsToken", "col", "Expr", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 0b4fc82f1..1f3cfa9ed 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -544,6 +544,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 6b7835470..935e07eec 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -67,6 +67,7 @@ from ..table import ( AsyncTable, BlobMode, Branches, + CompactionOptions, IndexStatistics, Query, Table, @@ -1071,6 +1072,7 @@ class RemoteTable(Table): cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, retrain: bool = False, + compaction_options: Optional[CompactionOptions] = None, ): """optimize() is not supported for remote tables.""" return LOOP.run( diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index f1e538c14..048c46a6a 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -25,6 +25,7 @@ from typing import ( Optional, Sequence, Tuple, + TypedDict, Union, overload, ) @@ -285,6 +286,91 @@ IndexConfigType = Union[ FTS, ] + +class CompactionOptions(TypedDict, total=False): + """Options that control file compaction during table optimization. + + Unspecified options use Lance's defaults. + + Compaction planning is row based. Lowering ``target_rows_per_fragment`` + based on the expected row size can bound later compaction passes once + oversized fragments have been rewritten. It does not split an existing + fragment, so the first pass over an oversized fragment is not subject to + that bound. ``max_bytes_per_file`` limits output file size, not compaction + memory. Source budgets keep whole planned tasks; if the first task exceeds + a budget, that run performs no compaction work. + + Examples + -------- + Derive a steady-state row target from the expected row size: + + >>> desired_fragment_bytes = 750 * 1024 * 1024 + >>> average_row_bytes = 1_500_000 + >>> options: CompactionOptions = { + ... "target_rows_per_fragment": max( + ... 1, desired_fragment_bytes // average_row_bytes + ... ), + ... } + >>> await table.optimize(compaction_options=options) # doctest: +SKIP + """ + + target_rows_per_fragment: int + """Target rows per fragment; existing oversized fragments are not split.""" + + max_rows_per_group: int + """Maximum number of rows per row group (default: 1,024).""" + + max_bytes_per_file: Optional[int] + """Maximum output data-file size; this does not bound compaction memory.""" + + 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_source_rows: Optional[int] + """Maximum live source rows per run, applied to whole planned tasks.""" + + max_source_bytes: Optional[int] + """Maximum source bytes per run, applied to whole planned tasks.""" + + excluded_fragment_ids: List[int] + """Fragment IDs to leave unchanged and use as planning boundaries.""" + + max_overlays_per_fragment: Optional[int] + """Maximum overlays before a fragment is fully compacted.""" + + data_storage_version: str + """Data-file version or release selector for compacted output.""" + + # Known distance metrics for legacy API detection KNOWN_METRICS = {"l2", "cosine", "dot", "hamming"} @@ -2767,6 +2853,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. @@ -2806,6 +2893,11 @@ 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. For large rows, derive a lower + ``target_rows_per_fragment`` from the expected row size to bound later + passes. This does not cap the first pass over an existing oversized + fragment; see [CompactionOptions][lancedb.table.CompactionOptions]. Notes ----- @@ -4893,6 +4985,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. @@ -4932,6 +5025,11 @@ 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. For large rows, derive a lower + ``target_rows_per_fragment`` from the expected row size to bound later + passes. This does not cap the first pass over an existing oversized + fragment; see [CompactionOptions][lancedb.table.CompactionOptions]. Notes ----- @@ -4945,6 +5043,7 @@ class LanceTable(Table): self._table._do_optimize( _optimize_cleanup_since_ms(cleanup_older_than, retrain), delete_unverified, + compaction_options, ) ) @@ -7620,6 +7719,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. @@ -7659,6 +7759,11 @@ class AsyncTable: retrain: bool, default False This parameter is no longer used and is deprecated. + compaction_options: CompactionOptions, optional + Options that control file compaction. For large rows, derive a lower + ``target_rows_per_fragment`` from the expected row size to bound later + passes. This does not cap the first pass over an existing oversized + fragment; see [CompactionOptions][lancedb.table.CompactionOptions]. Notes ----- @@ -7671,14 +7776,19 @@ class AsyncTable: return await self._do_optimize( _optimize_cleanup_since_ms(cleanup_older_than, retrain), delete_unverified, + compaction_options, ) async def _do_optimize( - self, cleanup_since_ms: Optional[int], delete_unverified: bool + self, + cleanup_since_ms: Optional[int], + delete_unverified: bool, + compaction_options: Optional[CompactionOptions] = None, ) -> OptimizeStats: 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 97f65b180..458f419ae 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -14,7 +14,7 @@ from concurrent.futures import ThreadPoolExecutor from datetime import date, datetime, timedelta from decimal import Decimal, localcontext from time import sleep -from typing import List +from typing import Any, List from unittest.mock import MagicMock, patch import lancedb @@ -5014,6 +5014,134 @@ 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}) + + +def test_optimize_compaction_options_sync(mem_db: DBConnection): + table = mem_db.create_table("test", data=[{"x": 1}]) + with pytest.raises(ValueError, match="target_rows_per_fragment must be between"): + table.optimize(compaction_options={"target_rows_per_fragment": 0}) + + +@pytest.mark.parametrize("option", ["max_source_rows", "max_source_bytes"]) +@pytest.mark.asyncio +async def test_optimize_compaction_source_limits( + mem_db_async: AsyncConnection, option: str +): + 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": 3, + option: 1, + } + ) + assert stats.compaction.fragments_removed == 0 + assert stats.compaction.fragments_added == 0 + + +@pytest.mark.asyncio +async def test_optimize_compaction_data_storage_version( + 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": 3, + "data_storage_version": "2.2", + } + ) + assert stats.compaction.fragments_removed == 2 + assert stats.compaction.fragments_added == 1 + + table = await mem_db_async.create_table("cross_family", data=[{"x": 1}]) + await table.add([{"x": 2}]) + with pytest.raises(ValueError, match="V1 and V2 storage versions cannot be mixed"): + await table.optimize( + compaction_options={ + "target_rows_per_fragment": 3, + "data_storage_version": "legacy", + } + ) + + +@pytest.mark.asyncio +async def test_optimize_compaction_excluded_fragments(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": 3, + "excluded_fragment_ids": [0], + } + ) + 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 + + +@pytest.mark.parametrize( + ("option", "value", "message"), + [ + ("target_rows_per_fragment", 0, "must be between 1 and 4294967295"), + ("max_rows_per_group", 0, "must be between 1 and 4294967295"), + ("batch_size", 0, "must be between 1 and 4294967295"), + ("num_threads", 0, "must be greater than 0"), + ("target_rows_per_fragment", 2**32, "must be between 1 and 4294967295"), + ("max_rows_per_group", 2**32, "must be between 1 and 4294967295"), + ("batch_size", 2**32, "must be between 1 and 4294967295"), + ("io_buffer_size", 2**63, "must be at most 9223372036854775807"), + ("max_source_rows", 0, "must be greater than 0"), + ("max_source_bytes", 0, "must be greater than 0"), + ("data_storage_version", "invalid", "Unknown Lance storage version: invalid"), + ( + "excluded_fragment_ids", + [-1], + "must contain values between 0 and 4294967295", + ), + ( + "excluded_fragment_ids", + [2**32], + "must contain values between 0 and 4294967295", + ), + ], +) +@pytest.mark.asyncio +async def test_optimize_compaction_options_validation( + mem_db_async: AsyncConnection, option: str, value: Any, message: str +): + table = await mem_db_async.create_table("test", data=[{"x": 1}]) + with pytest.raises(ValueError, match=message): + await table.optimize(compaction_options={option: value}) + + @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 e58b8a139..62fc9e611 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -25,10 +25,12 @@ use arrow::{ pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; use lancedb::blob::{BlobDescriptor, BlobFile, BlobRangeRequest}; +use lancedb::connection::LanceFileVersion; 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::{ @@ -107,6 +109,136 @@ enum PredicateArg { Sql(String), } +fn validate_positive_u32(value: u64, name: &str) -> PyResult { + if !(1..=u32::MAX as u64).contains(&value) { + return Err(PyValueError::new_err(format!( + "{name} must be between 1 and {}", + u32::MAX + ))); + } + Ok(value as usize) +} + +fn positive_u32(value: &Bound<'_, PyAny>, name: &str) -> PyResult { + validate_positive_u32(value.extract()?, name) +} + +fn optional_positive_u32(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + value + .extract::>()? + .map(|value| validate_positive_u32(value, name)) + .transpose() +} + +fn optional_positive_usize(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + let value: Option = value.extract()?; + if value == Some(0) { + return Err(PyValueError::new_err(format!( + "{name} must be greater than 0" + ))); + } + Ok(value) +} + +fn optional_positive_u64(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + let value: Option = value.extract()?; + if value == Some(0) { + return Err(PyValueError::new_err(format!( + "{name} must be greater than 0" + ))); + } + Ok(value) +} + +fn u32_list(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + value + .extract::>()? + .into_iter() + .map(|value| { + u32::try_from(value).map_err(|_| { + PyValueError::new_err(format!( + "{name} must contain values between 0 and {}", + u32::MAX + )) + }) + }) + .collect() +} + +fn optional_i64_bounded_u64(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + let value: Option = value.extract()?; + if value.is_some_and(|value| value > i64::MAX as u64) { + return Err(PyValueError::new_err(format!( + "{name} must be at most {}", + i64::MAX + ))); + } + Ok(value) +} + +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 = positive_u32(&value, &key)? + } + "max_rows_per_group" => parsed.max_rows_per_group = positive_u32(&value, &key)?, + "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 = optional_positive_usize(&value, &key)?, + "batch_size" => parsed.batch_size = optional_positive_u32(&value, &key)?, + "io_buffer_size" => parsed.io_buffer_size = optional_i64_bounded_u64(&value, &key)?, + "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_source_rows" => parsed.max_source_rows = optional_positive_usize(&value, &key)?, + "max_source_bytes" => parsed.max_source_bytes = optional_positive_u64(&value, &key)?, + "excluded_fragment_ids" => parsed.excluded_fragment_ids = u32_list(&value, &key)?, + "max_overlays_per_fragment" => parsed.max_overlays_per_fragment = value.extract()?, + "data_storage_version" => { + let version: String = value.extract()?; + parsed.data_storage_version = Some( + version + .parse::() + .map_err(|err| PyValueError::new_err(err.to_string()))?, + ); + } + _ => { + 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)] @@ -1569,13 +1701,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!( @@ -1591,7 +1725,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 3fcf4bdb7..d6575b959 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -107,7 +107,9 @@ use futures::future::join_all; pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTags}; pub use lance::dataset::scanner::DatasetRecordBatchStream; pub use lance_index::optimize::OptimizeOptions; -pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats}; +pub use optimize::{ + CompactionMode, CompactionOptions, IndexRemapMode, OptimizeAction, OptimizeStats, +}; pub use refresh::RefreshColumnResult; pub use rte_stats::{MemtableStats, RTEStats, SsTableStats, TableShardStats}; pub use schema_evolution::{ diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index ebe03359c..ee86e7989 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -16,7 +16,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; diff --git a/rust/lancedb/tests/doc_storage_examples.rs b/rust/lancedb/tests/doc_storage_examples.rs new file mode 100644 index 000000000..dd3acc597 --- /dev/null +++ b/rust/lancedb/tests/doc_storage_examples.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::ops::Range; +use std::sync::Arc; + +use arrow_array::{ + ArrayRef, BinaryArray, Int64Array, RecordBatch, RecordBatchIterator, RecordBatchReader, +}; + +fn wide_rows(ids: Range) -> anyhow::Result> { + let payload = [b'x'; 1016]; + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int64Array::from_iter_values(ids.clone())) as ArrayRef, + ), + ( + "payload", + Arc::new(BinaryArray::from_iter_values( + ids.map(|_| payload.as_slice()), + )) as ArrayRef, + ), + ])?; + let schema = batch.schema(); + Ok(Box::new(RecordBatchIterator::new( + std::iter::once(Ok(batch)), + schema, + ))) +} + +async fn tune_compaction(table: &lancedb::Table) -> lancedb::Result<()> { + // --8<-- [start:storage_compaction_tuning] + use lancedb::table::{CompactionOptions, OptimizeAction}; + + let desired_fragment_bytes: usize = 64 * 1024; + let average_row_bytes: usize = 1024; // Estimate from representative rows. + let options = CompactionOptions { + target_rows_per_fragment: (desired_fragment_bytes / average_row_bytes).max(1), + batch_size: Some(16), + num_threads: Some(1), + ..Default::default() + }; + table + .optimize(OptimizeAction::Compact { + options, + remap_options: None, + }) + .await?; + // --8<-- [end:storage_compaction_tuning] + Ok(()) +} + +#[tokio::test] +async fn compaction_tuning() -> anyhow::Result<()> { + let directory = tempfile::tempdir()?; + let db = lancedb::connect(directory.path().to_str().unwrap()) + .execute() + .await?; + let table = db + .create_table("compaction_tuning", wide_rows(0..32)?) + .execute() + .await?; + table.add(wide_rows(32..64)?).execute().await?; + assert_eq!(table.stats().await?.fragment_stats.num_fragments, 2); + + tune_compaction(&table).await?; + assert_eq!(table.stats().await?.fragment_stats.num_fragments, 1); + assert_eq!(table.count_rows(None).await?, 64); + + table.add(wide_rows(64..65)?).execute().await?; + tune_compaction(&table).await?; + assert_eq!(table.stats().await?.fragment_stats.num_fragments, 2); + assert_eq!(table.count_rows(None).await?, 65); + Ok(()) +}