mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat: support remote table compaction via SQL OPTIMIZE TABLE statement (#4802)
`Table::optimize()` previously returned `NotSupported` for remote tables. This PR wires it to the LanceDB Enterprise server's `OPTIMIZE TABLE ... REWRITE DATA` SQL statement, submitted over Flight SQL. Only `OptimizeAction::Compact` without a custom `remap_options` or non-default `CompactionOptions` is supported remotely. `OptimizeAction::All`, `::Prune`, and `::Index` all return a clear `NotSupported` error instead of silently doing the wrong thing. A branch-scoped table handle submits a `WITH (branch = '<branch>')` clause; the statement doesn't validate the branch exists itself, so an unknown branch surfaces as an async job failure rather than an error from this call. The server's SQL job adapter doesn't expose compaction metrics, so `OptimizeStats::compaction` is `Some(CompactionMetrics::default())`, a placeholder rather than real numbers — it must still be `Some`, since callers (including the Python/Node bindings) unwrap it unconditionally on success. This mirrors how `refresh_materialized_view_async` already handles the same limitation for its own metrics. The Python and Node `optimize()` convenience methods call a new shared `Table::optimize_full`, which runs `Compact` and then, on a remote table, skips `Prune` and `Index` instead of attempting and failing on them — returning zeroed prune stats rather than submitting a compaction job and then raising. Python's `RemoteTable.optimize()` also now forwards `compaction_options` to the Rust layer, so a non-default value is rejected before submission instead of being silently dropped. --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 5
parent
dcda5224ee
commit
a9ad499781
@@ -1083,12 +1083,21 @@ class RemoteTable(Table):
|
||||
retrain: bool = False,
|
||||
compaction_options: Optional[CompactionOptions] = None,
|
||||
):
|
||||
"""optimize() is not supported in LanceDB Enterprise."""
|
||||
"""Compact the table's data files via a durable, asynchronous
|
||||
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.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._table.optimize(
|
||||
cleanup_older_than=cleanup_older_than,
|
||||
delete_unverified=delete_unverified,
|
||||
retrain=retrain,
|
||||
compaction_options=compaction_options,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -2935,6 +2935,17 @@ class Table(ABC):
|
||||
"""
|
||||
Optimize the on-disk data and indices for better performance.
|
||||
|
||||
.. note::
|
||||
|
||||
On a LanceDB Enterprise (remote) table, only compaction is
|
||||
currently supported. Pruning old versions and index optimization
|
||||
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.
|
||||
|
||||
Modeled after ``VACUUM`` in PostgreSQL.
|
||||
|
||||
Optimization covers three operations:
|
||||
@@ -8195,6 +8206,17 @@ class AsyncTable:
|
||||
"""
|
||||
Optimize the on-disk data and indices for better performance.
|
||||
|
||||
.. note::
|
||||
|
||||
On a LanceDB Enterprise (remote) table, only compaction is
|
||||
currently supported. Pruning old versions and index optimization
|
||||
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.
|
||||
|
||||
Modeled after ``VACUUM`` in PostgreSQL.
|
||||
|
||||
Optimization covers three operations:
|
||||
|
||||
@@ -2494,7 +2494,6 @@ def remote_unsupported_operations():
|
||||
return [
|
||||
("uses_v2_manifest_paths", (), {}, "uses_v2_manifest_paths"),
|
||||
("migrate_v2_manifest_paths", (), {}, "migrate_manifest_paths_v2"),
|
||||
("optimize", (), {}, "optimize"),
|
||||
("create_index", ("vector",), {"config": HnswPq()}, "IVF_HNSW_PQ"),
|
||||
]
|
||||
|
||||
@@ -5626,3 +5625,134 @@ def test_local_connections_refuse_property_graphs(tmp_path):
|
||||
list(db.list_property_graphs())
|
||||
with pytest.raises(NotImplementedError, match=refused):
|
||||
db.drop_property_graph("social")
|
||||
|
||||
|
||||
def _create_table_handler(request):
|
||||
if request.path == "/v1/table/test/create/?mode=create":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b"{}")
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
|
||||
@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)."""
|
||||
with http.server.ThreadingHTTPServer(
|
||||
("127.0.0.1", 0), make_mock_http_handler(_create_table_handler)
|
||||
) as rest_server:
|
||||
rest_port = rest_server.server_address[1]
|
||||
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}])
|
||||
finally:
|
||||
rest_server.shutdown()
|
||||
handle.join()
|
||||
|
||||
|
||||
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
|
||||
request pruning, which is not supported remotely. Retrying that failure
|
||||
resubmitted another compaction job. It should instead succeed, with
|
||||
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
|
||||
`do_exchange`: the client sends a ticket wrapping the statement and reads
|
||||
back whatever the server returns as the job's result stream.
|
||||
"""
|
||||
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_table_with_sql(server) as table:
|
||||
stats = table.optimize()
|
||||
|
||||
assert len(server.tickets) == 1
|
||||
assert b"OPTIMIZE TABLE" 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():
|
||||
"""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
|
||||
unsupported options never saw it. It should now reach that guard and
|
||||
raise before any SQL statement is submitted."""
|
||||
from lancedb.table import CompactionOptions
|
||||
|
||||
flight = pytest.importorskip("pyarrow.flight")
|
||||
|
||||
class UnreachableServer(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 = UnreachableServer()
|
||||
options: CompactionOptions = {"target_rows_per_fragment": 5}
|
||||
with _mock_remote_table_with_sql(server) as table:
|
||||
with pytest.raises(NotImplementedError, match="CompactionOptions"):
|
||||
table.optimize(compaction_options=options)
|
||||
|
||||
assert server.tickets == []
|
||||
|
||||
|
||||
def test_remote_optimize_rejects_explicit_pruning_before_submitting():
|
||||
"""Regression test: a remote table's `optimize()` used to silently drop
|
||||
an explicit pruning request (`cleanup_older_than`), returning success
|
||||
with zero bytes removed even though nothing was pruned. It should now
|
||||
raise before any SQL statement is submitted, the same way a non-default
|
||||
`compaction_options` does."""
|
||||
from datetime import timedelta
|
||||
|
||||
flight = pytest.importorskip("pyarrow.flight")
|
||||
|
||||
class UnreachableServer(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 = UnreachableServer()
|
||||
with _mock_remote_table_with_sql(server) as table:
|
||||
with pytest.raises(NotImplementedError, match="pruning"):
|
||||
table.optimize(cleanup_older_than=timedelta(days=1))
|
||||
with pytest.raises(NotImplementedError, match="pruning"):
|
||||
table.optimize(delete_unverified=True)
|
||||
|
||||
assert server.tickets == []
|
||||
|
||||
@@ -192,7 +192,10 @@ def test_remote_optimize_options():
|
||||
with pytest.raises(NotImplementedError, match="remote optimize"):
|
||||
table.optimize(cleanup_older_than=age, delete_unverified=True, retrain=True)
|
||||
handle.optimize.assert_awaited_once_with(
|
||||
cleanup_older_than=age, delete_unverified=True, retrain=True
|
||||
cleanup_older_than=age,
|
||||
delete_unverified=True,
|
||||
retrain=True,
|
||||
compaction_options=None,
|
||||
)
|
||||
|
||||
|
||||
|
||||
+6
-25
@@ -29,8 +29,8 @@ use lancedb::connection::LanceFileVersion;
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, CompactionMode, CompactionOptions, Duration,
|
||||
FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform,
|
||||
OptimizeAction, OptimizeOptions, Ref, RefreshColumnOptions, Table as LanceDbTable,
|
||||
FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform, Ref,
|
||||
RefreshColumnOptions, Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
@@ -1783,31 +1783,12 @@ impl Table {
|
||||
None
|
||||
};
|
||||
future_into_py(self_.py(), async move {
|
||||
let compaction_stats = inner
|
||||
.optimize(OptimizeAction::Compact {
|
||||
options: compaction_options,
|
||||
remap_options: None,
|
||||
})
|
||||
.await
|
||||
.infer_error()?
|
||||
.compaction
|
||||
.unwrap();
|
||||
let prune_stats = inner
|
||||
.optimize(OptimizeAction::Prune {
|
||||
older_than,
|
||||
delete_unverified,
|
||||
error_if_tagged_old_versions: None,
|
||||
})
|
||||
.await
|
||||
.infer_error()?
|
||||
.prune
|
||||
.unwrap();
|
||||
inner
|
||||
.optimize(lancedb::table::OptimizeAction::Index(
|
||||
OptimizeOptions::default(),
|
||||
))
|
||||
let stats = inner
|
||||
.optimize_full(compaction_options, older_than, None, delete_unverified)
|
||||
.await
|
||||
.infer_error()?;
|
||||
let compaction_stats = stats.compaction.unwrap();
|
||||
let prune_stats = stats.prune.unwrap();
|
||||
Ok(OptimizeStats {
|
||||
compaction: CompactionStats {
|
||||
files_added: compaction_stats.files_added as u64,
|
||||
|
||||
Reference in New Issue
Block a user