Compare commits

...

2 Commits

Author SHA1 Message Date
geruh cfcbcfbc92 feat: list_bases returns registered table storage prefixes
Return the additional storage bases for the current table snapshot
on native, memory, namespace, and Cloud clients.
2026-08-21 00:48:38 -07:00
geruh 047f431837 feat: add_bases registers extra table storage prefixes
TableBase plus add_bases on native, memory, namespace, and Cloud
clients.
2026-08-21 00:48:38 -07:00
17 changed files with 1052 additions and 4 deletions
+107
View File
@@ -1001,4 +1001,111 @@ describe("remote connection jobs surface", () => {
},
);
});
it("addBases posts the bases array", async () => {
const postedBodies: unknown[] = [];
await withMockDatabase(
(req, res) => {
const path = req.url ?? "";
if (path.endsWith("/describe/")) {
res.writeHead(200, { "Content-Type": "application/json" }).end(
JSON.stringify({
name: "photos",
version: 1,
schema: { fields: [] },
}),
);
return;
}
if (path.endsWith("/bases/")) {
const chunks: Buffer[] = [];
req.on("data", (chunk) => chunks.push(chunk));
req.on("end", () => {
postedBodies.push(JSON.parse(Buffer.concat(chunks).toString()));
res
.writeHead(200, { "Content-Type": "application/json" })
.end(JSON.stringify({ version: 2 }));
});
return;
}
if (path.endsWith("/bases/list/")) {
res.writeHead(200, { "Content-Type": "application/json" }).end(
JSON.stringify({
bases: [
{
path: "s3://bucket/media/",
isDatasetRoot: false,
},
],
}),
);
return;
}
res.writeHead(404).end();
},
async (db) => {
const table = await db.openTable("photos");
await table.addBases({ path: "s3://bucket/media/" });
expect(await table.listBases()).toEqual([
{
path: "s3://bucket/media/",
isDatasetRoot: false,
},
]);
},
);
expect(postedBodies).toEqual([
{
bases: [
{
path: "s3://bucket/media/",
isDatasetRoot: false,
},
],
},
]);
});
it("listBases returns a named dataset-root base", async () => {
await withMockDatabase(
(req, res) => {
const path = req.url ?? "";
if (path.endsWith("/describe/")) {
res.writeHead(200, { "Content-Type": "application/json" }).end(
JSON.stringify({
name: "photos",
version: 1,
schema: { fields: [] },
}),
);
return;
}
if (path.endsWith("/bases/list/")) {
res.writeHead(200, { "Content-Type": "application/json" }).end(
JSON.stringify({
bases: [
{
path: "s3://bucket/archive/",
name: "archive",
isDatasetRoot: true,
},
],
}),
);
return;
}
res.writeHead(404).end();
},
async (db) => {
const table = await db.openTable("photos");
expect(await table.listBases()).toEqual([
{
path: "s3://bucket/archive/",
name: "archive",
isDatasetRoot: true,
},
]);
},
);
});
});
+26
View File
@@ -4,6 +4,7 @@
import * as fs from "fs";
import * as path from "path";
import * as tmp from "tmp";
import { pathToFileURL } from "url";
import * as arrow15 from "apache-arrow-15";
import * as arrow16 from "apache-arrow-16";
@@ -3404,3 +3405,28 @@ describe("computed columns", () => {
expect(rows.map((r) => r.doubled).sort()).toEqual([10, 2]);
});
});
describe("table bases", () => {
let tmpDir: tmp.DirResult;
beforeEach(() => {
tmpDir = tmp.dirSync({ unsafeCleanup: true });
});
afterEach(() => tmpDir.removeCallback());
it("listBases reflects added bases", async () => {
const conn = await connect(tmpDir.name);
const table = await conn.createEmptyTable(
"photos",
new arrow.Schema([new arrow.Field("id", new arrow.Int64(), false)]),
);
const media = path.join(tmpDir.name, "media");
fs.mkdirSync(media);
const location = pathToFileURL(media).toString();
expect(await table.listBases()).toEqual([]);
await table.addBases(location);
expect(await table.listBases()).toEqual([
{ path: location, isDatasetRoot: false },
]);
});
});
+1
View File
@@ -130,6 +130,7 @@ export {
export {
Table,
TableBase,
Branches,
BranchColumnSummary,
BranchColumnChange,
+56
View File
@@ -78,6 +78,25 @@ export interface WriteProgress {
done: boolean;
}
/**
* An extra storage prefix registered on a table.
*
* `path` is an object-store URI. `name` is an optional alias. `isDatasetRoot`
* is true when `path` points to a Lance dataset root. When false, `path`
* points directly to the directory containing the referenced files.
*/
export interface TableBase {
/** Object store URI such as `s3://bucket/media/`. */
path: string;
/** Optional alias. */
name?: string;
/**
* True when `path` is a Lance dataset root. When false, `path` is the
* directory containing the referenced files.
*/
isDatasetRoot?: boolean;
}
/**
* Options for adding data to a table.
*/
@@ -563,6 +582,18 @@ export abstract class Table {
| { computed: AddColumnsSql[] },
): Promise<AddColumnsResult>;
/**
* Register additional storage bases for this table.
*
* A URI string is a non-root base with no alias.
*/
abstract addBases(
bases: string | TableBase | Array<string | TableBase>,
): Promise<void>;
/** Return the additional storage bases for the current table snapshot. */
abstract listBases(): Promise<TableBase[]>;
/**
* Fill the rows of a computed column that hold no value yet.
*
@@ -1196,6 +1227,16 @@ export class LocalTable extends Table {
throw new Error("Invalid input type for addColumns");
}
async addBases(
bases: string | TableBase | Array<string | TableBase>,
): Promise<void> {
await this.inner.addBases(normalizeBases(bases));
}
async listBases(): Promise<TableBase[]> {
return await this.inner.listBases();
}
async refreshColumn(column: string): Promise<RefreshColumnResult> {
return await this.inner.refreshColumn(column);
}
@@ -1396,6 +1437,21 @@ export class LocalTable extends Table {
}
}
function normalizeBases(
bases: string | TableBase | Array<string | TableBase>,
): TableBase[] {
const baseInputs = Array.isArray(bases) ? bases : [bases];
return baseInputs.map((base) =>
typeof base === "string"
? { path: base, isDatasetRoot: false }
: {
path: base.path,
name: base.name,
isDatasetRoot: base.isDatasetRoot ?? false,
},
);
}
/**
* A definition of a column alteration. The alteration changes the column at
* `path` to have the new name `name`, to be nullable if `nullable` is true,
+47
View File
@@ -10,6 +10,7 @@ use lancedb::table::{
AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration,
FieldMetadataUpdate as LanceFieldMetadataUpdate, FtsToken as LanceDbFtsToken,
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
TableBase as LanceTableBase,
};
use napi::bindgen_prelude::*;
use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode};
@@ -446,6 +447,30 @@ impl Table {
Ok(res.into())
}
#[napi(catch_unwind)]
pub async fn add_bases(&self, bases: Vec<TableBase>) -> napi::Result<()> {
self.inner_ref()?
.add_bases(bases.into_iter().map(|base| LanceTableBase {
path: base.path,
name: base.name,
is_dataset_root: base.is_dataset_root,
}))
.await
.default_error()
}
#[napi(catch_unwind)]
pub async fn list_bases(&self) -> napi::Result<Vec<TableBase>> {
Ok(self
.inner_ref()?
.list_bases()
.await
.default_error()?
.into_iter()
.map(TableBase::from)
.collect())
}
#[napi(catch_unwind)]
pub async fn drop_columns(&self, columns: Vec<String>) -> napi::Result<DropColumnsResult> {
let col_refs = columns.iter().map(String::as_str).collect::<Vec<_>>();
@@ -700,6 +725,28 @@ impl Table {
}
}
#[napi(object)]
/// An extra storage prefix registered on a table.
pub struct TableBase {
/// Object store URI such as `s3://bucket/media/`.
pub path: String,
/// Optional alias.
pub name: Option<String>,
/// True when `path` is a Lance dataset root. When false, `path` is the
/// directory containing the referenced files.
pub is_dataset_root: bool,
}
impl From<LanceTableBase> for TableBase {
fn from(base: LanceTableBase) -> Self {
Self {
path: base.path,
name: base.name,
is_dataset_root: base.is_dataset_root,
}
}
}
#[napi(object)]
/// A description of an index currently configured on a column
pub struct IndexConfig {
+2 -1
View File
@@ -21,7 +21,7 @@ from .remote.db import RemoteDBConnection
from .expr import Expr, col, lit, func
from .schema import blob, vector, BlobType
from .job import AsyncJob, Job
from .table import AsyncTable, Table
from .table import AsyncTable, Table, TableBase
from .types import BaseTokenizerType
from ._lancedb import Session
from .namespace import (
@@ -521,5 +521,6 @@ __all__ = [
"RemoteDBConnection",
"Session",
"Table",
"TableBase",
"__version__",
]
+2
View File
@@ -377,6 +377,8 @@ class Table:
def take_offsets(self, offsets: list[int]) -> TakeQuery: ...
def take_row_ids(self, row_ids: list[int]) -> TakeQuery: ...
async def blob_columns(self) -> list[str]: ...
async def add_bases(self, bases: list[Any]) -> None: ...
async def list_bases(self) -> list[tuple[str, Optional[str], bool]]: ...
async def fetch_blobs(
self, column: str, row_ids: list[int]
) -> pa.LargeBinaryArray: ...
+12 -1
View File
@@ -50,7 +50,7 @@ from lancedb.index import (
)
from lancedb.job import Job
from lancedb.remote.db import LOOP
from lancedb.table import IndexConfigType, KNOWN_METRICS
from lancedb.table import IndexConfigType, KNOWN_METRICS, TableBase
import pyarrow as pa
from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME
@@ -1082,6 +1082,17 @@ class RemoteTable(Table):
def blob_columns(self) -> list[str]:
return LOOP.run(self._table.blob_columns())
def add_bases(
self,
bases: Union[str, TableBase, Iterable[Union[str, TableBase]]],
) -> None:
"""Register additional storage bases for this table."""
LOOP.run(self._table.add_bases(bases))
def list_bases(self) -> list[TableBase]:
"""Return the additional storage bases for the current table snapshot."""
return LOOP.run(self._table.list_bases())
def fetch_blobs(
self, column: str, row_ids: Union[list[int], pa.Table]
) -> pa.LargeBinaryArray:
+85
View File
@@ -19,6 +19,7 @@ from typing import (
Iterable,
List,
Literal,
Mapping,
Optional,
Sequence,
Tuple,
@@ -710,6 +711,21 @@ def _normalize_progress(progress):
return progress, False
@dataclass
class TableBase:
"""An extra storage prefix registered on a table.
``path`` is an object-store URI. ``name`` is an optional alias.
``is_dataset_root`` is true when ``path`` points to a Lance dataset
root. When false, ``path`` points directly to the directory containing
the referenced files.
"""
path: str
name: Optional[str] = None
is_dataset_root: bool = False
class Table(ABC):
"""
A Table is a collection of Records in a LanceDB Database.
@@ -1568,6 +1584,22 @@ class Table(ABC):
def blob_columns(self) -> list[str]:
"""Names of the blob v2 columns declared on this table."""
def add_bases(
self,
bases: Union[str, TableBase, Iterable[Union[str, TableBase]]],
) -> None:
"""Register additional storage bases for this table.
A URI string is a non-root base with no alias::
table.add_bases("s3://bucket/media/")
"""
raise NotImplementedError
def list_bases(self) -> list[TableBase]:
"""Return the additional storage bases for the current table snapshot."""
raise NotImplementedError
@abstractmethod
def fetch_blobs(
self, column: str, row_ids: Union[list[int], pa.Table]
@@ -2414,6 +2446,16 @@ class LanceTable(Table):
def blob_columns(self) -> list[str]:
return LOOP.run(self._table.blob_columns())
def add_bases(
self,
bases: Union[str, TableBase, Iterable[Union[str, TableBase]]],
) -> None:
LOOP.run(self._table.add_bases(bases))
def list_bases(self) -> list[TableBase]:
"""Return the additional storage bases for the current table snapshot."""
return LOOP.run(self._table.list_bases())
def fetch_blobs(
self, column: str, row_ids: Union[list[int], pa.Table]
) -> pa.LargeBinaryArray:
@@ -6266,6 +6308,25 @@ class AsyncTable:
async def blob_columns(self) -> list[str]:
return await self._inner.blob_columns()
async def add_bases(
self,
bases: Union[str, TableBase, Iterable[Union[str, TableBase]]],
) -> None:
"""Register additional storage bases for this table.
A URI string is a non-root base with no alias::
await table.add_bases("s3://bucket/media/")
"""
await self._inner.add_bases(_normalize_bases(bases))
async def list_bases(self) -> list[TableBase]:
"""Return the additional storage bases for the current table snapshot."""
return [
TableBase(path=path, name=name, is_dataset_root=is_dataset_root)
for path, name, is_dataset_root in await self._inner.list_bases()
]
async def fetch_blobs(
self, column: str, row_ids: Union[list[int], pa.Table]
) -> pa.LargeBinaryArray:
@@ -6484,6 +6545,30 @@ class AsyncTable:
await self._inner.replace_field_metadata(field_name, new_metadata)
def _normalize_bases(
base_inputs: Union[str, TableBase, Iterable[Union[str, TableBase]]],
) -> list[TableBase]:
if isinstance(base_inputs, (str, TableBase)):
items: Iterable[Union[str, TableBase]] = [base_inputs]
elif isinstance(base_inputs, Mapping):
raise TypeError(
"Expected a URI string, TableBase, or an iterable of those values"
)
else:
items = base_inputs
normalized_bases: list[TableBase] = []
for base in items:
if isinstance(base, str):
normalized_bases.append(TableBase(path=base))
elif isinstance(base, TableBase):
normalized_bases.append(base)
else:
raise TypeError(
f"Expected a URI string or TableBase, got {type(base).__name__}"
)
return normalized_bases
@dataclass
class IndexStatistics:
"""
+97
View File
@@ -0,0 +1,97 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import pyarrow as pa
import pytest
import lancedb
def test_list_bases_reflects_added_bases(tmp_path):
media = tmp_path / "media"
media.mkdir()
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64())])
table = db.create_table("photos", schema=schema)
assert table.list_bases() == []
table.add_bases(media.as_uri())
assert table.list_bases() == [lancedb.TableBase(path=media.as_uri())]
def test_add_bases_accepts_two_unnamed_paths(tmp_path):
media = tmp_path / "media"
other = tmp_path / "other"
media.mkdir()
other.mkdir()
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64())])
table = db.create_table("photos", schema=schema)
table.add_bases([media.as_uri(), other.as_uri()])
assert table.list_bases() == [
lancedb.TableBase(path=media.as_uri()),
lancedb.TableBase(path=other.as_uri()),
]
def test_add_bases_records_name_and_dataset_root(tmp_path):
media = tmp_path / "media"
parent = tmp_path / "parent"
media.mkdir()
parent.mkdir()
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64())])
table = db.create_table("photos", schema=schema)
table.add_bases(
[
lancedb.TableBase(path=media.as_uri(), name="media", is_dataset_root=False),
lancedb.TableBase(
path=parent.as_uri(), name="parent", is_dataset_root=True
),
]
)
assert table.list_bases() == [
lancedb.TableBase(path=media.as_uri(), name="media", is_dataset_root=False),
lancedb.TableBase(path=parent.as_uri(), name="parent", is_dataset_root=True),
]
def test_add_bases_rejects_dict_input(tmp_path):
db = lancedb.connect(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64())])
table = db.create_table("photos", schema=schema)
with pytest.raises(TypeError, match="TableBase"):
table.add_bases({"path": "s3://bucket/media/"})
@pytest.mark.asyncio
async def test_async_add_bases_accepts_file_uri(tmp_path):
media = tmp_path / "media"
media.mkdir()
db = await lancedb.connect_async(tmp_path / "db")
schema = pa.schema([pa.field("id", pa.int64())])
table = await db.create_table("photos", schema=schema)
assert await table.list_bases() == []
await table.add_bases(media.as_uri())
assert await table.list_bases() == [lancedb.TableBase(path=media.as_uri())]
def test_memory_add_bases_accepts_file_uri(tmp_path):
media = tmp_path / "media"
media.mkdir()
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64())])
table = db.create_table("photos", schema=schema)
assert table.list_bases() == []
table.add_bases(media.as_uri())
assert table.list_bases() == [lancedb.TableBase(path=media.as_uri())]
def test_namespace_add_bases_accepts_file_uri(tmp_path):
media = tmp_path / "media"
media.mkdir()
db = lancedb.connect_namespace("dir", {"root": str(tmp_path / "ns")})
schema = pa.schema([pa.field("id", pa.int64())])
table = db.create_table("photos", schema=schema)
assert table.list_bases() == []
table.add_bases(media.as_uri())
assert table.list_bases() == [lancedb.TableBase(path=media.as_uri())]
+76
View File
@@ -2306,3 +2306,79 @@ def test_remote_connection_jobs_surface():
assert job.status() == "failed"
with pytest.raises(JobFailedError, match="worker died"):
job.wait(timeout=timedelta(seconds=5))
def test_remote_add_and_list_bases():
captured_body = {}
def handler(request):
if request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
elif request.path == "/v1/table/test/bases/":
content_len = int(request.headers.get("Content-Length", 0))
captured_body.update(json.loads(request.rfile.read(content_len)))
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 2}')
elif request.path == "/v1/table/test/bases/list/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
b'{"bases":[{"path":"s3://bucket/media/","isDatasetRoot":false}]}'
)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.open_table("test")
table.add_bases(lancedb.TableBase(path="s3://bucket/media/"))
assert table.list_bases() == [
lancedb.TableBase(
path="s3://bucket/media/",
name=None,
is_dataset_root=False,
)
]
assert captured_body["bases"] == [
{
"path": "s3://bucket/media/",
"isDatasetRoot": False,
}
]
def test_remote_list_bases_returns_named_dataset_root():
def handler(request):
if request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
elif request.path == "/v1/table/test/bases/list/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
b'{"bases":[{"path":"s3://bucket/archive/","name":"archive","isDatasetRoot":true}]}'
)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.open_table("test")
assert table.list_bases() == [
lancedb.TableBase(
path="s3://bucket/archive/",
name="archive",
is_dataset_root=True,
)
]
+38
View File
@@ -22,6 +22,7 @@ use lancedb::index::scalar::FtsIndexBuilder;
use lancedb::table::{
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
TableBase as LanceTableBase,
};
use lancedb::tokenize as lancedb_tokenize;
use pyo3::{
@@ -94,6 +95,13 @@ fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult
Ok(out.unbind())
}
#[derive(FromPyObject)]
pub(crate) struct PyTableBase {
path: String,
name: Option<String>,
is_dataset_root: bool,
}
#[derive(FromPyObject)]
enum PredicateArg {
Expr(PyExpr),
@@ -1238,6 +1246,36 @@ impl Table {
})
}
#[pyo3(signature = (bases))]
pub fn add_bases(
self_: PyRef<'_, Self>,
bases: Vec<PyTableBase>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
let bases: Vec<LanceTableBase> = bases
.into_iter()
.map(|base| LanceTableBase {
path: base.path,
name: base.name,
is_dataset_root: base.is_dataset_root,
})
.collect();
future_into_py(self_.py(), async move {
inner.add_bases(bases).await.infer_error()
})
}
pub fn list_bases(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let bases = inner.list_bases().await.infer_error()?;
Ok(bases
.into_iter()
.map(|base| (base.path, base.name, base.is_dataset_root))
.collect::<Vec<_>>())
})
}
/// Read blob bytes for `row_ids` from blob v2 column `column`.
#[pyo3(signature = (column, row_ids))]
pub fn fetch_blobs(
+1 -1
View File
@@ -214,7 +214,7 @@ use lance_linalg::distance::DistanceType as LanceDistanceType;
/// a built-in pull-based adapter.
#[cfg(feature = "metrics")]
pub use metrics;
pub use table::{FtsToken, Table};
pub use table::{FtsToken, Table, TableBase};
/// Tokenize a full-text search query using an explicit FTS tokenizer configuration.
///
+85
View File
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
pub mod bases;
pub mod blobs;
pub mod insert;
@@ -2244,6 +2245,14 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
self.blob_columns_impl().await
}
async fn add_bases(&self, bases: &[crate::table::TableBase]) -> Result<()> {
self.add_bases_impl(bases).await
}
async fn list_bases(&self) -> Result<Vec<crate::table::TableBase>> {
self.list_bases_impl().await
}
async fn fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result<LargeBinaryArray> {
self.fetch_blobs_impl(column, row_ids).await
}
@@ -4089,6 +4098,82 @@ mod tests {
.unwrap()
}
#[tokio::test]
async fn test_add_bases_posts_the_bases_array() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/bases/");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(
body["bases"],
serde_json::json!([{
"path": "s3://bucket/media/",
"isDatasetRoot": false
}])
);
http::Response::builder()
.status(200)
.body(r#"{"version": 4}"#)
.unwrap()
});
table.add_bases(["s3://bucket/media/"]).await.unwrap();
}
#[tokio::test]
async fn list_bases_posts_and_returns_the_bases() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/bases/list/");
http::Response::builder()
.status(200)
.body(r#"{"bases":[{"path":"s3://bucket/media/","isDatasetRoot":false}]}"#)
.unwrap()
});
let bases = table.list_bases().await.unwrap();
assert_eq!(
bases,
vec![crate::table::TableBase::from("s3://bucket/media/")]
);
}
#[tokio::test]
async fn list_bases_returns_named_dataset_root_entries() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/v1/table/my_table/bases/list/");
http::Response::builder()
.status(200)
.body(r#"{"bases":[{"path":"s3://bucket/archive/","name":"archive","isDatasetRoot":true}]}"#)
.unwrap()
});
let bases = table.list_bases().await.unwrap();
assert_eq!(
bases,
vec![crate::table::TableBase {
path: "s3://bucket/archive/".into(),
name: Some("archive".into()),
is_dataset_root: true,
}]
);
}
#[tokio::test]
async fn test_add_bases_rejects_empty_response() {
let table = Table::new_with_handler("my_table", |_request| {
http::Response::builder().status(200).body("").unwrap()
});
let err = table.add_bases(["s3://bucket/media/"]).await.unwrap_err();
assert!(
err.to_string()
.contains("invalid response while registering table bases"),
"{err}"
);
}
#[rstest]
#[case(semver::Version::new(0, 1, 0))]
#[case(semver::Version::new(0, 5, 0))]
+68
View File
@@ -0,0 +1,68 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Cloud HTTP for registering and listing extra table storage bases.
use serde::Deserialize;
use crate::Error;
use crate::error::Result;
use crate::remote::client::{HttpSend, RequestResultExt};
use super::RemoteTable;
#[derive(Debug, Deserialize)]
struct AddBasesResponse {
version: u64,
}
#[derive(Debug, Deserialize)]
struct ListBasesResponse {
bases: Vec<crate::table::TableBase>,
}
impl<S: HttpSend> RemoteTable<S> {
pub(super) async fn add_bases_impl(&self, bases: &[crate::table::TableBase]) -> Result<()> {
self.check_mutable().await?;
let mut body = serde_json::json!({ "bases": bases });
self.apply_branch_body(&mut body);
let request = self
.client
.post(&format!("/v1/table/{}/bases/", self.identifier))
.json(&body);
let (request_id, response) = self.send(request, true).await?;
let response = self.check_table_response(&request_id, response).await?;
let body = response.text().await.err_to_http(request_id.clone())?;
let parsed: AddBasesResponse = serde_json::from_str(&body).map_err(|e| Error::Http {
source: format!(
"The server returned an invalid response while registering table bases: {e}"
)
.into(),
request_id,
status_code: None,
})?;
self.track_write_version(parsed.version);
Ok(())
}
pub(super) async fn list_bases_impl(&self) -> Result<Vec<crate::table::TableBase>> {
let version = self.current_version().await;
let mut body = serde_json::json!({ "version": version });
self.apply_branch_body(&mut body);
let request = self
.post_read(&format!("/v1/table/{}/bases/list/", self.identifier))
.json(&body);
let (request_id, response) = self.send(request, true).await?;
let response = self.check_table_response(&request_id, response).await?;
let body = response.text().await.err_to_http(request_id.clone())?;
let parsed: ListBasesResponse = serde_json::from_str(&body).map_err(|e| Error::Http {
source: format!(
"The server returned an invalid response while listing table bases: {e}"
)
.into(),
request_id,
status_code: None,
})?;
Ok(parsed.bases)
}
}
+118 -1
View File
@@ -34,7 +34,7 @@ use lance_index::scalar::inverted::query::collect_query_tokens;
use lance_namespace::LanceNamespace;
use lance_namespace::error::NamespaceError;
use lance_namespace::models::DescribeTableRequest;
use lance_table::format::Manifest;
use lance_table::format::{BasePath, Manifest};
use lance_table::io::commit::CommitHandler;
use lance_table::io::commit::ManifestNamingScheme;
use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
@@ -711,6 +711,18 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
message: "blob_columns is not supported on this table type".into(),
})
}
/// Register additional storage bases for this table.
async fn add_bases(&self, _bases: &[TableBase]) -> Result<()> {
Err(Error::NotSupported {
message: "Registering table bases is not supported for this table type.".into(),
})
}
/// Return the additional storage bases for the current table snapshot.
async fn list_bases(&self) -> Result<Vec<TableBase>> {
Err(Error::NotSupported {
message: "Listing table bases is not supported for this table type.".into(),
})
}
/// Materialize blob bytes for the given row ids. See [`Table::fetch_blobs`].
async fn fetch_blobs(&self, _column: &str, _row_ids: &[u64]) -> Result<LargeBinaryArray> {
Err(Error::NotSupported {
@@ -889,6 +901,54 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
}
}
/// An extra storage prefix registered on a table.
///
/// `path` is an object-store URI. `name` is an optional alias. `is_dataset_root`
/// is true when `path` points to a Lance dataset root. When false, `path`
/// points directly to the directory containing the referenced files.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TableBase {
/// Object store URI such as `s3://bucket/media/`.
pub path: String,
/// Optional alias.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
/// True when `path` is a Lance dataset root. When false, `path` is the
/// directory containing the referenced files.
#[serde(default)]
pub is_dataset_root: bool,
}
impl TableBase {
/// A non-root base with no alias.
pub fn new(path: impl Into<String>) -> Self {
Self {
path: path.into(),
name: None,
is_dataset_root: false,
}
}
}
impl From<&str> for TableBase {
fn from(path: &str) -> Self {
Self::new(path)
}
}
impl From<&String> for TableBase {
fn from(path: &String) -> Self {
Self::new(path.as_str())
}
}
impl From<String> for TableBase {
fn from(path: String) -> Self {
Self::new(path)
}
}
/// A Table is a collection of strong typed Rows.
///
/// The type of the each row is defined in Apache Arrow [Schema].
@@ -1126,6 +1186,30 @@ impl Table {
self.inner.blob_columns().await
}
/// Register additional storage bases for this table.
///
/// A URI string is a non-root base with no alias.
///
/// ```
/// # use lancedb::Table;
/// # async fn register(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
/// table.add_bases(["s3://bucket/media/"]).await?;
/// # Ok(())
/// # }
/// ```
pub async fn add_bases(
&self,
bases: impl IntoIterator<Item = impl Into<TableBase>>,
) -> Result<()> {
let bases: Vec<TableBase> = bases.into_iter().map(Into::into).collect();
self.inner.add_bases(&bases).await
}
/// Return the additional storage bases for the current table snapshot.
pub async fn list_bases(&self) -> Result<Vec<TableBase>> {
self.inner.list_bases().await
}
/// Materialize blob bytes for the given row ids.
///
/// Output matches `row_ids` in length and order. Null blobs are null;
@@ -3364,6 +3448,39 @@ impl BaseTable for NativeTable {
Ok(crate::blob::blob_column_names(schema.as_ref()))
}
async fn add_bases(&self, bases: &[TableBase]) -> Result<()> {
self.dataset.ensure_mutable()?;
let dataset = self.dataset.get().await?;
let new_bases = bases
.iter()
.map(|base| {
BasePath::new(
0,
base.path.clone(),
base.name.clone(),
base.is_dataset_root,
)
})
.collect();
let dataset = dataset.add_bases(new_bases, None).await?;
self.dataset.update(dataset);
Ok(())
}
async fn list_bases(&self) -> Result<Vec<TableBase>> {
let dataset = self.dataset.get().await?;
let mut bases: Vec<&BasePath> = dataset.manifest().base_paths.values().collect();
bases.sort_by_key(|base| base.id);
Ok(bases
.into_iter()
.map(|base| TableBase {
path: base.path.clone(),
name: base.name.clone(),
is_dataset_root: base.is_dataset_root,
})
.collect())
}
async fn fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result<LargeBinaryArray> {
let dataset = self.dataset.get().await?;
crate::blob::take_blobs_aligned(&dataset, column, row_ids).await
+231
View File
@@ -0,0 +1,231 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Arc;
use arrow_array::{Int64Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use lance::dataset::{WriteMode, WriteParams};
use lancedb::{Result, TableBase, connect, connect_namespace, table::WriteOptions};
use tempfile::tempdir;
use url::Url;
fn empty_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]))
}
fn file_uri(path: &std::path::Path) -> String {
Url::from_file_path(path)
.unwrap_or_else(|_| panic!("not an absolute path: {}", path.display()))
.to_string()
}
#[tokio::test]
async fn test_list_bases_reflects_added_bases() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().join("db").to_str().unwrap())
.execute()
.await?;
let table = db.create_empty_table("t", empty_schema()).execute().await?;
assert!(table.list_bases().await?.is_empty());
let media = tmp.path().join("media");
std::fs::create_dir_all(&media).unwrap();
let media_uri = file_uri(&media);
table.add_bases([&media_uri]).await?;
assert_eq!(
table.list_bases().await?,
vec![TableBase {
path: media_uri,
name: None,
is_dataset_root: false,
}]
);
Ok(())
}
#[tokio::test]
async fn test_add_bases_accepts_two_unnamed_paths() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().join("db").to_str().unwrap())
.execute()
.await?;
let table = db.create_empty_table("t", empty_schema()).execute().await?;
let media = tmp.path().join("media");
let other = tmp.path().join("other");
std::fs::create_dir_all(&media).unwrap();
std::fs::create_dir_all(&other).unwrap();
let media_uri = file_uri(&media);
let other_uri = file_uri(&other);
table.add_bases([&media_uri, &other_uri]).await?;
assert_eq!(
table.list_bases().await?,
vec![
TableBase {
path: media_uri,
name: None,
is_dataset_root: false,
},
TableBase {
path: other_uri,
name: None,
is_dataset_root: false,
},
]
);
Ok(())
}
#[tokio::test]
async fn test_add_bases_write_and_read_through_registered_base() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().join("db").to_str().unwrap())
.execute()
.await?;
let table = db.create_empty_table("t", empty_schema()).execute().await?;
let media = tmp.path().join("media");
std::fs::create_dir_all(&media).unwrap();
let media_uri = file_uri(&media);
table.add_bases([&media_uri]).await?;
let batch = RecordBatch::try_new(
empty_schema(),
vec![Arc::new(Int64Array::from(vec![1, 2, 3]))],
)
.unwrap();
table
.add(batch)
.write_options(WriteOptions {
lance_write_params: Some(WriteParams {
mode: WriteMode::Append,
target_base_names_or_paths: Some(vec![media_uri.clone()]),
..Default::default()
}),
})
.execute()
.await?;
assert_eq!(table.count_rows(None).await?, 3);
let dataset = table.dataset().unwrap().get().await?;
let registered = dataset
.manifest()
.base_paths
.values()
.find(|base| base.path == media_uri)
.expect("registered base");
assert_ne!(registered.id, 0);
assert!(registered.name.is_none());
assert!(
dataset.get_fragments().iter().any(|fragment| {
fragment
.metadata()
.files
.iter()
.any(|file| file.base_id == Some(registered.id))
}),
"written fragment should reference the registered base"
);
assert!(
std::fs::read_dir(&media)
.unwrap()
.filter_map(|entry| entry.ok())
.any(|entry| entry.path().extension().is_some_and(|ext| ext == "lance")),
"data file should land under the registered base"
);
Ok(())
}
#[tokio::test]
async fn test_add_bases_accepts_named_and_dataset_root_entries() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect(tmp.path().join("db").to_str().unwrap())
.execute()
.await?;
let table = db.create_empty_table("t", empty_schema()).execute().await?;
let media = tmp.path().join("media");
let parent = tmp.path().join("parent");
std::fs::create_dir_all(&media).unwrap();
std::fs::create_dir_all(&parent).unwrap();
let media_uri = file_uri(&media);
let parent_uri = file_uri(&parent);
table
.add_bases([
TableBase {
path: media_uri.clone(),
name: Some("media".into()),
is_dataset_root: false,
},
TableBase {
path: parent_uri.clone(),
name: Some("parent".into()),
is_dataset_root: true,
},
])
.await?;
assert_eq!(
table.list_bases().await?,
vec![
TableBase {
path: media_uri,
name: Some("media".into()),
is_dataset_root: false,
},
TableBase {
path: parent_uri,
name: Some("parent".into()),
is_dataset_root: true,
},
]
);
Ok(())
}
#[tokio::test]
async fn test_memory_list_bases_reflects_added_bases() -> Result<()> {
let tmp = tempdir().unwrap();
let db = connect("memory://").execute().await?;
let table = db.create_empty_table("t", empty_schema()).execute().await?;
assert!(table.list_bases().await?.is_empty());
let media = tmp.path().join("media");
std::fs::create_dir_all(&media).unwrap();
let media_uri = file_uri(&media);
table.add_bases([&media_uri]).await?;
assert_eq!(
table.list_bases().await?,
vec![TableBase {
path: media_uri,
name: None,
is_dataset_root: false,
}]
);
Ok(())
}
#[tokio::test]
async fn test_namespace_list_bases_reflects_added_bases() -> Result<()> {
let tmp = tempdir().unwrap();
let mut properties = std::collections::HashMap::new();
properties.insert("root".to_string(), tmp.path().to_str().unwrap().to_string());
let db = connect_namespace("dir", properties).execute().await?;
let table = db.create_empty_table("t", empty_schema()).execute().await?;
assert!(table.list_bases().await?.is_empty());
let media = tmp.path().join("media");
std::fs::create_dir_all(&media).unwrap();
let media_uri = file_uri(&media);
table.add_bases([&media_uri]).await?;
assert_eq!(
table.list_bases().await?,
vec![TableBase {
path: media_uri,
name: None,
is_dataset_root: false,
}]
);
Ok(())
}