From 546a28ad2c8e8bd8a95540afc80f27091de514db Mon Sep 17 00:00:00 2001 From: Bruno Ramirez Date: Mon, 9 Nov 2026 11:19:47 -0800 Subject: [PATCH] feat: support remote compaction tuning and table inheritance (#4886) --- docs/src/js/classes/Table.md | 3 + docs/src/js/globals.md | 1 + docs/src/js/interfaces/CompactionOptions.md | 49 +++++++++ docs/src/js/interfaces/OptimizeOptions.md | 22 ++++ docs/src/python/python.md | 24 +++- docs/web/maintenance/compaction.mdx | 7 +- nodejs/__test__/table.test.ts | 89 ++++++++++++++- nodejs/lancedb/index.ts | 1 + nodejs/lancedb/table.ts | 41 +++++++ nodejs/src/table.rs | 65 ++++++++++- python/python/lancedb/remote/table.py | 11 +- python/python/lancedb/table.py | 32 ++++-- python/python/tests/test_remote_db.py | 116 +++++++++++++++++--- python/src/table.rs | 24 +++- rust/lancedb/src/remote/sql_test.rs | 112 ++++++++++++++++++- rust/lancedb/src/remote/table.rs | 56 ++++++---- rust/lancedb/src/remote/table/optimize.rs | 98 +++++++++++++++++ rust/lancedb/src/table.rs | 69 +++++++++++- rust/lancedb/src/table/optimize.rs | 85 +++++++++++++- 19 files changed, 830 insertions(+), 75 deletions(-) create mode 100644 docs/src/js/interfaces/CompactionOptions.md create mode 100644 rust/lancedb/src/remote/table/optimize.rs diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 8bdc6b5f8..1fad723fe 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -948,6 +948,9 @@ supported. Pruning old versions and index optimization are silently skipped there instead of being attempted, so this call resolves successfully, with zeroed pruning statistics, once the compaction job has been submitted; it does not wait for that job to finish. +`compactionOptions` can tune the four per-job settings documented in +[CompactionOptions](../interfaces/CompactionOptions.md). Older-server errors are propagated without +retrying with the settings removed. #### Parameters diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 46f8db833..dab1c61bc 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -74,6 +74,7 @@ - [ClientConfig](interfaces/ClientConfig.md) - [ColumnAlteration](interfaces/ColumnAlteration.md) - [ColumnOrdering](interfaces/ColumnOrdering.md) +- [CompactionOptions](interfaces/CompactionOptions.md) - [CompactionStats](interfaces/CompactionStats.md) - [ConnectNamespaceOptions](interfaces/ConnectNamespaceOptions.md) - [ConnectionOptions](interfaces/ConnectionOptions.md) diff --git a/docs/src/js/interfaces/CompactionOptions.md b/docs/src/js/interfaces/CompactionOptions.md new file mode 100644 index 000000000..135266567 --- /dev/null +++ b/docs/src/js/interfaces/CompactionOptions.md @@ -0,0 +1,49 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / CompactionOptions + +# Interface: CompactionOptions + +Per-invocation compaction settings. Values must be safe integers. + +## Properties + +### batchSize? + +```ts +optional batchSize: number; +``` + +Positive scanner batch size in rows. + +*** + +### ioBufferSize? + +```ts +optional ioBufferSize: number; +``` + +Scanner I/O read-ahead budget in bytes; zero disables byte-based I/O backpressure. + +*** + +### maxSourceRows? + +```ts +optional maxSourceRows: number; +``` + +Positive live-source-row budget across whole tasks; an oversized first task can produce no work. + +*** + +### targetRowsPerFragment? + +```ts +optional targetRowsPerFragment: number; +``` + +Positive target rows per output fragment; also affects candidate selection. diff --git a/docs/src/js/interfaces/OptimizeOptions.md b/docs/src/js/interfaces/OptimizeOptions.md index 110eb3813..faf1d2724 100644 --- a/docs/src/js/interfaces/OptimizeOptions.md +++ b/docs/src/js/interfaces/OptimizeOptions.md @@ -33,6 +33,28 @@ tbl.optimize({cleanupOlderThan: new Date()}); *** +### compactionOptions? + +```ts +optional compactionOptions: CompactionOptions; +``` + +Compaction settings for this invocation. Omitted remote values inherit table +tuning, then deployment/built-in defaults. +Remote tables require a server supporting these SQL options; settings do +not change table policy. This is not a runtime or total-memory limit. + +#### Example + +```ts +await tbl.optimize({ compactionOptions: { + targetRowsPerFragment: 1000, maxSourceRows: 5000, + ioBufferSize: 64 * 1024 * 1024, batchSize: 128, +} }); +``` + +*** + ### deleteUnverified ```ts diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 3658f2ad3..7e5140928 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -352,9 +352,27 @@ accept the LanceDB Embedded options: cleanup and compaction warn that they are n `optimize` submits a durable compaction job and returns once that submission is accepted, not once compaction itself finishes. Pruning and index optimization are not yet supported remotely, so it silently skips them and returns with -zeroed pruning statistics rather than attempting and failing. A non-default -`compaction_options` raises `NotImplementedError` before anything is -submitted, since there is no way yet to forward those settings to the server. +zeroed pruning statistics. Remote `compaction_options` support per-job +`target_rows_per_fragment`, `max_source_rows`, `io_buffer_size`, and `batch_size` +on compatible servers. Other non-default settings raise `NotImplementedError` +before submission; older-server errors are returned without dropping options. +The source-row budget selects whole tasks and can produce no work if the first +task exceeds it. An I/O buffer size of zero disables byte-based I/O backpressure. +Use a positive size to limit I/O read-ahead; this is not a total-memory limit. +Each field resolves as invocation option → table tuning → deployment/built-in +default. These options apply only to the submitted job and leave table policy +unchanged. Table `enabled` controls automatic scheduling; explicit manual runs +remain allowed. + +```python +table.optimize(compaction_options={ + "target_rows_per_fragment": 100_000, + "max_source_rows": 1_000_000, + "io_buffer_size": 64 * 1024 * 1024, + "batch_size": 1024, +}) +``` + Removed Tantivy options produce the same explanatory errors on both backends. The old LanceDB Enterprise keyword names `predicate` for `delete`, `index_name` for diff --git a/docs/web/maintenance/compaction.mdx b/docs/web/maintenance/compaction.mdx index 3110c0858..9c8765293 100644 --- a/docs/web/maintenance/compaction.mdx +++ b/docs/web/maintenance/compaction.mdx @@ -150,9 +150,10 @@ task at a time can also reduce memory use. 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 tables in LanceDB Embedded. +TypeScript's `optimize()` exposes the fragment row target through +`compactionOptions.targetRowsPerFragment` and the scanner batch size through +`compactionOptions.batchSize`. Compaction thread count is not exposed. The tuning +examples above use Python and Rust tables in LanceDB Embedded. Planning remains row based. The row target bounds later passes after existing oversized fragments have been rewritten; the first pass rewrites those diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index af66ebb60..c3888d6c0 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -11,6 +11,7 @@ import { AutoQuery, type BlobInput, type BlobMode, + type CompactionOptions, Connection, MatchQuery, PhraseQuery, @@ -3997,7 +3998,93 @@ it("passes cleanupOlderThan to the native binding as an absolute timestamp", asy await table.optimize({ cleanupOlderThan: cutoff, deleteUnverified: true }); - expect(optimize).toHaveBeenCalledWith(cutoff.getTime(), true); + expect(optimize).toHaveBeenCalledWith(cutoff.getTime(), true, undefined); +}); + +it("forwards per-job compaction options to the native binding", async () => { + const optimize = jest.fn().mockResolvedValue({}); + const table = new LocalTable({ optimize } as never); + const compactionOptions = { + targetRowsPerFragment: 20, + maxSourceRows: 25, + ioBufferSize: 0, + batchSize: 8, + }; + await table.optimize({ compactionOptions }); + expect(optimize).toHaveBeenCalledWith( + undefined, + undefined, + compactionOptions, + ); +}); + +it("rejects unsupported JavaScript compaction options before native submission", async () => { + const optimize = jest.fn(); + const table = new LocalTable({ optimize } as never); + await expect( + table.optimize({ + compactionOptions: { + maxSourceBytes: 1024, + } as unknown as CompactionOptions, + }), + ).rejects.toThrow("Unsupported compaction option: maxSourceBytes"); + expect(optimize).not.toHaveBeenCalled(); +}); + +describe("per-job compaction options", () => { + let tmpDir: tmp.DirResult; + let table: Table; + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + table = await db.createTable( + "bounded", + Array.from({ length: 10 }, (_, id) => ({ id })), + ); + for (let f = 1; f < 8; f++) { + await table.add( + Array.from({ length: 10 }, (_, id) => ({ id: f * 10 + id })), + ); + } + }); + afterEach(() => { + tmpDir.removeCallback(); + }); + it("bounds source work and preserves zero I/O buffering", async () => { + const stats = await table.optimize({ + compactionOptions: { + targetRowsPerFragment: 20, + maxSourceRows: 25, + ioBufferSize: 0, + batchSize: 8, + }, + }); + expect(stats.compaction.fragmentsRemoved).toBe(2); + expect(stats.compaction.fragmentsAdded).toBe(1); + expect(await table.countRows()).toBe(80); + const defaults = await table.optimize(); + expect(defaults.compaction.fragmentsRemoved).toBe(7); + expect(defaults.compaction.fragmentsAdded).toBe(1); + expect(await table.countRows()).toBe(80); + }); + const invalid: CompactionOptions[] = [ + { targetRowsPerFragment: 0 }, + { maxSourceRows: 0 }, + { batchSize: 0 }, + { ioBufferSize: -1 }, + { batchSize: 1.5 }, + { maxSourceRows: Number.NaN }, + { ioBufferSize: Number.POSITIVE_INFINITY }, + { maxSourceRows: Number.MAX_SAFE_INTEGER + 1 }, + { targetRowsPerFragment: 2 ** 32 }, + { batchSize: 2 ** 32 }, + ]; + it.each(invalid)("rejects invalid settings %p", async (compactionOptions) => { + await expect(table.optimize({ compactionOptions })).rejects.toThrow( + /safe integer/, + ); + expect(await table.countRows()).toBe(80); + }); }); describe.each(arrowVersions)( diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 0e054d70e..9644b0060 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -165,6 +165,7 @@ export { AddDataOptions, UpdateOptions, OptimizeOptions, + CompactionOptions, Version, WriteProgress, FtsToken, diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index de22ef0f8..5ef298c8c 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -152,7 +152,31 @@ export interface UpdateOptions { where: string; } +/** Per-invocation compaction settings. Values must be safe integers. */ +export interface CompactionOptions { + /** Positive target rows per output fragment; also affects candidate selection. */ + targetRowsPerFragment?: number; + /** Positive live-source-row budget across whole tasks; an oversized first task can produce no work. */ + maxSourceRows?: number; + /** Scanner I/O read-ahead budget in bytes; zero disables byte-based I/O backpressure. */ + ioBufferSize?: number; + /** Positive scanner batch size in rows. */ + batchSize?: number; +} + export interface OptimizeOptions { + /** + * Compaction settings for this invocation. Omitted remote values inherit table + * tuning, then deployment/built-in defaults. + * Remote tables require a server supporting these SQL options; settings do + * not change table policy. This is not a runtime or total-memory limit. + * @example + * await tbl.optimize({ compactionOptions: { + * targetRowsPerFragment: 1000, maxSourceRows: 5000, + * ioBufferSize: 64 * 1024 * 1024, batchSize: 128, + * } }); + */ + compactionOptions?: CompactionOptions; /** * If set then all versions older than the given date * be removed. The current version will never be removed. @@ -1047,6 +1071,9 @@ export abstract class Table { * skipped there instead of being attempted, so this call resolves * successfully, with zeroed pruning statistics, once the compaction job * has been submitted; it does not wait for that job to finish. + * `compactionOptions` can tune the four per-job settings documented in + * {@link CompactionOptions}. Older-server errors are propagated without + * retrying with the settings removed. */ abstract optimize(options?: Partial): Promise; /** @@ -1646,9 +1673,23 @@ export class LocalTable extends Table { } async optimize(options?: Partial): Promise { + // JavaScript callers can supply extra keys that TypeScript would reject. + // Refuse them before the native binding can silently discard a work limit. + const supported = new Set([ + "targetRowsPerFragment", + "maxSourceRows", + "ioBufferSize", + "batchSize", + ]); + for (const key of Object.keys(options?.compactionOptions ?? {})) { + if (!supported.has(key)) { + throw new Error(`Unsupported compaction option: ${key}`); + } + } return await this.inner.optimize( options?.cleanupOlderThan?.getTime(), options?.deleteUnverified, + options?.compactionOptions, ); } diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 5c3805f51..ebbc4a060 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -23,6 +23,63 @@ use crate::merge::NativeMergeInsertBuilder; use crate::query::{Query, TakeQuery, VectorQuery}; use crate::util::schema_to_buffer; +/// Per-invocation file-compaction settings. All numbers must be safe integers. +#[napi(object)] +pub struct CompactionOptions { + pub target_rows_per_fragment: Option, + pub max_source_rows: Option, + pub io_buffer_size: Option, + pub batch_size: Option, +} + +impl CompactionOptions { + fn into_overrides(self) -> napi::Result { + const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991; + fn integer( + value: Option, + name: &str, + max: u64, + allow_zero: bool, + ) -> napi::Result> { + let Some(value) = value else { + return Ok(None); + }; + let min = if allow_zero { 0 } else { 1 }; + let max = max.min(MAX_SAFE_INTEGER); + if !value.is_finite() + || value.fract() != 0.0 + || value < min as f64 + || value > max as f64 + { + return Err(napi::Error::from_reason(format!( + "{name} must be a safe integer between {min} and {max}" + ))); + } + Ok(Some(value as u64)) + } + let mut options = lancedb::table::optimize::CompactionOverrides::default(); + if let Some(target) = integer( + self.target_rows_per_fragment, + "targetRowsPerFragment", + u32::MAX as u64, + false, + )? { + options.target_rows_per_fragment = Some(target as usize); + } + options.max_source_rows = integer( + self.max_source_rows, + "maxSourceRows", + usize::MAX as u64, + false, + )? + .map(|v| v as usize); + options.io_buffer_size = integer(self.io_buffer_size, "ioBufferSize", u64::MAX, true)?; + options.batch_size = + integer(self.batch_size, "batchSize", u32::MAX as u64, false)?.map(|v| v as usize); + Ok(options) + } +} + #[napi] pub struct Table { // We keep a duplicate of the table name so we can use it for error @@ -732,6 +789,7 @@ impl Table { &self, before_timestamp_ms: Option, delete_unverified: Option, + compaction_options: Option, ) -> napi::Result { let inner = self.inner_ref()?; @@ -746,8 +804,11 @@ impl Table { .transpose()?; let stats = inner - .optimize_full( - lancedb::table::CompactionOptions::default(), + .optimize_full_with_overrides( + compaction_options + .map(CompactionOptions::into_overrides) + .transpose()? + .unwrap_or_default(), None, before_timestamp, delete_unverified, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 6bef59922..8f1f741fb 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1091,10 +1091,13 @@ class RemoteTable(Table): Enterprise job. Pruning old versions and index optimization are not yet supported remotely, so this call skips them and returns successfully with zeroed pruning statistics once the compaction job - has been submitted; it does not wait for that job to finish. Passing - a non-default ``compaction_options`` raises ``NotImplementedError`` - before anything is submitted, since there is no way yet to forward - those settings to the server. + has been submitted; it does not wait for that job to finish. The + per-job options ``target_rows_per_fragment``, ``max_source_rows``, + ``io_buffer_size``, and ``batch_size`` are supported on compatible + servers. Omitted fields inherit table tuning, then deployment/built-in + defaults. Other non-default settings raise ``NotImplementedError`` + before submission. Errors from older servers are propagated without + retrying with the settings removed. """ return LOOP.run( self._table.optimize( diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index f43df8b40..75975a2c3 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -307,7 +307,12 @@ IndexConfigType = Union[ class CompactionOptions(TypedDict, total=False): """Options that control file compaction during table optimization. - Unspecified options use Lance's defaults. + Unspecified options use Lance's defaults for local tables. Remote tables + forward ``target_rows_per_fragment``, ``max_source_rows``, ``io_buffer_size``, + and ``batch_size`` as per-job SQL options; other non-default settings are + unsupported. Omitted remote options inherit table tuning, then deployment/built-in + defaults. The server must support these SQL options. These settings do not + change table policy. Compaction planning is row based. Lowering ``target_rows_per_fragment`` based on the expected row size can bound later compaction passes once @@ -353,7 +358,10 @@ class CompactionOptions(TypedDict, total=False): """Number of rows per input scan batch.""" io_buffer_size: Optional[int] - """Maximum number of bytes queued in the input scan I/O buffer.""" + """Input scan I/O read-ahead budget in bytes. + + Zero disables byte-based I/O backpressure. + """ defer_index_remap: bool """Whether to defer index remapping during compaction (default: False).""" @@ -2943,9 +2951,13 @@ class Table(ABC): are silently skipped there instead of being attempted, so this call returns successfully, with zeroed pruning statistics, once the compaction job has been submitted; it does not wait for that - job to finish. A non-default ``compaction_options`` raises - ``NotImplementedError`` before anything is submitted, since - there is no way yet to forward those settings to the server. + job to finish. Remote compaction supports per-job + ``target_rows_per_fragment``, ``max_source_rows``, ``io_buffer_size``, + and ``batch_size``. Omitted fields inherit table tuning, then + deployment/built-in defaults. Other non-default compaction settings raise + ``NotImplementedError`` before submission. The server must support + these SQL options; older-server errors are propagated without + dropping the requested settings. Modeled after ``VACUUM`` in PostgreSQL. @@ -8477,9 +8489,13 @@ class AsyncTable: are silently skipped there instead of being attempted, so this call returns successfully, with zeroed pruning statistics, once the compaction job has been submitted; it does not wait for that - job to finish. A non-default ``compaction_options`` raises - ``NotImplementedError`` before anything is submitted, since - there is no way yet to forward those settings to the server. + job to finish. Remote compaction supports per-job + ``target_rows_per_fragment``, ``max_source_rows``, ``io_buffer_size``, + and ``batch_size``. Omitted fields inherit table tuning, then + deployment/built-in defaults. Other non-default compaction settings raise + ``NotImplementedError`` before submission. The server must support + these SQL options; older-server errors are propagated without + dropping the requested settings. Modeled after ``VACUUM`` in PostgreSQL. diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 8a65a2a1c..25d9db22a 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -5837,10 +5837,8 @@ def _create_table_handler(request): @contextlib.contextmanager -def _mock_remote_table_with_sql(flight_server): - """A remote table backed by a mock REST server (for create/open) and the - given mock Flight SQL server (for the SQL statements `optimize()` submits - as jobs).""" +def _mock_remote_sql_connection(flight_server): + """Connection arguments for mock REST and Flight SQL servers.""" with http.server.ThreadingHTTPServer( ("127.0.0.1", 0), make_mock_http_handler(_create_table_handler) ) as rest_server: @@ -5848,18 +5846,106 @@ def _mock_remote_table_with_sql(flight_server): handle = start_server_thread(rest_server) try: with flight_server: - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=f"http://127.0.0.1:{rest_port}", - sql_host_override=f"grpc://127.0.0.1:{flight_server.port}", - ) - yield db.create_table("test", [{"id": 1}]) + yield { + "uri": "db://dev", + "api_key": "fake", + "host_override": f"http://127.0.0.1:{rest_port}", + "sql_host_override": f"grpc://127.0.0.1:{flight_server.port}", + } finally: rest_server.shutdown() handle.join() +@contextlib.contextmanager +def _mock_remote_table_with_sql(flight_server): + with _mock_remote_sql_connection(flight_server) as kwargs: + yield lancedb.connect(**kwargs).create_table("test", [{"id": 1}]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "options,sql_options", + [ + ({"target_rows_per_fragment": 100}, [b"target_rows_per_fragment = 100"]), + ( + {"target_rows_per_fragment": 1_048_576}, + [b"target_rows_per_fragment = 1048576"], + ), + ({"max_source_rows": 500}, [b"max_source_rows = 500"]), + ({"io_buffer_size": 0}, [b"io_buffer_size = 0"]), + ({"batch_size": 16}, [b"batch_size = 16"]), + ( + { + "target_rows_per_fragment": 100, + "max_source_rows": 500, + "io_buffer_size": 65536, + "batch_size": 16, + }, + [ + b"target_rows_per_fragment = 100", + b"max_source_rows = 500", + b"io_buffer_size = 65536", + b"batch_size = 16", + ], + ), + ({}, []), + ], +) +async def test_remote_optimize_forwards_per_job_compaction_options( + asynchronous, options, sql_options +): + flight = pytest.importorskip("pyarrow.flight") + + class AckServer(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") + self.tickets = [] + + def do_get(self, context, ticket): + self.tickets.append(ticket.ticket) + return flight.RecordBatchStream( + pa.table({"ok": pa.array([], type=pa.bool_())}) + ) + + server = AckServer() + with _mock_remote_sql_connection(server) as kwargs: + if asynchronous: + db = await lancedb.connect_async(**kwargs) + table = await db.create_table("test", [{"id": 1}]) + stats = await table.optimize(compaction_options=options) + else: + table = lancedb.connect(**kwargs).create_table("test", [{"id": 1}]) + stats = table.optimize(compaction_options=options) + assert len(server.tickets) == 1 + assert b"OPTIMIZE" in server.tickets[0] + assert all(option in server.tickets[0] for option in sql_options) + if not sql_options: + assert b"WITH (" not in server.tickets[0] + assert stats.prune.bytes_removed == 0 + + +def test_remote_optimize_does_not_drop_options_after_server_rejection(): + flight = pytest.importorskip("pyarrow.flight") + + class OldServer(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") + self.tickets = [] + + def do_get(self, context, ticket): + self.tickets.append(ticket.ticket) + raise flight.FlightServerError("Unknown OPTIMIZE option 'max_source_rows'") + + server = OldServer() + with _mock_remote_table_with_sql(server) as table: + with pytest.raises(Exception, match="max_source_rows"): + table.optimize(compaction_options={"max_source_rows": 500}) + assert server.tickets + assert all(b"max_source_rows = 500" in ticket for ticket in server.tickets) + + def test_remote_optimize_skips_prune_and_index(): """Regression test: a remote table's `optimize()` used to submit a compaction job and then raise `NotImplementedError` once it went on to @@ -5868,7 +5954,7 @@ def test_remote_optimize_skips_prune_and_index(): zeroed pruning statistics, once compaction is submitted. The SQL job submission for a query without parameters (like - `OPTIMIZE TABLE ... REWRITE DATA`) goes over plain Flight `do_get`, not + `OPTIMIZE ... REWRITE DATA`) goes over plain Flight `do_get`, not `do_exchange`: the client sends a ticket wrapping the statement and reads back whatever the server returns as the job's result stream. """ @@ -5890,12 +5976,12 @@ def test_remote_optimize_skips_prune_and_index(): stats = table.optimize() assert len(server.tickets) == 1 - assert b"OPTIMIZE TABLE" in server.tickets[0] + assert b"OPTIMIZE" in server.tickets[0] assert stats.prune.bytes_removed == 0 assert stats.prune.old_versions_removed == 0 -def test_remote_optimize_rejects_non_default_compaction_options_before_submitting(): +def test_remote_optimize_rejects_unsupported_compaction_options_before_submitting(): """Regression test: `RemoteTable.optimize()` used to drop `compaction_options` instead of forwarding it to the Rust layer, so a non-default value was silently ignored -- the server-side guard against @@ -5917,7 +6003,7 @@ def test_remote_optimize_rejects_non_default_compaction_options_before_submittin ) server = UnreachableServer() - options: CompactionOptions = {"target_rows_per_fragment": 5} + options: CompactionOptions = {"max_source_bytes": 1024} with _mock_remote_table_with_sql(server) as table: with pytest.raises(NotImplementedError, match="CompactionOptions"): table.optimize(compaction_options=options) diff --git a/python/src/table.rs b/python/src/table.rs index 6371d94c0..db3a88f7f 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1838,6 +1838,10 @@ impl Table { compaction_options: Option<&Bound<'py, PyDict>>, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); + let target_specified = compaction_options + .map(|options| options.contains("target_rows_per_fragment")) + .transpose()? + .unwrap_or(false); let compaction_options = parse_compaction_options(compaction_options)?; let older_than = if let Some(ms) = cleanup_since_ms { if ms > i64::MAX as u64 { @@ -1852,10 +1856,22 @@ impl Table { None }; future_into_py(self_.py(), async move { - let stats = inner - .optimize_full(compaction_options, older_than, None, delete_unverified) - .await - .infer_error()?; + let stats = if inner.as_native().is_some() { + inner + .optimize_full(compaction_options, older_than, None, delete_unverified) + .await + } else { + let mut overrides = + lancedb::table::optimize::CompactionOverrides::try_from(compaction_options) + .infer_error()?; + if !target_specified { + overrides.target_rows_per_fragment = None; + } + inner + .optimize_full_with_overrides(overrides, older_than, None, delete_unverified) + .await + } + .infer_error()?; let compaction_stats = stats.compaction.unwrap(); let prune_stats = stats.prune.unwrap(); Ok(OptimizeStats { diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 1211263d8..aa41f2c86 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -2318,7 +2318,7 @@ async fn remote_optimize_submits_a_compaction_job_over_sql() { assert!(stats.prune.is_none()); assert_eq!( *statements.lock().unwrap(), - ["OPTIMIZE TABLE \"events\" REWRITE DATA"] + ["OPTIMIZE \"events\" REWRITE DATA"] ); assert!( headers @@ -2362,7 +2362,7 @@ async fn remote_optimize_on_a_branch_submits_a_with_clause_over_sql() { assert_eq!( *statements.lock().unwrap(), - ["OPTIMIZE TABLE \"events\" REWRITE DATA WITH (branch = 'o''brien')"] + ["OPTIMIZE \"events\" REWRITE DATA WITH (branch = 'o''brien')"] ); shutdown.send(()).unwrap(); } @@ -2426,7 +2426,7 @@ async fn remote_optimize_full_skips_prune_and_index() { ); assert_eq!( *statements.lock().unwrap(), - ["OPTIMIZE TABLE \"events\" REWRITE DATA"] + ["OPTIMIZE \"events\" REWRITE DATA"] ); shutdown.send(()).unwrap(); } @@ -3612,3 +3612,109 @@ fn sql_status_reports_structured_resource_identity() { assert_eq!(resource.name, "missing"); assert_eq!(resource.namespace_path, ["ns"]); } + +#[tokio::test] +async fn remote_optimize_forwards_per_job_options_on_main_and_branch() { + for branch in [None, Some("o'brien")] { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let client = client_with_handler(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()); + }); + let table = RemoteTable::new_with_sql_client( + client, + "analytics".into(), + "events".into(), + vec!["team".into()], + "team$events".into(), + Default::default(), + Some(sql_client), + ) + .with_branch(branch.map(str::to_string)); + table + .optimize(crate::table::optimize::OptimizeAction::Compact { + options: crate::table::optimize::CompactionOptions { + target_rows_per_fragment: 100, + max_source_rows: Some(500), + io_buffer_size: Some(0), + batch_size: Some(16), + ..Default::default() + }, + remap_options: None, + }) + .await + .unwrap(); + let prefix = if branch.is_some() { + "branch = 'o''brien', " + } else { + "" + }; + assert_eq!( + *statements.lock().unwrap(), + [format!( + "OPTIMIZE \"events\" REWRITE DATA WITH ({prefix}target_rows_per_fragment = 100, max_source_rows = 500, io_buffer_size = 0, batch_size = 16)" + )] + ); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|h| h.namespace_path == "team") + ); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn remote_compact_preserves_omissions_and_explicit_default_values() { + use crate::table::{CompactionOptions, CompactionOverrides}; + let target = CompactionOptions::default().target_rows_per_fragment; + for (options, suffix) in [ + (CompactionOverrides::default(), String::new()), + ( + CompactionOverrides { + target_rows_per_fragment: Some(target), + ..Default::default() + }, + format!(" WITH (target_rows_per_fragment = {target})"), + ), + ( + CompactionOverrides { + io_buffer_size: Some(0), + ..Default::default() + }, + " WITH (io_buffer_size = 0)".into(), + ), + ] { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let statements = service.statements.clone(); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let client = client_with_handler(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()); + }); + let table = RemoteTable::new_with_sql_client( + client, + "analytics".into(), + "events".into(), + vec![], + "events".into(), + Default::default(), + Some(sql_client), + ); + table.compact(options).await.unwrap(); + assert_eq!( + *statements.lock().unwrap(), + [format!("OPTIMIZE \"events\" REWRITE DATA{suffix}")] + ); + shutdown.send(()).unwrap(); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b781531fe..3801405eb 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3,6 +3,7 @@ pub mod blobs; pub mod insert; +mod optimize; use self::insert::{RemoteWriteExec, WriteOp}; use super::client::RequestResultExt; @@ -3662,9 +3663,9 @@ impl BaseTable for RemoteTable { Ok(()) } - /// Only `OptimizeAction::Compact` with default `options` and no custom + /// Only `OptimizeAction::Compact` with supported `options` and no custom /// `remap_options` is supported remotely: it submits a durable - /// compaction job through the server's `OPTIMIZE TABLE ... REWRITE DATA` + /// compaction job through the server's `OPTIMIZE ... REWRITE DATA` /// SQL statement, the same atomic, single-worker, index-safe path the /// server uses for its own automatic compaction. `.wait()` below waits /// for that SQL statement's own execution to finish, which only confirms @@ -3680,11 +3681,12 @@ impl BaseTable for RemoteTable { /// optimization. A branch-scoped table handle submits a /// `WITH (branch = '')` clause; the statement does not itself /// validate the branch exists, so a nonexistent branch surfaces as an - /// async job failure rather than an error from this call. A non-default - /// `CompactionOptions` (for example `target_rows_per_fragment` or - /// `max_source_bytes`) is rejected: the SQL statement has no `WITH` - /// clause to carry it yet, so accepting it would silently ignore - /// settings a caller may be relying on to bound the job. + /// async job failure rather than an error from this call. Per-job + /// `target_rows_per_fragment`, `max_source_rows`, `io_buffer_size`, and + /// `batch_size` are forwarded in the SQL `WITH` clause. Other non-default + /// options are rejected before submission. Servers must support these + /// SQL options; an older server's error is returned without retrying + /// with the settings removed. async fn optimize(&self, action: OptimizeAction) -> Result { self.check_mutable().await?; let OptimizeAction::Compact { @@ -3707,26 +3709,34 @@ impl BaseTable for RemoteTable { .into(), }); } - if *options != lance::dataset::optimize::CompactionOptions::default() { - return Err(Error::NotSupported { - message: "custom CompactionOptions are not supported in LanceDB \ - Enterprise; the OPTIMIZE TABLE ... REWRITE DATA SQL \ - statement has no WITH clause to carry them yet, so a \ - non-default value would be silently ignored rather than \ - applied. Omit options to use the server's managed \ - compaction settings." - .into(), - }); + let mut overrides = crate::table::optimize::CompactionOverrides::try_from(options.clone())?; + // The legacy Lance options type has no presence bit for this field. + // Use Table::compact to explicitly send a target equal to Lance's default. + if options.target_rows_per_fragment + == crate::table::CompactionOptions::default().target_rows_per_fragment + { + overrides.target_rows_per_fragment = None; } + self.compact(overrides).await + } + + async fn compact( + &self, + options: crate::table::optimize::CompactionOverrides, + ) -> Result { + self.check_mutable().await?; + let compaction_options = optimize::compaction_sql_options(&options)?; let sql_client = self.sql_client.clone().ok_or_else(|| Error::NotSupported { message: "SQL is unavailable for this remote table client".to_string(), })?; - let mut statement = format!( - "OPTIMIZE TABLE {} REWRITE DATA", - quote_sql_identifier(&self.name) - ); + let mut statement = format!("OPTIMIZE {} REWRITE DATA", quote_sql_identifier(&self.name)); + let mut sql_options = Vec::new(); if let Some(branch) = &self.branch { - statement.push_str(&format!(" WITH (branch = {})", quote_sql_string(branch))); + sql_options.push(format!("branch = {}", quote_sql_string(branch))); + } + sql_options.extend(compaction_options); + if !sql_options.is_empty() { + statement.push_str(&format!(" WITH ({})", sql_options.join(", "))); } // `OptimizeStats` carries `lance::dataset::optimize` metric types that // are not `Serialize`/`Deserialize`/`Clone`, which `Job` requires, @@ -15863,7 +15873,7 @@ mod tests { } #[tokio::test] - async fn test_optimize_rejects_non_default_compaction_options() { + async fn test_optimize_rejects_unsupported_compaction_options() { let table = Table::new_with_handler("my_table", move |_| -> http::Response { panic!("custom compaction options must be rejected before any request is sent") }); diff --git a/rust/lancedb/src/remote/table/optimize.rs b/rust/lancedb/src/remote/table/optimize.rs new file mode 100644 index 000000000..2a42dd3ab --- /dev/null +++ b/rust/lancedb/src/remote/table/optimize.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Compaction options supported by the remote SQL submission protocol. +use crate::error::{Error, Result}; + +pub(super) fn compaction_sql_options( + options: &crate::table::optimize::CompactionOverrides, +) -> Result> { + for (name, value) in [ + ("target_rows_per_fragment", options.target_rows_per_fragment), + ("max_source_rows", options.max_source_rows), + ("batch_size", options.batch_size), + ] { + if value == Some(0) { + return Err(Error::InvalidInput { + message: format!("{name} must be greater than 0"), + }); + } + } + let mut result = Vec::new(); + for (name, value) in [ + ( + "target_rows_per_fragment", + options.target_rows_per_fragment.map(|v| v as u64), + ), + ("max_source_rows", options.max_source_rows.map(|v| v as u64)), + ("io_buffer_size", options.io_buffer_size), + ("batch_size", options.batch_size.map(|v| v as u64)), + ] { + if let Some(value) = value { + result.push(format!("{name} = {value}")); + } + } + Ok(result) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::table::{CompactionOptions, CompactionOverrides}; + + #[test] + fn omissions_and_explicit_default_target_are_distinct() { + assert!( + compaction_sql_options(&CompactionOverrides::default()) + .unwrap() + .is_empty() + ); + let target = CompactionOptions::default().target_rows_per_fragment; + assert_eq!( + compaction_sql_options(&CompactionOverrides { + target_rows_per_fragment: Some(target), + io_buffer_size: Some(0), + ..Default::default() + }) + .unwrap(), + [ + format!("target_rows_per_fragment = {target}"), + "io_buffer_size = 0".into() + ] + ); + } + + #[test] + fn unsupported_settings_are_rejected() { + assert!(matches!( + CompactionOverrides::try_from(CompactionOptions { + max_source_bytes: Some(1024), + ..Default::default() + }), + Err(Error::NotSupported { .. }) + )); + } + + #[test] + fn invalid_row_settings_are_rejected() { + for options in [ + CompactionOverrides { + target_rows_per_fragment: Some(0), + ..Default::default() + }, + CompactionOverrides { + max_source_rows: Some(0), + ..Default::default() + }, + CompactionOverrides { + batch_size: Some(0), + ..Default::default() + }, + ] { + assert!(matches!( + compaction_sql_options(&options), + Err(Error::InvalidInput { .. }) + )); + } + } +} diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index aefc0ac91..4691d33f1 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -111,7 +111,8 @@ pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTa pub use lance::dataset::scanner::DatasetRecordBatchStream; pub use lance_index::optimize::OptimizeOptions; pub use optimize::{ - CompactionMode, CompactionOptions, IndexRemapMode, OptimizeAction, OptimizeStats, + CompactionMode, CompactionOptions, CompactionOverrides, IndexRemapMode, OptimizeAction, + OptimizeStats, }; pub use refresh::{RefreshColumnOptions, RefreshColumnResult}; pub use rte_stats::{MemtableStats, RTEStats, SsTableStats, TableShardStats}; @@ -842,6 +843,14 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn tags(&self) -> Result>; /// Optimize the dataset. async fn optimize(&self, action: OptimizeAction) -> Result; + /// Compact with optional fields, preserving remote policy inheritance. + async fn compact(&self, options: optimize::CompactionOverrides) -> Result { + self.optimize(OptimizeAction::Compact { + options: options.into_options(), + remap_options: None, + }) + .await + } /// Add columns to the table. async fn add_columns( &self, @@ -1925,14 +1934,62 @@ impl Table { /// /// # LanceDB Enterprise (remote tables) /// - /// Only [`OptimizeAction::Compact`] with default `options` and no custom - /// `remap_options` is supported; `All`, `Prune`, and `Index` return + /// [`OptimizeAction::Compact`] supports per-job `target_rows_per_fragment`, + /// `max_source_rows`, `io_buffer_size`, and `batch_size`, with no custom + /// `remap_options`. Other non-default options, `All`, `Prune`, and `Index` return /// [`crate::Error::NotSupported`]. See [`OptimizeAction::Compact`] for /// what a remote call submits and what it does and does not wait for. pub async fn optimize(&self, action: OptimizeAction) -> Result { self.inner.optimize(action).await } + /// Compact data files with optional tuning fields. On remote tables, + /// omitted fields inherit table tuning, then deployment/built-in defaults. + /// Returns after job submission; remote statistics are placeholders. + /// + /// ``` + /// use lancedb::{Table, table::optimize::CompactionOverrides}; + /// # async fn compact(table: &Table) -> lancedb::Result<()> { + /// table.compact(CompactionOverrides { + /// io_buffer_size: Some(64 * 1024 * 1024), + /// ..Default::default() + /// }).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn compact(&self, options: optimize::CompactionOverrides) -> Result { + self.inner.compact(options).await + } + + /// Language-binding entry point retaining optional compaction fields. + #[doc(hidden)] + pub async fn optimize_full_with_overrides( + &self, + options: optimize::CompactionOverrides, + older_than: Option, + before_timestamp: Option>, + delete_unverified: Option, + ) -> Result { + if self.as_native().is_some() { + return self + .optimize_full( + options.into_options(), + older_than, + before_timestamp, + delete_unverified, + ) + .await; + } + if older_than.is_some() || before_timestamp.is_some() || delete_unverified == Some(true) { + return Err(Error::NotSupported { message: "pruning is not supported in LanceDB Enterprise; omit pruning options to run compaction only".into() }); + } + let stats = self.compact(options).await?; + Ok(OptimizeStats { + compaction: stats.compaction, + prune: Some(RemovalStats::default()), + }) + } + /// Prune versions committed before an absolute timestamp. /// /// This is an internal entry point for language bindings whose public API @@ -1966,10 +2023,10 @@ impl Table { /// /// This is an internal entry point for language bindings. On a remote /// table (LanceDB Enterprise) only [`OptimizeAction::Compact`] with its - /// defaults is supported. A caller asking for anything pruning-related + /// supported per-job options is supported. A caller asking for anything pruning-related /// (`older_than`, `before_timestamp`, or `delete_unverified: Some(true)`) - /// gets `NotSupported` before anything is submitted, the same way a - /// non-default `CompactionOptions` is already rejected -- rather than a + /// gets `NotSupported` before anything is submitted, the same way + /// unsupported `CompactionOptions` are rejected -- rather than a /// success that silently reports zero bytes pruned when nothing was /// pruned at all. A plain call with none of those set only compacts: /// `Prune` and `Index` are skipped here rather than attempted and diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index a5346afef..ab34e05c6 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -184,6 +184,53 @@ async fn validate_sq_index_remapping( Ok(plan) } +/// Optional compaction tuning. Remote tables inherit omitted values from table +/// policy, then deployment/built-in defaults. Explicit values, including a zero +/// I/O budget, take precedence. Local tables use Lance defaults for omissions. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CompactionOverrides { + pub target_rows_per_fragment: Option, + pub max_source_rows: Option, + /// Read-ahead budget in bytes; zero disables byte-based I/O backpressure. + pub io_buffer_size: Option, + pub batch_size: Option, +} + +impl CompactionOverrides { + pub(crate) fn into_options(self) -> CompactionOptions { + let mut options = CompactionOptions::default(); + if let Some(target) = self.target_rows_per_fragment { + options.target_rows_per_fragment = target; + } + options.max_source_rows = self.max_source_rows; + options.io_buffer_size = self.io_buffer_size; + options.batch_size = self.batch_size; + options + } +} + +impl TryFrom for CompactionOverrides { + type Error = Error; + + fn try_from(options: CompactionOptions) -> Result { + let mut remaining = options.clone(); + let defaults = CompactionOptions::default(); + remaining.target_rows_per_fragment = defaults.target_rows_per_fragment; + remaining.max_source_rows = defaults.max_source_rows; + remaining.io_buffer_size = defaults.io_buffer_size; + remaining.batch_size = defaults.batch_size; + if remaining != defaults { + return Err(Error::NotSupported { message: "remote CompactionOptions support only target_rows_per_fragment, max_source_rows, io_buffer_size, and batch_size; other options must retain their defaults".into() }); + } + Ok(Self { + target_rows_per_fragment: Some(options.target_rows_per_fragment), + max_source_rows: options.max_source_rows, + io_buffer_size: options.io_buffer_size, + batch_size: options.batch_size, + }) + } +} + /// Optimize the dataset. /// /// Similar to `VACUUM` in PostgreSQL, it offers different options to @@ -209,9 +256,24 @@ pub enum OptimizeAction { /// /// # LanceDB Enterprise (remote tables) /// - /// A remote table only supports `Compact` with default `options` and no - /// custom `remap_options`; both are rejected with - /// [`crate::Error::NotSupported`] otherwise. It submits a durable + /// A remote table supports per-job `target_rows_per_fragment`, + /// `max_source_rows`, `io_buffer_size`, and `batch_size`. Other non-default + /// options and custom `remap_options` return [`crate::Error::NotSupported`] + /// before submission. The server must support these SQL options; errors + /// from older servers are returned without dropping the settings. + /// The source-row budget admits whole tasks and may perform no work when + /// the first task exceeds it. Zero I/O buffer size disables byte-based I/O + /// backpressure. Use a positive size to limit I/O read-ahead; this is not a + /// total-memory limit. Row targets, source-row budgets, and batch sizes must + /// be positive. + /// The legacy Lance options type cannot distinguish an omitted fragment target + /// from its default value; that default is treated as unspecified remotely. + /// Use [`crate::Table::compact`] with [`CompactionOverrides`] to explicitly + /// override a table target with the Lance default value. + /// Omitted options inherit table tuning, then deployment/built-in defaults; options + /// affect only this job, without changing ongoing table policy. + /// + /// It submits a durable /// compaction job to the server over SQL and waits for that submission to /// be confirmed, not for compaction itself to finish: the job then runs /// asynchronously on the server, so this call can return well before the @@ -219,6 +281,23 @@ pub enum OptimizeAction { /// [`OptimizeStats::compaction`] is a zeroed [`CompactionMetrics`] /// placeholder rather than the real counts a local table reports. /// + /// ``` + /// use lancedb::{Table, table::{CompactionOptions, OptimizeAction}}; + /// # async fn compact(table: &Table) -> lancedb::Result<()> { + /// table.optimize(OptimizeAction::Compact { + /// options: CompactionOptions { + /// target_rows_per_fragment: 100_000, + /// max_source_rows: Some(1_000_000), + /// io_buffer_size: Some(64 * 1024 * 1024), + /// batch_size: Some(1024), + /// ..Default::default() + /// }, + /// remap_options: None, + /// }).await?; + /// # Ok(()) + /// # } + /// ``` + /// /// A branch-scoped table handle compacts that branch. The branch is not /// validated to exist before the job is submitted, so a nonexistent /// branch fails the compaction job asynchronously rather than failing