|
| 1 | +import json |
1 | 2 | import pathlib |
2 | 3 |
|
3 | 4 | import pyarrow as pa |
@@ -305,3 +306,207 @@ def test_set_table_properties_enable_dv(tmp_path: pathlib.Path, sample_table: pa |
305 | 306 | assert protocol.min_writer_version == 7 |
306 | 307 | assert protocol.writer_features == ["deletionVectors"] |
307 | 308 | assert protocol.reader_features == ["deletionVectors"] |
| 309 | + |
| 310 | + |
| 311 | +data = pa.Table.from_pydict( |
| 312 | + { |
| 313 | + "foo": [1], |
| 314 | + "bar": [{"baz": {"nested": 1, "nested_02": 1}, "foo": 2}], |
| 315 | + "barz": [{"baz": {"nested": 1}}], |
| 316 | + } |
| 317 | +) |
| 318 | + |
| 319 | + |
| 320 | +def test_drop_columns_single_col(tmp_path): |
| 321 | + write_deltalake(tmp_path, data) |
| 322 | + |
| 323 | + dt = DeltaTable(tmp_path) |
| 324 | + dt.alter.drop_columns("foo") |
| 325 | + expected = { |
| 326 | + "type": "struct", |
| 327 | + "fields": [ |
| 328 | + { |
| 329 | + "name": "bar", |
| 330 | + "type": { |
| 331 | + "type": "struct", |
| 332 | + "fields": [ |
| 333 | + { |
| 334 | + "name": "baz", |
| 335 | + "type": { |
| 336 | + "type": "struct", |
| 337 | + "fields": [ |
| 338 | + { |
| 339 | + "name": "nested", |
| 340 | + "type": "long", |
| 341 | + "nullable": True, |
| 342 | + "metadata": {}, |
| 343 | + }, |
| 344 | + { |
| 345 | + "name": "nested_02", |
| 346 | + "type": "long", |
| 347 | + "nullable": True, |
| 348 | + "metadata": {}, |
| 349 | + }, |
| 350 | + ], |
| 351 | + }, |
| 352 | + "nullable": True, |
| 353 | + "metadata": {}, |
| 354 | + }, |
| 355 | + { |
| 356 | + "name": "foo", |
| 357 | + "type": "long", |
| 358 | + "nullable": True, |
| 359 | + "metadata": {}, |
| 360 | + }, |
| 361 | + ], |
| 362 | + }, |
| 363 | + "nullable": True, |
| 364 | + "metadata": {}, |
| 365 | + }, |
| 366 | + { |
| 367 | + "name": "barz", |
| 368 | + "type": { |
| 369 | + "type": "struct", |
| 370 | + "fields": [ |
| 371 | + { |
| 372 | + "name": "baz", |
| 373 | + "type": { |
| 374 | + "type": "struct", |
| 375 | + "fields": [ |
| 376 | + { |
| 377 | + "name": "nested", |
| 378 | + "type": "long", |
| 379 | + "nullable": True, |
| 380 | + "metadata": {}, |
| 381 | + } |
| 382 | + ], |
| 383 | + }, |
| 384 | + "nullable": True, |
| 385 | + "metadata": {}, |
| 386 | + } |
| 387 | + ], |
| 388 | + }, |
| 389 | + "nullable": True, |
| 390 | + "metadata": {}, |
| 391 | + }, |
| 392 | + ], |
| 393 | + } |
| 394 | + assert json.loads(dt.schema().to_json()) == expected |
| 395 | + |
| 396 | + |
| 397 | +def test_drop_columns_multiple_cols(tmp_path): |
| 398 | + write_deltalake(tmp_path, data) |
| 399 | + |
| 400 | + dt = DeltaTable(tmp_path) |
| 401 | + dt.alter.drop_columns(["foo", "bar"]) |
| 402 | + expected = { |
| 403 | + "type": "struct", |
| 404 | + "fields": [ |
| 405 | + { |
| 406 | + "name": "barz", |
| 407 | + "type": { |
| 408 | + "type": "struct", |
| 409 | + "fields": [ |
| 410 | + { |
| 411 | + "name": "baz", |
| 412 | + "type": { |
| 413 | + "type": "struct", |
| 414 | + "fields": [ |
| 415 | + { |
| 416 | + "name": "nested", |
| 417 | + "type": "long", |
| 418 | + "nullable": True, |
| 419 | + "metadata": {}, |
| 420 | + } |
| 421 | + ], |
| 422 | + }, |
| 423 | + "nullable": True, |
| 424 | + "metadata": {}, |
| 425 | + } |
| 426 | + ], |
| 427 | + }, |
| 428 | + "nullable": True, |
| 429 | + "metadata": {}, |
| 430 | + } |
| 431 | + ], |
| 432 | + } |
| 433 | + assert json.loads(dt.schema().to_json()) == expected |
| 434 | + |
| 435 | + |
| 436 | +def test_drop_columns_multiple_cols_nested(tmp_path): |
| 437 | + write_deltalake(tmp_path, data) |
| 438 | + |
| 439 | + dt = DeltaTable(tmp_path) |
| 440 | + dt.alter.drop_columns(["foo", "bar.baz.nested"]) |
| 441 | + expected = { |
| 442 | + "type": "struct", |
| 443 | + "fields": [ |
| 444 | + { |
| 445 | + "name": "bar", |
| 446 | + "type": { |
| 447 | + "type": "struct", |
| 448 | + "fields": [ |
| 449 | + { |
| 450 | + "name": "baz", |
| 451 | + "type": { |
| 452 | + "type": "struct", |
| 453 | + "fields": [ |
| 454 | + { |
| 455 | + "name": "nested_02", |
| 456 | + "type": "long", |
| 457 | + "nullable": True, |
| 458 | + "metadata": {}, |
| 459 | + } |
| 460 | + ], |
| 461 | + }, |
| 462 | + "nullable": True, |
| 463 | + "metadata": {}, |
| 464 | + }, |
| 465 | + { |
| 466 | + "name": "foo", |
| 467 | + "type": "long", |
| 468 | + "nullable": True, |
| 469 | + "metadata": {}, |
| 470 | + }, |
| 471 | + ], |
| 472 | + }, |
| 473 | + "nullable": True, |
| 474 | + "metadata": {}, |
| 475 | + }, |
| 476 | + { |
| 477 | + "name": "barz", |
| 478 | + "type": { |
| 479 | + "type": "struct", |
| 480 | + "fields": [ |
| 481 | + { |
| 482 | + "name": "baz", |
| 483 | + "type": { |
| 484 | + "type": "struct", |
| 485 | + "fields": [ |
| 486 | + { |
| 487 | + "name": "nested", |
| 488 | + "type": "long", |
| 489 | + "nullable": True, |
| 490 | + "metadata": {}, |
| 491 | + } |
| 492 | + ], |
| 493 | + }, |
| 494 | + "nullable": True, |
| 495 | + "metadata": {}, |
| 496 | + } |
| 497 | + ], |
| 498 | + }, |
| 499 | + "nullable": True, |
| 500 | + "metadata": {}, |
| 501 | + }, |
| 502 | + ], |
| 503 | + } |
| 504 | + assert json.loads(dt.schema().to_json()) == expected |
| 505 | + |
| 506 | + |
| 507 | +def test_drop_columns_raise_not_exists(tmp_path): |
| 508 | + write_deltalake(tmp_path, data) |
| 509 | + |
| 510 | + dt = DeltaTable(tmp_path) |
| 511 | + with pytest.raises(DeltaError): |
| 512 | + dt.alter.drop_columns(["foo", "not_exists"], raise_if_not_exists=True) |
0 commit comments