feat: support remote compaction tuning and table inheritance (#4886)

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(

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.
 

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)

python/src/table.rs

@@ -1838,6 +1838,10 @@ impl Table {
         compaction_options: Option<&Bound<'py, PyDict>>,
     ) -> PyResult<Bound<'py, PyAny>> {
         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 {