diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c1489b9af..9d951f559 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2617,6 +2617,10 @@ class Table(ABC): to nullable. Currently, you cannot change a nullable column to non-nullable. + Each alteration must specify at least one of "rename", "data_type", + or "nullable". The legacy key "name" is also accepted for renaming. + Unknown keys raise ValueError before any alterations are applied. + Returns ------- AlterColumnsResult @@ -6874,6 +6878,10 @@ class AsyncTable: to nullable. Currently, you cannot change a nullable column to non-nullable. + Each alteration must specify at least one of "rename", "data_type", + or "nullable". The legacy key "name" is also accepted for renaming. + Unknown keys raise ValueError before any alterations are applied. + Returns ------- AlterColumnsResult diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index fb88afd40..47f1f866e 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -119,6 +119,77 @@ async def test_async_remote_db(): assert table_names == [] +@pytest.mark.parametrize( + "alteration, match", + [ + ({"path": "id"}, "One of rename, nullable or data_type"), + ( + {"path": "id", "nulable": False}, # spellchecker:disable-line + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ( + { + "path": "id", + "rename": "new_id", + "nulable": False, # spellchecker:disable-line + }, + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ], +) +def test_remote_alter_columns_rejects_invalid_before_request(alteration, match): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + with mock_lancedb_connection(handler) as db: + table = db.open_table("test") + requests.clear() + + with pytest.raises(ValueError, match=match): + table.alter_columns(alteration) + + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "alteration, match", + [ + ({"path": "id"}, "One of rename, nullable or data_type"), + ( + {"path": "id", "nulable": False}, # spellchecker:disable-line + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ], +) +async def test_async_remote_alter_columns_rejects_invalid_before_request( + alteration, match +): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test") + requests.clear() + + with pytest.raises(ValueError, match=match): + await table.alter_columns(alteration) + + assert requests == [] + + @pytest.mark.asyncio async def test_async_checkout(): def handler(request): diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 2fe6bc6ff..0c12c5c20 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4170,14 +4170,74 @@ async def test_add_columns_with_schema(mem_db_async: AsyncConnection): ) -def test_alter_columns(mem_db: DBConnection): +@pytest.mark.parametrize("rename_key", ["rename", "name"]) +def test_alter_columns(mem_db: DBConnection, rename_key): data = pa.table({"id": [0, 1]}) table = mem_db.create_table("my_table", data=data) - alter_columns_res = table.alter_columns({"path": "id", "rename": "new_id"}) + alter_columns_res = table.alter_columns({"path": "id", rename_key: "new_id"}) assert alter_columns_res.version == 2 assert table.to_arrow().column_names == ["new_id"] +INVALID_COLUMN_ALTERATIONS = [ + (({"path": "id"},), "One of rename, nullable or data_type"), + ( + ({"path": "id", "nulable": False},), # spellchecker:disable-line + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ( + ( + { + "path": "id", + "rename": "new_id", + "nulable": False, # spellchecker:disable-line + }, + ), + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ( + ({"path": "id", "rename": "new_id"}, {"path": "id"}), + "One of rename, nullable or data_type", + ), +] + + +def test_alter_columns_nullable_false(mem_db: DBConnection): + table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]})) + result = table.alter_columns({"path": "id", "nullable": False}) + assert result.version == 2 + assert not table.schema.field("id").nullable + + +@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS) +def test_alter_columns_rejects_invalid(mem_db: DBConnection, alterations, match): + table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]})) + initial_version = table.version + initial_schema = table.schema + + with pytest.raises(ValueError, match=match): + table.alter_columns(*alterations) + + assert table.version == initial_version + assert table.schema == initial_schema + + +@pytest.mark.asyncio +@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS) +async def test_alter_columns_rejects_invalid_async( + mem_db_async: AsyncConnection, alterations, match +): + table = await mem_db_async.create_table("my_table", data=pa.table({"id": [0, 1]})) + initial_version = await table.version() + initial_schema = await table.schema() + + with pytest.raises(ValueError, match=match): + await table.alter_columns(*alterations) + + assert await table.version() == initial_version + assert await table.schema() == initial_schema + + def test_update_field_metadata(mem_db: DBConnection): data = pa.table({"id": [0, 1], "category": ["a", "b"]}) table = mem_db.create_table("my_table", data=data) @@ -4203,10 +4263,11 @@ def test_update_field_metadata(mem_db: DBConnection): @pytest.mark.asyncio -async def test_alter_columns_async(mem_db_async: AsyncConnection): +@pytest.mark.parametrize("rename_key", ["rename", "name"]) +async def test_alter_columns_async(mem_db_async: AsyncConnection, rename_key): data = pa.table({"id": [0, 1]}) table = await mem_db_async.create_table("my_table", data=data) - alter_columns_res = await table.alter_columns({"path": "id", "rename": "new_id"}) + alter_columns_res = await table.alter_columns({"path": "id", rename_key: "new_id"}) assert alter_columns_res.version == 2 assert (await table.to_arrow()).column_names == ["new_id"] alter_columns_res = await table.alter_columns( diff --git a/python/src/table.rs b/python/src/table.rs index e7b6f1c6f..d7519bdab 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1897,6 +1897,17 @@ impl Table { let alterations = alterations .iter() .map(|alteration| { + for key in alteration.keys().iter() { + let key = key.extract::()?; + if !matches!( + key.as_str(), + "path" | "rename" | "name" | "nullable" | "data_type" + ) { + return Err(PyValueError::new_err(format!( + "Unknown column alteration key '{key}'" + ))); + } + } let path = alteration .get_item("path")? .ok_or_else(|| PyValueError::new_err("Missing path"))? diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index d43bab9db..74f4b70ef 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -4694,6 +4694,25 @@ mod tests { assert_eq!(result.rows_updated, if old_server { 0 } else { 5 }); } + #[tokio::test] + async fn test_alter_columns_rejects_missing_changes_before_request() { + let table = Table::new_with_handler::("my_table", |request| { + panic!("Unexpected request: {}", request.url().path()) + }); + + for alterations in [ + vec![ColumnAlteration::new("id".into())], + vec![ + ColumnAlteration::new("id".into()).rename("new_id".into()), + ColumnAlteration::new("id".into()), + ], + ] { + let err = table.alter_columns(&alterations).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("path 'id'")); + } + } + #[rstest] #[case(true)] #[case(false)] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index a99f086ef..a95d8ddf3 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1919,11 +1919,28 @@ impl Table { self.inner.function_errors(&request).await } - /// Change a column's name or nullability. + /// Change a column's name, data type, or nullability. + /// + /// 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. pub async fn alter_columns( &self, alterations: &[ColumnAlteration], ) -> Result { + for alteration in alterations { + if alteration.rename.is_none() + && alteration.nullable.is_none() + && alteration.data_type.is_none() + { + return Err(Error::InvalidInput { + message: format!( + "One of rename, nullable or data_type must be specified for path '{}'", + alteration.path + ), + }); + } + } self.inner.alter_columns(alterations).await } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 4d0e2815d..4990901ae 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -807,6 +807,29 @@ mod tests { // Alter Columns Tests + #[tokio::test] + async fn test_alter_columns_rejects_missing_changes() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap(); + let table = conn.create_table("test", batch).execute().await.unwrap(); + let initial_version = table.version().await.unwrap(); + let initial_schema = table.schema().await.unwrap(); + + for alterations in [ + vec![ColumnAlteration::new("id".into())], + vec![ + ColumnAlteration::new("id".into()).rename("new_id".into()), + ColumnAlteration::new("id".into()), + ], + ] { + let err = table.alter_columns(&alterations).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("path 'id'")); + assert_eq!(table.version().await.unwrap(), initial_version); + assert_eq!(table.schema().await.unwrap(), initial_schema); + } + } + #[tokio::test] async fn test_alter_column_rename() { let conn = connect("memory://").execute().await.unwrap();