diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 7455ab294..a2a15838a 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -126,6 +126,9 @@ abstract alterColumns(columnAlterations): Promise Alter the name or nullability of columns. +Local tables reject renaming source or vector columns referenced by +`embedding_functions` schema metadata, preserving the embedding binding. + #### Parameters * **columnAlterations**: [`ColumnAlteration`](../interfaces/ColumnAlteration.md)[] @@ -504,6 +507,9 @@ underlying storage. In order to remove the data, you must subsequently call ``compact_files`` to rewrite the data without the removed columns and then call ``cleanup_files`` to remove the old files. +Local tables reject dropping source or vector columns referenced by +`embedding_functions` schema metadata, preserving the embedding binding. + #### Parameters * **columnNames**: `string`[] diff --git a/docs/web/management/table.mdx b/docs/web/management/table.mdx index 0b71fb315..fc20a3f2e 100644 --- a/docs/web/management/table.mdx +++ b/docs/web/management/table.mdx @@ -1116,6 +1116,21 @@ For transformations that are not simple casts (for example, converting `"$100"` #### Alter embedding types and dimensions {#schema-alter-embedding-types-and-dimensions} +{/* +Draft for the embedding schema-evolution protection release. + + +Local tables reject renaming or dropping source and vector columns referenced by +`embedding_functions` schema metadata or native Rust embedding definitions in +`lancedb::column_definitions`. The operation fails before changing any columns, +preserving the configuration used to generate embeddings on later writes. + +The add/drop/rename recipe below applies only to vector columns without a registered +embedding function. To change the dimension of a registered embedding column, create +a new table with the updated embedding configuration and regenerate its embeddings. + +*/} + It's quite common to need to change an embedding column's schema, in case a new model becomes available with a different embedding dimension. - In Python, the example shows an in-place type update when the cast is compatible. - In TypeScript and Rust, the example shows a dimension change (`384 -> 1024`), which cannot be cast in-place. diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index db0898f4b..3dbf914fe 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -713,6 +713,9 @@ export abstract class Table { /** * Alter the name or nullability of columns. + * + * Local tables reject renaming source or vector columns referenced by + * `embedding_functions` schema metadata, preserving the embedding binding. * @param {ColumnAlteration[]} columnAlterations One or more alterations to * apply to columns. * @returns {Promise} A promise that resolves to an object @@ -750,6 +753,9 @@ export abstract class Table { * underlying storage. In order to remove the data, you must subsequently * call ``compact_files`` to rewrite the data without the removed columns and * then call ``cleanup_files`` to remove the old files. + * + * Local tables reject dropping source or vector columns referenced by + * `embedding_functions` schema metadata, preserving the embedding binding. * @param {string[]} columnNames The names of the columns to drop. These can * be nested column references (e.g. "a.b.c") or top-level column names * (e.g. "a"). diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 8a7571fe6..676d70ac1 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3280,6 +3280,9 @@ class Table(ABC): or "nullable". The legacy key "name" is also accepted for renaming. Unknown keys raise ValueError before any alterations are applied. + Local tables reject renaming source or vector columns used by an + embedding function with ValueError, preserving the embedding binding. + Returns ------- AlterColumnsResult @@ -3348,6 +3351,9 @@ class Table(ABC): """ Drop columns from the table. + Local tables reject dropping source or vector columns used by an + embedding function with ValueError, preserving the embedding binding. + Parameters ---------- columns : Iterable[str] @@ -7744,6 +7750,9 @@ class AsyncTable: or "nullable". The legacy key "name" is also accepted for renaming. Unknown keys raise ValueError before any alterations are applied. + Local tables reject renaming source or vector columns used by an + embedding function with ValueError, preserving the embedding binding. + Returns ------- AlterColumnsResult @@ -7762,6 +7771,9 @@ class AsyncTable: """ Drop columns from the table. + Local tables reject dropping source or vector columns used by an + embedding function with ValueError, preserving the embedding binding. + Parameters ---------- columns : Iterable[str] diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index c18ddd644..78cd76ad2 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -33,6 +33,78 @@ def mock_embed_func(input_data): return [np.random.randn(128).tolist() for _ in range(len(input_data))] +@pytest.fixture +def embedding_schema(): + func = MockTextEmbeddingFunction.create() + + class Doc(LanceModel): + id: int + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + return Doc + + +@pytest.mark.parametrize("column", ["text", "vector", "`text`", "`vector`"]) +@pytest.mark.parametrize("operation", ["rename", "drop"]) +def test_embedding_columns_schema_evolution( + tmp_path, embedding_schema, column, operation +): + db = lancedb.connect(tmp_path) + table = db.create_table("docs", schema=embedding_schema) + table.add([{"id": 1, "text": "hello"}]) + schema, version = table.schema, table.version + + with pytest.raises(ValueError, match="used by an embedding function"): + if operation == "rename": + table.alter_columns( + {"path": "id", "rename": "doc_id"}, + {"path": column, "rename": "renamed"}, + ) + else: + table.drop_columns(["id", column]) + + assert table.schema == schema + assert table.version == version + assert table.embedding_functions["vector"].source_column == "text" + reopened = db.open_table("docs") + reopened.add([{"id": 2, "text": "world"}]) + assert reopened.count_rows() == 2 + assert reopened.to_arrow()["vector"].null_count == 0 + + # Columns outside the embedding binding can still evolve. + reopened.alter_columns({"path": "id", "rename": "doc_id"}) + reopened.drop_columns(["doc_id"]) + reopened.add([{"text": "again"}]) + assert reopened.count_rows() == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("column", ["text", "vector"]) +@pytest.mark.parametrize("operation", ["rename", "drop"]) +async def test_async_embedding_columns_schema_evolution( + tmp_path, embedding_schema, column, operation +): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("docs", schema=embedding_schema) + await table.add([{"id": 1, "text": "hello"}]) + schema, version = await table.schema(), await table.version() + + with pytest.raises(ValueError, match="used by an embedding function"): + if operation == "rename": + await table.alter_columns({"path": column, "rename": "renamed"}) + else: + await table.drop_columns([column]) + + assert await table.schema() == schema + assert await table.version() == version + assert (await table.embedding_functions())["vector"].source_column == "text" + reopened = await db.open_table("docs") + await reopened.add([{"id": 2, "text": "world"}]) + assert await reopened.count_rows() == 2 + assert (await reopened.to_arrow())["vector"].null_count == 0 + + def test_embedding_function(tmp_path): registry = EmbeddingFunctionRegistry.get_instance() diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 0c28820b5..8b008b4ec 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -2085,6 +2085,10 @@ impl Table { /// Each alteration must specify at least one of `rename`, `data_type`, or /// `nullable`. Returns [`Error::InvalidInput`] before applying any changes /// if an alteration does not specify any of these fields. + /// + /// Local tables reject renaming source or vector columns referenced by + /// [`EmbeddingDefinition`] or `embedding_functions` schema metadata to + /// preserve the embedding binding. pub async fn alter_columns( &self, alterations: &[ColumnAlteration], @@ -2128,6 +2132,10 @@ impl Table { } /// Remove columns from the table. + /// + /// Local tables reject dropping source or vector columns referenced by + /// [`EmbeddingDefinition`] or `embedding_functions` schema metadata to + /// preserve the embedding binding. pub async fn drop_columns(&self, columns: &[&str]) -> Result { self.inner.drop_columns(columns).await } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index d78fc5102..8723505ff 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -8,7 +8,7 @@ //! - [`alter_columns`](execute_alter_columns): Rename columns, change types, or modify nullability //! - [`drop_columns`](execute_drop_columns): Remove columns from the table -use arrow_schema::Schema as ArrowSchema; +use arrow_schema::{Schema as ArrowSchema, SchemaRef}; use lance::dataset::transaction::{Operation, Transaction, UpdateMap, UpdateMapEntry}; use lance::dataset::{ColumnAlteration, CommitBuilder, Dataset, NewColumnTransform}; use serde::{Deserialize, Serialize}; @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::sync::Arc; use super::computed_columns; -use super::{BaseTable, NativeTable}; +use super::{BaseTable, ColumnKind, NativeTable, TableDefinition}; use crate::{Error, Result}; /// The result of an add columns operation. @@ -201,6 +201,63 @@ pub(crate) async fn commit_add_columns( Ok(AddColumnsResult { version }) } +/// Native and SDK embedding configurations refer to top-level columns by name. +/// Lance's schema evolution preserves this metadata, so reject operations that +/// would leave those names dangling without loading an embedding function. +fn ensure_not_embedding_column<'a>( + schema: &SchemaRef, + operation: &str, + paths: impl IntoIterator, +) -> Result<()> { + #[derive(Deserialize)] + struct EmbeddingColumns { + #[serde(alias = "sourceColumn")] + source_column: String, + #[serde(alias = "vectorColumn", alias = "dest_column", alias = "destColumn")] + vector_column: String, + } + let mut embeddings: Vec = match schema.metadata.get("embedding_functions") { + Some(raw) => serde_json::from_str(raw).map_err(|error| Error::InvalidInput { + message: format!("Invalid embedding_functions metadata: {error}"), + })?, + None => Vec::new(), + }; + let definition = TableDefinition::try_from_rich_schema(schema.clone())?; + embeddings.extend( + definition + .column_definitions + .into_iter() + .filter_map(|column| { + let ColumnKind::Embedding(embedding) = column.kind else { + return None; + }; + let vector_column = embedding + .dest_column + .unwrap_or_else(|| format!("{}_embedding", embedding.source_column)); + Some(EmbeddingColumns { + source_column: embedding.source_column, + vector_column, + }) + }), + ); + for path in paths { + // Use Lance's path grammar so quoting or a nested path cannot bypass + // the check. Altering a child also changes the embedding's input. + let column = computed_columns::field_root(path)?; + if embeddings + .iter() + .any(|embedding| embedding.source_column == column || embedding.vector_column == column) + { + return Err(Error::InvalidInput { + message: format!( + "Cannot {operation} column '{path}' because it is used by an embedding function" + ), + }); + } + } + Ok(()) +} + /// Internal implementation of the alter columns logic. /// /// Alters existing columns in the table (rename, change type, or modify nullability). @@ -229,6 +286,14 @@ pub(crate) async fn execute_alter_columns( std::iter::once(alteration.path.as_str()).chain(alteration.rename.as_deref()) }), )?; + ensure_not_embedding_column( + &schema, + "rename", + alterations + .iter() + .filter(|alteration| alteration.rename.is_some()) + .map(|alteration| alteration.path.as_str()), + )?; // Nullability is not part of what an expression resolves against, so only // a rename or a retype can invalidate a binding. let rebinding = alterations @@ -271,6 +336,7 @@ pub(crate) async fn execute_drop_columns( let mut names = columns.iter().map(|c| c.to_string()).collect::>(); names.extend(unbinding.assignment_columns.iter().cloned()); let dropped = names.iter().map(String::as_str).collect::>(); + ensure_not_embedding_column(&schema, "drop", dropped.iter().copied())?; computed_columns::ensure_not_bound_by(&unbinding.retained, "schema evolution", &dropped)?; computed_columns::ensure_not_an_input_of(&schema, &dropped, &dropped)?; @@ -783,6 +849,85 @@ mod tests { // Alter Columns Tests + #[rstest::rstest] + #[case("source_column", "vector_column", "text")] + #[case("sourceColumn", "vectorColumn", "text")] + #[case("source_column", "vector_column", "text with space")] + #[tokio::test] + async fn test_schema_evolution_preserves_embedding_bindings( + #[case] source_key: &str, + #[case] vector_key: &str, + #[case] source_column: &str, + ) { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!( + ("id", Int32, [1]), + (source_column, Utf8, ["hello"]), + ("vector", Int32, [1]) + ) + .unwrap(); + let metadata = serde_json::json!([{ + "name": "test", + (source_key): source_column, + (vector_key): "vector", + "model": {}, + }]); + let schema = batch + .schema() + .as_ref() + .clone() + .with_metadata(HashMap::from([( + "embedding_functions".to_string(), + metadata.to_string(), + )])); + let batch = batch.with_schema(std::sync::Arc::new(schema)).unwrap(); + let table = conn + .create_table("test", batch.clone()) + .execute() + .await + .unwrap(); + let schema = table.schema().await.unwrap(); + let version = table.version().await.unwrap(); + + for path in [format!("`{source_column}`"), "vector".to_string()] { + let err = table + .alter_columns(&[ + ColumnAlteration::new("id".into()).rename("new_id".into()), + ColumnAlteration::new(path.clone()).rename("renamed".into()), + ]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("used by an embedding function")); + + let err = table.drop_columns(&["id", &path]).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("used by an embedding function")); + assert_eq!(table.schema().await.unwrap(), schema); + assert_eq!(table.version().await.unwrap(), version); + } + + let reopened = conn.open_table("test").execute().await.unwrap(); + assert_eq!(reopened.schema().await.unwrap(), schema); + reopened.add(batch).execute().await.unwrap(); + assert_eq!(reopened.count_rows(None).await.unwrap(), 2); + + // Nullability and operations on unrelated columns remain supported. + reopened + .alter_columns(&[ColumnAlteration::new("vector".into()).set_nullable(true)]) + .await + .unwrap(); + reopened + .alter_columns(&[ColumnAlteration::new("id".into()).rename("new_id".into())]) + .await + .unwrap(); + reopened.drop_columns(&["new_id"]).await.unwrap(); + assert_eq!( + reopened.schema().await.unwrap().metadata(), + schema.metadata() + ); + } + #[tokio::test] async fn test_alter_columns_rejects_missing_changes() { let conn = connect("memory://").execute().await.unwrap(); diff --git a/rust/lancedb/tests/embedding_registry_test.rs b/rust/lancedb/tests/embedding_registry_test.rs index 8e1a80ec4..e06f5e5b5 100644 --- a/rust/lancedb/tests/embedding_registry_test.rs +++ b/rust/lancedb/tests/embedding_registry_test.rs @@ -14,6 +14,7 @@ use arrow_array::{ }; use arrow_schema::{DataType, Field, Schema}; use futures::StreamExt; +use lance::dataset::ColumnAlteration; use lancedb::{ Error, Result, connect, embeddings::{EmbeddingDefinition, EmbeddingFunction, EmbeddingRegistry}, @@ -194,6 +195,69 @@ async fn test_open_table_embeddings() -> Result<()> { Ok(()) } +#[rstest::rstest] +#[case(Some("embeddings"), "embeddings")] +#[case(None, "text_embedding")] +#[tokio::test] +async fn test_schema_evolution_preserves_native_embeddings( + #[case] dest_column: Option<&str>, + #[case] vector_column: &str, +) -> Result<()> { + let tempdir = tempfile::tempdir().unwrap(); + let uri = tempdir.path().to_str().unwrap(); + let db = connect(uri).execute().await?; + db.embedding_registry() + .register("embed_fun", Arc::new(MockEmbed::new("embed_fun".into(), 2)))?; + let table = db + .create_table("test", create_some_records()?) + .add_embedding(EmbeddingDefinition::new("text", "embed_fun", dest_column))? + .execute() + .await?; + let schema = table.schema().await?; + let version = table.version().await?; + assert!( + schema + .metadata() + .contains_key("lancedb::column_definitions") + ); + assert!(!schema.metadata().contains_key("embedding_functions")); + + // Schema protection must not require the embedding function to be loaded. + let unregistered_db = connect(uri).execute().await?; + let table = unregistered_db.open_table("test").execute().await?; + for column in ["text", vector_column] { + let err = table + .alter_columns(&[ + ColumnAlteration::new("id".into()).rename("new_id".into()), + ColumnAlteration::new(column.into()).rename("renamed".into()), + ]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("used by an embedding function")); + assert_eq!(table.schema().await?, schema); + assert_eq!(table.version().await?, version); + + let err = table.drop_columns(&["id", column]).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("used by an embedding function")); + assert_eq!(table.schema().await?, schema); + assert_eq!(table.version().await?, version); + } + + let table = db.open_table("test").execute().await?; + assert_eq!(table.schema().await?, schema); + table.add(create_some_records()?).execute().await?; + assert_eq!(table.count_rows(None).await?, 4); + + table + .alter_columns(&[ColumnAlteration::new("id".into()).rename("new_id".into())]) + .await?; + table.drop_columns(&["new_id"]).await?; + assert_eq!(table.schema().await?.metadata(), schema.metadata()); + Ok(()) +} + #[tokio::test] async fn test_no_func_in_registry() -> Result<()> { let tempdir = tempfile::tempdir().unwrap();