Check a ORC file against the expected columns dictionary.
(orc_path, expected_df, need_fix=False)
| 91 | |
| 92 | |
| 93 | def check_example_file(orc_path, expected_df, need_fix=False): |
| 94 | """ |
| 95 | Check a ORC file against the expected columns dictionary. |
| 96 | """ |
| 97 | from pyarrow import orc |
| 98 | |
| 99 | orc_file = orc.ORCFile(orc_path) |
| 100 | # Exercise ORCFile.read() |
| 101 | table = orc_file.read() |
| 102 | assert isinstance(table, pa.Table) |
| 103 | table.validate() |
| 104 | |
| 105 | # This workaround needed because of ARROW-3080 |
| 106 | orc_df = pd.DataFrame(table.to_pydict()) |
| 107 | |
| 108 | assert set(expected_df.columns) == set(orc_df.columns) |
| 109 | |
| 110 | # reorder columns if necessary |
| 111 | if not orc_df.columns.equals(expected_df.columns): |
| 112 | expected_df = expected_df.reindex(columns=orc_df.columns) |
| 113 | |
| 114 | if need_fix: |
| 115 | fix_example_values(orc_df, expected_df) |
| 116 | |
| 117 | check_example_values(orc_df, expected_df) |
| 118 | # Exercise ORCFile.read_stripe() |
| 119 | json_pos = 0 |
| 120 | for i in range(orc_file.nstripes): |
| 121 | batch = orc_file.read_stripe(i) |
| 122 | check_example_values(pd.DataFrame(batch.to_pydict()), |
| 123 | expected_df, |
| 124 | start=json_pos, |
| 125 | stop=json_pos + len(batch)) |
| 126 | json_pos += len(batch) |
| 127 | assert json_pos == orc_file.nrows |
| 128 | |
| 129 | |
| 130 | @pytest.mark.pandas |
no test coverage detected