|
6 | 6 | from vgi.catalog import ( |
7 | 7 | AttachId, |
8 | 8 | CatalogAttachResult, |
| 9 | + ExtensionOption, |
9 | 10 | FunctionInfo, |
10 | 11 | FunctionType, |
11 | 12 | ScanFunctionResult, |
@@ -71,6 +72,44 @@ def test_all_flags_true(self) -> None: |
71 | 72 | assert restored.supports_time_travel is True |
72 | 73 | assert restored.catalog_version_frozen is True |
73 | 74 |
|
| 75 | + def test_with_settings(self) -> None: |
| 76 | + """Test with settings list.""" |
| 77 | + # Create some serialized ExtensionOption bytes |
| 78 | + option_bytes_1 = b"serialized_option_1" |
| 79 | + option_bytes_2 = b"serialized_option_2" |
| 80 | + |
| 81 | + original = CatalogAttachResult( |
| 82 | + attach_id=AttachId(b"\x01\x02"), |
| 83 | + supports_transactions=True, |
| 84 | + supports_time_travel=False, |
| 85 | + catalog_version_frozen=False, |
| 86 | + catalog_version=1, |
| 87 | + settings=[option_bytes_1, option_bytes_2], |
| 88 | + ) |
| 89 | + serialized = original.serialize() |
| 90 | + batch = deserialize_record_batch(serialized) |
| 91 | + restored = CatalogAttachResult.deserialize(batch) |
| 92 | + |
| 93 | + assert len(restored.settings) == 2 |
| 94 | + assert restored.settings[0] == option_bytes_1 |
| 95 | + assert restored.settings[1] == option_bytes_2 |
| 96 | + |
| 97 | + def test_with_empty_settings(self) -> None: |
| 98 | + """Test with empty settings list.""" |
| 99 | + original = CatalogAttachResult( |
| 100 | + attach_id=AttachId(b"\x01\x02"), |
| 101 | + supports_transactions=True, |
| 102 | + supports_time_travel=False, |
| 103 | + catalog_version_frozen=False, |
| 104 | + catalog_version=1, |
| 105 | + settings=[], |
| 106 | + ) |
| 107 | + serialized = original.serialize() |
| 108 | + batch = deserialize_record_batch(serialized) |
| 109 | + restored = CatalogAttachResult.deserialize(batch) |
| 110 | + |
| 111 | + assert restored.settings == [] |
| 112 | + |
74 | 113 |
|
75 | 114 | class TestSchemaInfoSerialization: |
76 | 115 | """Test SchemaInfo serialization round-trip.""" |
@@ -307,6 +346,140 @@ def test_table_function(self) -> None: |
307 | 346 |
|
308 | 347 | assert restored.function_type == FunctionType.TABLE |
309 | 348 |
|
| 349 | + def test_examples_with_descriptions(self) -> None: |
| 350 | + """Test structured examples with descriptions.""" |
| 351 | + from vgi.catalog import CatalogExample |
| 352 | + |
| 353 | + args_schema = schema(value=pa.int64()) |
| 354 | + output_schema = schema(result=pa.int64()) |
| 355 | + |
| 356 | + original = FunctionInfo( |
| 357 | + name="double", |
| 358 | + schema_name="main", |
| 359 | + function_type=FunctionType.SCALAR, |
| 360 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 361 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 362 | + comment="Double the input", |
| 363 | + tags={}, |
| 364 | + examples=[ |
| 365 | + CatalogExample( |
| 366 | + sql="SELECT double(5)", |
| 367 | + description="Double the number 5", |
| 368 | + expected_output="10", |
| 369 | + ), |
| 370 | + CatalogExample( |
| 371 | + sql="SELECT double(-3)", |
| 372 | + description="Double a negative number", |
| 373 | + expected_output=None, |
| 374 | + ), |
| 375 | + ], |
| 376 | + ) |
| 377 | + serialized = original.serialize() |
| 378 | + batch = deserialize_record_batch(serialized) |
| 379 | + restored = FunctionInfo.deserialize(batch) |
| 380 | + |
| 381 | + assert len(restored.examples) == 2 |
| 382 | + ex0 = restored.examples[0] |
| 383 | + ex1 = restored.examples[1] |
| 384 | + assert isinstance(ex0, CatalogExample) |
| 385 | + assert isinstance(ex1, CatalogExample) |
| 386 | + assert ex0.sql == "SELECT double(5)" |
| 387 | + assert ex0.description == "Double the number 5" |
| 388 | + assert ex0.expected_output == "10" |
| 389 | + assert ex1.sql == "SELECT double(-3)" |
| 390 | + assert ex1.description == "Double a negative number" |
| 391 | + assert ex1.expected_output is None |
| 392 | + |
| 393 | + def test_examples_empty_descriptions(self) -> None: |
| 394 | + """Test examples with empty descriptions.""" |
| 395 | + from vgi.catalog import CatalogExample |
| 396 | + |
| 397 | + args_schema = schema(value=pa.int64()) |
| 398 | + output_schema = schema(result=pa.int64()) |
| 399 | + |
| 400 | + original = FunctionInfo( |
| 401 | + name="echo", |
| 402 | + schema_name="main", |
| 403 | + function_type=FunctionType.SCALAR, |
| 404 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 405 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 406 | + comment=None, |
| 407 | + tags={}, |
| 408 | + examples=[ |
| 409 | + CatalogExample(sql="SELECT echo('hi')", description=""), |
| 410 | + ], |
| 411 | + ) |
| 412 | + serialized = original.serialize() |
| 413 | + batch = deserialize_record_batch(serialized) |
| 414 | + restored = FunctionInfo.deserialize(batch) |
| 415 | + |
| 416 | + assert len(restored.examples) == 1 |
| 417 | + ex = restored.examples[0] |
| 418 | + assert isinstance(ex, CatalogExample) |
| 419 | + assert ex.sql == "SELECT echo('hi')" |
| 420 | + assert ex.description == "" |
| 421 | + |
| 422 | + def test_examples_empty_list(self) -> None: |
| 423 | + """Test with no examples.""" |
| 424 | + args_schema = schema(value=pa.int64()) |
| 425 | + output_schema = schema(result=pa.int64()) |
| 426 | + |
| 427 | + original = FunctionInfo( |
| 428 | + name="test_func", |
| 429 | + schema_name="main", |
| 430 | + function_type=FunctionType.SCALAR, |
| 431 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 432 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 433 | + comment=None, |
| 434 | + tags={}, |
| 435 | + examples=[], |
| 436 | + ) |
| 437 | + serialized = original.serialize() |
| 438 | + batch = deserialize_record_batch(serialized) |
| 439 | + restored = FunctionInfo.deserialize(batch) |
| 440 | + |
| 441 | + assert restored.examples == [] |
| 442 | + |
| 443 | + def test_tags_dict_serialization(self) -> None: |
| 444 | + """Test tags dict serialization and deserialization.""" |
| 445 | + args_schema = schema(value=pa.int64()) |
| 446 | + output_schema = schema(result=pa.int64()) |
| 447 | + |
| 448 | + original = FunctionInfo( |
| 449 | + name="tagged_func", |
| 450 | + schema_name="main", |
| 451 | + function_type=FunctionType.SCALAR, |
| 452 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 453 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 454 | + comment=None, |
| 455 | + tags={"category": "math", "type": "scalar", "version": "1.0"}, |
| 456 | + ) |
| 457 | + serialized = original.serialize() |
| 458 | + batch = deserialize_record_batch(serialized) |
| 459 | + restored = FunctionInfo.deserialize(batch) |
| 460 | + |
| 461 | + assert restored.tags == {"category": "math", "type": "scalar", "version": "1.0"} |
| 462 | + |
| 463 | + def test_empty_tags_dict(self) -> None: |
| 464 | + """Test with empty tags dict.""" |
| 465 | + args_schema = schema(value=pa.int64()) |
| 466 | + output_schema = schema(result=pa.int64()) |
| 467 | + |
| 468 | + original = FunctionInfo( |
| 469 | + name="no_tags_func", |
| 470 | + schema_name="main", |
| 471 | + function_type=FunctionType.SCALAR, |
| 472 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 473 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 474 | + comment=None, |
| 475 | + tags={}, |
| 476 | + ) |
| 477 | + serialized = original.serialize() |
| 478 | + batch = deserialize_record_batch(serialized) |
| 479 | + restored = FunctionInfo.deserialize(batch) |
| 480 | + |
| 481 | + assert restored.tags == {} |
| 482 | + |
310 | 483 |
|
311 | 484 | class TestScanFunctionResultSerialization: |
312 | 485 | """Test ScanFunctionResult serialization round-trip.""" |
@@ -340,19 +513,150 @@ def test_none_invocation_id(self) -> None: |
340 | 513 | assert restored.invocation_id is None |
341 | 514 |
|
342 | 515 |
|
| 516 | +class TestExtensionOptionSerialization: |
| 517 | + """Test ExtensionOption serialization round-trip.""" |
| 518 | + |
| 519 | + def test_basic_round_trip(self) -> None: |
| 520 | + """Test basic serialization and deserialization.""" |
| 521 | + type_schema = pa.schema([pa.field("value", pa.bool_())]) |
| 522 | + original = ExtensionOption( |
| 523 | + name="vgi_verbose_mode", |
| 524 | + description="Enable verbose output", |
| 525 | + type=SerializedSchema(type_schema.serialize().to_pybytes()), |
| 526 | + default_value=None, |
| 527 | + ) |
| 528 | + serialized = original.serialize() |
| 529 | + batch = deserialize_record_batch(serialized) |
| 530 | + restored = ExtensionOption.deserialize(batch) |
| 531 | + |
| 532 | + assert restored.name == original.name |
| 533 | + assert restored.description == original.description |
| 534 | + assert restored.type == original.type |
| 535 | + assert restored.default_value is None |
| 536 | + |
| 537 | + def test_with_default_value(self) -> None: |
| 538 | + """Test with a default value.""" |
| 539 | + type_schema = pa.schema([pa.field("value", pa.string())]) |
| 540 | + # Create a default value batch |
| 541 | + default_batch = pa.RecordBatch.from_pydict( |
| 542 | + {"value": ["default"]}, schema=type_schema |
| 543 | + ) |
| 544 | + from vgi.ipc_utils import serialize_record_batch |
| 545 | + |
| 546 | + default_bytes = serialize_record_batch(default_batch) |
| 547 | + |
| 548 | + original = ExtensionOption( |
| 549 | + name="vgi_log_level", |
| 550 | + description="Logging level", |
| 551 | + type=SerializedSchema(type_schema.serialize().to_pybytes()), |
| 552 | + default_value=default_bytes, |
| 553 | + ) |
| 554 | + serialized = original.serialize() |
| 555 | + batch = deserialize_record_batch(serialized) |
| 556 | + restored = ExtensionOption.deserialize(batch) |
| 557 | + |
| 558 | + assert restored.name == original.name |
| 559 | + assert restored.description == original.description |
| 560 | + assert restored.default_value == default_bytes |
| 561 | + |
| 562 | + def test_integer_type(self) -> None: |
| 563 | + """Test with integer type.""" |
| 564 | + type_schema = pa.schema([pa.field("value", pa.int64())]) |
| 565 | + original = ExtensionOption( |
| 566 | + name="vgi_max_workers", |
| 567 | + description="Maximum worker count", |
| 568 | + type=SerializedSchema(type_schema.serialize().to_pybytes()), |
| 569 | + ) |
| 570 | + serialized = original.serialize() |
| 571 | + batch = deserialize_record_batch(serialized) |
| 572 | + restored = ExtensionOption.deserialize(batch) |
| 573 | + |
| 574 | + # Verify the type can be deserialized back |
| 575 | + restored_type = pa.ipc.read_schema(pa.py_buffer(restored.type)) |
| 576 | + assert restored_type.field("value").type == pa.int64() |
| 577 | + |
| 578 | + |
| 579 | +class TestFunctionInfoRequiredSettings: |
| 580 | + """Test FunctionInfo with required_settings field.""" |
| 581 | + |
| 582 | + def test_empty_required_settings(self) -> None: |
| 583 | + """Test with no required settings.""" |
| 584 | + args_schema = schema(value=pa.int64()) |
| 585 | + output_schema = schema(result=pa.int64()) |
| 586 | + |
| 587 | + original = FunctionInfo( |
| 588 | + name="echo", |
| 589 | + schema_name="main", |
| 590 | + function_type=FunctionType.SCALAR, |
| 591 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 592 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 593 | + comment=None, |
| 594 | + tags={}, |
| 595 | + required_settings=[], |
| 596 | + ) |
| 597 | + serialized = original.serialize() |
| 598 | + batch = deserialize_record_batch(serialized) |
| 599 | + restored = FunctionInfo.deserialize(batch) |
| 600 | + |
| 601 | + assert restored.required_settings == [] |
| 602 | + |
| 603 | + def test_single_required_setting(self) -> None: |
| 604 | + """Test with a single required setting.""" |
| 605 | + args_schema = schema(count=pa.int64()) |
| 606 | + output_schema = schema(result=pa.string()) |
| 607 | + |
| 608 | + original = FunctionInfo( |
| 609 | + name="settings_aware", |
| 610 | + schema_name="main", |
| 611 | + function_type=FunctionType.TABLE, |
| 612 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 613 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 614 | + comment=None, |
| 615 | + tags={}, |
| 616 | + required_settings=["vgi_verbose_mode"], |
| 617 | + ) |
| 618 | + serialized = original.serialize() |
| 619 | + batch = deserialize_record_batch(serialized) |
| 620 | + restored = FunctionInfo.deserialize(batch) |
| 621 | + |
| 622 | + assert restored.required_settings == ["vgi_verbose_mode"] |
| 623 | + |
| 624 | + def test_multiple_required_settings(self) -> None: |
| 625 | + """Test with multiple required settings.""" |
| 626 | + args_schema = schema(value=pa.int64()) |
| 627 | + output_schema = schema(result=pa.int64()) |
| 628 | + |
| 629 | + original = FunctionInfo( |
| 630 | + name="logging_function", |
| 631 | + schema_name="main", |
| 632 | + function_type=FunctionType.TABLE, |
| 633 | + arguments=SerializedSchema(args_schema.serialize().to_pybytes()), |
| 634 | + output_schema=SerializedSchema(output_schema.serialize().to_pybytes()), |
| 635 | + comment=None, |
| 636 | + tags={}, |
| 637 | + required_settings=["vgi_log_level", "vgi_log_format"], |
| 638 | + ) |
| 639 | + serialized = original.serialize() |
| 640 | + batch = deserialize_record_batch(serialized) |
| 641 | + restored = FunctionInfo.deserialize(batch) |
| 642 | + |
| 643 | + assert restored.required_settings == ["vgi_log_level", "vgi_log_format"] |
| 644 | + |
| 645 | + |
343 | 646 | class TestArrowSchemaCorrectness: |
344 | 647 | """Test that Arrow schemas are correct for each type.""" |
345 | 648 |
|
346 | 649 | def test_catalog_attach_result_schema(self) -> None: |
347 | 650 | """Verify CatalogAttachResult Arrow schema.""" |
348 | 651 | schema = CatalogAttachResult.ARROW_SCHEMA |
349 | | - assert len(schema) == 6 |
| 652 | + assert len(schema) == 7 |
350 | 653 | assert schema.field("attach_id").type == pa.binary() |
351 | 654 | assert schema.field("supports_transactions").type == pa.bool_() |
352 | 655 | assert schema.field("supports_time_travel").type == pa.bool_() |
353 | 656 | assert schema.field("catalog_version_frozen").type == pa.bool_() |
354 | 657 | assert schema.field("catalog_version").type == pa.int64() |
355 | 658 | assert schema.field("attach_id_required").type == pa.bool_() |
| 659 | + assert schema.field("settings").type == pa.list_(pa.binary()) |
356 | 660 |
|
357 | 661 | def test_schema_info_schema(self) -> None: |
358 | 662 | """Verify SchemaInfo Arrow schema.""" |
|
0 commit comments