| 601 | |
| 602 | |
| 603 | def test_column_selection(tempdir): |
| 604 | from pyarrow import orc |
| 605 | |
| 606 | # create a table with nested types |
| 607 | inner = pa.field('inner', pa.int64()) |
| 608 | middle = pa.field('middle', pa.struct([inner])) |
| 609 | fields = [ |
| 610 | pa.field('basic', pa.int32()), |
| 611 | pa.field( |
| 612 | 'list', pa.list_(pa.field('item', pa.int32())) |
| 613 | ), |
| 614 | pa.field( |
| 615 | 'struct', pa.struct([middle, pa.field('inner2', pa.int64())]) |
| 616 | ), |
| 617 | pa.field( |
| 618 | 'list-struct', pa.list_(pa.field( |
| 619 | 'item', pa.struct([ |
| 620 | pa.field('inner1', pa.int64()), |
| 621 | pa.field('inner2', pa.int64()) |
| 622 | ]) |
| 623 | )) |
| 624 | ), |
| 625 | pa.field('basic2', pa.int64()), |
| 626 | ] |
| 627 | arrs = [ |
| 628 | [0], [[1, 2]], [{"middle": {"inner": 3}, "inner2": 4}], |
| 629 | [[{"inner1": 5, "inner2": 6}, {"inner1": 7, "inner2": 8}]], [9]] |
| 630 | table = pa.table(arrs, schema=pa.schema(fields)) |
| 631 | |
| 632 | path = str(tempdir / 'test.orc') |
| 633 | orc.write_table(table, path) |
| 634 | orc_file = orc.ORCFile(path) |
| 635 | |
| 636 | # default selecting all columns |
| 637 | result1 = orc_file.read() |
| 638 | assert result1.equals(table) |
| 639 | |
| 640 | # selecting with columns names |
| 641 | result2 = orc_file.read(columns=["basic", "basic2"]) |
| 642 | assert result2.equals(table.select(["basic", "basic2"])) |
| 643 | |
| 644 | result3 = orc_file.read(columns=["list", "struct", "basic2"]) |
| 645 | assert result3.equals(table.select(["list", "struct", "basic2"])) |
| 646 | |
| 647 | # using dotted paths |
| 648 | result4 = orc_file.read(columns=["struct.middle.inner"]) |
| 649 | expected4 = pa.table({"struct": [{"middle": {"inner": 3}}]}) |
| 650 | assert result4.equals(expected4) |
| 651 | |
| 652 | result5 = orc_file.read(columns=["struct.inner2"]) |
| 653 | expected5 = pa.table({"struct": [{"inner2": 4}]}) |
| 654 | assert result5.equals(expected5) |
| 655 | |
| 656 | result6 = orc_file.read( |
| 657 | columns=["list", "struct.middle.inner", "struct.inner2"] |
| 658 | ) |
| 659 | assert result6.equals(table.select(["list", "struct"])) |
| 660 | |