| 335 | |
| 336 | @pytest.mark.parquet |
| 337 | def test_filesystem_dataset(mockfs): |
| 338 | schema = pa.schema([ |
| 339 | pa.field('const', pa.int64()) |
| 340 | ]) |
| 341 | file_format = ds.ParquetFileFormat() |
| 342 | paths = ['subdir/1/xxx/file0.parquet', 'subdir/2/yyy/file1.parquet'] |
| 343 | partitions = [ds.field('part') == x for x in range(1, 3)] |
| 344 | fragments = [file_format.make_fragment(path, mockfs, part) |
| 345 | for path, part in zip(paths, partitions)] |
| 346 | root_partition = ds.field('level') == ds.scalar(1337) |
| 347 | |
| 348 | dataset_from_fragments = ds.FileSystemDataset( |
| 349 | fragments, schema=schema, format=file_format, |
| 350 | filesystem=mockfs, root_partition=root_partition, |
| 351 | ) |
| 352 | dataset_from_paths = ds.FileSystemDataset.from_paths( |
| 353 | paths, schema=schema, format=file_format, filesystem=mockfs, |
| 354 | partitions=partitions, root_partition=root_partition, |
| 355 | ) |
| 356 | |
| 357 | for dataset in [dataset_from_fragments, dataset_from_paths]: |
| 358 | assert isinstance(dataset, ds.FileSystemDataset) |
| 359 | assert isinstance(dataset.format, ds.ParquetFileFormat) |
| 360 | assert dataset.partition_expression.equals(root_partition) |
| 361 | assert set(dataset.files) == set(paths) |
| 362 | |
| 363 | fragments = list(dataset.get_fragments()) |
| 364 | for fragment, partition, path in zip(fragments, partitions, paths): |
| 365 | assert fragment.partition_expression.equals(partition) |
| 366 | assert fragment.path == path |
| 367 | assert isinstance(fragment.format, ds.ParquetFileFormat) |
| 368 | assert isinstance(fragment, ds.ParquetFileFragment) |
| 369 | assert fragment.row_groups == [0] |
| 370 | assert fragment.num_row_groups == 1 |
| 371 | |
| 372 | row_group_fragments = list(fragment.split_by_row_group()) |
| 373 | assert fragment.num_row_groups == len(row_group_fragments) == 1 |
| 374 | assert isinstance(row_group_fragments[0], ds.ParquetFileFragment) |
| 375 | assert row_group_fragments[0].path == path |
| 376 | assert row_group_fragments[0].row_groups == [0] |
| 377 | assert row_group_fragments[0].num_row_groups == 1 |
| 378 | |
| 379 | fragments = list(dataset.get_fragments(filter=ds.field("const") == 0)) |
| 380 | assert len(fragments) == 2 |
| 381 | |
| 382 | # the root_partition keyword has a default |
| 383 | dataset = ds.FileSystemDataset( |
| 384 | fragments, schema=schema, format=file_format, filesystem=mockfs |
| 385 | ) |
| 386 | assert dataset.partition_expression.equals(ds.scalar(True)) |
| 387 | |
| 388 | # from_paths partitions have defaults |
| 389 | dataset = ds.FileSystemDataset.from_paths( |
| 390 | paths, schema=schema, format=file_format, filesystem=mockfs |
| 391 | ) |
| 392 | assert dataset.partition_expression.equals(ds.scalar(True)) |
| 393 | for fragment in dataset.get_fragments(): |
| 394 | assert fragment.partition_expression.equals(ds.scalar(True)) |