(mockfs, paths_or_selector, pre_buffer)
| 1139 | @pytest.mark.parametrize('pre_buffer', [False, True]) |
| 1140 | @pytest.mark.parquet |
| 1141 | def test_filesystem_factory(mockfs, paths_or_selector, pre_buffer): |
| 1142 | format = ds.ParquetFileFormat( |
| 1143 | read_options=ds.ParquetReadOptions(dictionary_columns={"str"}), |
| 1144 | pre_buffer=pre_buffer |
| 1145 | ) |
| 1146 | |
| 1147 | options = ds.FileSystemFactoryOptions('subdir') |
| 1148 | options.partitioning = ds.DirectoryPartitioning( |
| 1149 | pa.schema([ |
| 1150 | pa.field('group', pa.int32()), |
| 1151 | pa.field('key', pa.string()) |
| 1152 | ]) |
| 1153 | ) |
| 1154 | assert options.partition_base_dir == 'subdir' |
| 1155 | assert options.selector_ignore_prefixes == ['.', '_'] |
| 1156 | assert options.exclude_invalid_files is False |
| 1157 | |
| 1158 | factory = ds.FileSystemDatasetFactory( |
| 1159 | mockfs, paths_or_selector, format, options |
| 1160 | ) |
| 1161 | inspected_schema = factory.inspect() |
| 1162 | |
| 1163 | assert factory.inspect().equals(pa.schema([ |
| 1164 | pa.field('i64', pa.int64()), |
| 1165 | pa.field('f64', pa.float64()), |
| 1166 | pa.field('str', pa.dictionary(pa.int32(), pa.string())), |
| 1167 | pa.field('const', pa.int64()), |
| 1168 | pa.field('struct', pa.struct({'a': pa.int64(), |
| 1169 | 'b': pa.string()})), |
| 1170 | pa.field('group', pa.int32()), |
| 1171 | pa.field('key', pa.string()), |
| 1172 | ]), check_metadata=False) |
| 1173 | |
| 1174 | assert isinstance(factory.inspect_schemas(), list) |
| 1175 | assert isinstance(factory.finish(inspected_schema), |
| 1176 | ds.FileSystemDataset) |
| 1177 | assert factory.root_partition.equals(ds.scalar(True)) |
| 1178 | |
| 1179 | dataset = factory.finish() |
| 1180 | assert isinstance(dataset, ds.FileSystemDataset) |
| 1181 | |
| 1182 | scanner = dataset.scanner() |
| 1183 | expected_i64 = pa.array([0, 1, 2, 3, 4], type=pa.int64()) |
| 1184 | expected_f64 = pa.array([0, 1, 2, 3, 4], type=pa.float64()) |
| 1185 | expected_str = pa.DictionaryArray.from_arrays( |
| 1186 | pa.array([0, 1, 2, 3, 4], type=pa.int32()), |
| 1187 | pa.array("0 1 2 3 4".split(), type=pa.string()) |
| 1188 | ) |
| 1189 | expected_struct = pa.array([{'a': i % 3, 'b': str(i % 3)} |
| 1190 | for i in range(5)]) |
| 1191 | iterator = scanner.scan_batches() |
| 1192 | for (batch, fragment), group, key in zip(iterator, [1, 2], ['xxx', 'yyy']): |
| 1193 | expected_group = pa.array([group] * 5, type=pa.int32()) |
| 1194 | expected_key = pa.array([key] * 5, type=pa.string()) |
| 1195 | expected_const = pa.array([group - 1] * 5, type=pa.int64()) |
| 1196 | # Can't compare or really introspect expressions from Python |
| 1197 | assert fragment.partition_expression is not None |
| 1198 | assert batch.num_columns == 7 |
nothing calls this directly
no test coverage detected