diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index e766b3d2a..653d68ea1 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -1001,4 +1001,49 @@ 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; + } + res.writeHead(404).end(); + }, + async (db) => { + const table = await db.openTable("photos"); + await table.addBases({ path: "s3://bucket/media/" }); + }, + ); + expect(postedBodies).toEqual([ + { + bases: [ + { + path: "s3://bucket/media/", + isDatasetRoot: false, + }, + ], + }, + ]); + }); }); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 5396a251a..192170dd5 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -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,22 @@ 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("addBases accepts a file uri", 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); + await table.addBases(pathToFileURL(media).toString()); + }); +}); diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 9f2e97989..d30b101d7 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -130,6 +130,7 @@ export { export { Table, + TableBase, Branches, BranchColumnSummary, BranchColumnChange, diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index a7dc8def1..11c615db6 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -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,15 @@ export abstract class Table { | { computed: AddColumnsSql[] }, ): Promise; + /** + * Register additional storage bases for this table. + * + * A URI string is a non-root base with no alias. + */ + abstract addBases( + bases: string | TableBase | Array, + ): Promise; + /** * Fill the rows of a computed column that hold no value yet. * @@ -1196,6 +1224,12 @@ export class LocalTable extends Table { throw new Error("Invalid input type for addColumns"); } + async addBases( + bases: string | TableBase | Array, + ): Promise { + await this.inner.addBases(normalizeBases(bases)); + } + async refreshColumn(column: string): Promise { return await this.inner.refreshColumn(column); } @@ -1396,6 +1430,21 @@ export class LocalTable extends Table { } } +function normalizeBases( + bases: string | TableBase | Array, +): 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, diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 4c45be668..e4603d8e3 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -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,18 @@ impl Table { Ok(res.into()) } + #[napi(catch_unwind)] + pub async fn add_bases(&self, bases: Vec) -> 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 drop_columns(&self, columns: Vec) -> napi::Result { let col_refs = columns.iter().map(String::as_str).collect::>(); @@ -700,6 +713,18 @@ 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, + /// True when `path` is a Lance dataset root. When false, `path` is the + /// directory containing the referenced files. + pub is_dataset_root: bool, +} + #[napi(object)] /// A description of an index currently configured on a column pub struct IndexConfig { diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 235049f97..494bdd9a7 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -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__", ] diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 22878fd85..179c200b1 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -377,6 +377,7 @@ 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 fetch_blobs( self, column: str, row_ids: list[int] ) -> pa.LargeBinaryArray: ... diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index aa822b913..862950ba1 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -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,13 @@ 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 fetch_blobs( self, column: str, row_ids: Union[list[int], pa.Table] ) -> pa.LargeBinaryArray: diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 4ecf6e836..ead870e0b 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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,18 @@ 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 + @abstractmethod def fetch_blobs( self, column: str, row_ids: Union[list[int], pa.Table] @@ -2414,6 +2442,12 @@ 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 fetch_blobs( self, column: str, row_ids: Union[list[int], pa.Table] ) -> pa.LargeBinaryArray: @@ -6266,6 +6300,18 @@ 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 fetch_blobs( self, column: str, row_ids: Union[list[int], pa.Table] ) -> pa.LargeBinaryArray: @@ -6484,6 +6530,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: """ diff --git a/python/python/tests/test_bases.py b/python/python/tests/test_bases.py new file mode 100644 index 000000000..1b111adfe --- /dev/null +++ b/python/python/tests/test_bases.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pyarrow as pa +import pytest + +import lancedb + + +def test_add_bases_accepts_named_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 + ), + ] + ) + + +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()]) + + +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) + await table.add_bases(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) + table.add_bases(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) + table.add_bases(media.as_uri()) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index ce8d5bd6e..40173d8c8 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2306,3 +2306,36 @@ 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_bases_posts_the_bases_array(): + 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}') + 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 captured_body["bases"] == [ + { + "path": "s3://bucket/media/", + "isDatasetRoot": False, + } + ] + diff --git a/python/src/table.rs b/python/src/table.rs index 35ee92dc4..47c063cfd 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -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, + is_dataset_root: bool, +} + #[derive(FromPyObject)] enum PredicateArg { Expr(PyExpr), @@ -1238,6 +1246,25 @@ impl Table { }) } + #[pyo3(signature = (bases))] + pub fn add_bases( + self_: PyRef<'_, Self>, + bases: Vec, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + let bases: Vec = 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() + }) + } + /// Read blob bytes for `row_ids` from blob v2 column `column`. #[pyo3(signature = (column, row_ids))] pub fn fetch_blobs( diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 70d023ccc..e4e803d55 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -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. /// diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a0a4cebc2..9068e98f5 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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,10 @@ impl BaseTable for RemoteTable { self.blob_columns_impl().await } + async fn add_bases(&self, bases: &[crate::table::TableBase]) -> Result<()> { + self.add_bases_impl(bases).await + } + async fn fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result { self.fetch_blobs_impl(column, row_ids).await } @@ -4089,6 +4094,42 @@ 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 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))] diff --git a/rust/lancedb/src/remote/table/bases.rs b/rust/lancedb/src/remote/table/bases.rs new file mode 100644 index 000000000..5806aabb8 --- /dev/null +++ b/rust/lancedb/src/remote/table/bases.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Cloud HTTP for registering 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, +} + +impl RemoteTable { + 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(()) + } +} diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 2e16b0940..757899751 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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,12 @@ 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(), + }) + } /// Materialize blob bytes for the given row ids. See [`Table::fetch_blobs`]. async fn fetch_blobs(&self, _column: &str, _row_ids: &[u64]) -> Result { Err(Error::NotSupported { @@ -889,6 +895,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, + /// 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) -> 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 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 +1180,25 @@ 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> { + /// table.add_bases(["s3://bucket/media/"]).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn add_bases( + &self, + bases: impl IntoIterator>, + ) -> Result<()> { + let bases: Vec = bases.into_iter().map(Into::into).collect(); + self.inner.add_bases(&bases).await + } + /// Materialize blob bytes for the given row ids. /// /// Output matches `row_ids` in length and order. Null blobs are null; @@ -3364,6 +3437,25 @@ 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 fetch_blobs(&self, column: &str, row_ids: &[u64]) -> Result { let dataset = self.dataset.get().await?; crate::blob::take_blobs_aligned(&dataset, column, row_ids).await diff --git a/rust/lancedb/tests/table_bases.rs b/rust/lancedb/tests/table_bases.rs new file mode 100644 index 000000000..c75706943 --- /dev/null +++ b/rust/lancedb/tests/table_bases.rs @@ -0,0 +1,147 @@ +// 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 { + 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_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(); + + table + .add_bases([ + TableBase { + path: file_uri(&media), + name: Some("media".into()), + is_dataset_root: false, + }, + TableBase { + path: file_uri(&parent), + name: Some("parent".into()), + is_dataset_root: true, + }, + ]) + .await +} + +#[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(); + table + .add_bases([&file_uri(&media), &file_uri(&other)]) + .await +} + +#[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_memory_add_bases_accepts_a_file_uri() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect("memory://").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(); + table.add_bases([file_uri(&media)]).await +} + +#[tokio::test] +async fn test_namespace_add_bases_accepts_a_file_uri() -> 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?; + let media = tmp.path().join("media"); + std::fs::create_dir_all(&media).unwrap(); + table.add_bases([file_uri(&media)]).await +}