(tempdir, dataset_reader, pickle_module)
| 1447 | |
| 1448 | @pytest.mark.parquet |
| 1449 | def test_fragments_reconstruct(tempdir, dataset_reader, pickle_module): |
| 1450 | table, dataset = _create_dataset_for_fragments(tempdir) |
| 1451 | |
| 1452 | def assert_yields_projected(fragment, row_slice, |
| 1453 | columns=None, filter=None): |
| 1454 | actual = fragment.to_table( |
| 1455 | schema=table.schema, columns=columns, filter=filter) |
| 1456 | column_names = columns if columns else table.column_names |
| 1457 | assert actual.column_names == column_names |
| 1458 | |
| 1459 | expected = table.slice(*row_slice).select(column_names) |
| 1460 | assert actual.equals(expected) |
| 1461 | |
| 1462 | fragment = list(dataset.get_fragments())[0] |
| 1463 | parquet_format = fragment.format |
| 1464 | |
| 1465 | # test pickle roundtrip |
| 1466 | pickled_fragment = pickle_module.loads(pickle_module.dumps(fragment)) |
| 1467 | assert dataset_reader.to_table( |
| 1468 | pickled_fragment) == dataset_reader.to_table(fragment) |
| 1469 | |
| 1470 | # manually re-construct a fragment, with explicit schema |
| 1471 | new_fragment = parquet_format.make_fragment( |
| 1472 | fragment.path, fragment.filesystem, |
| 1473 | partition_expression=fragment.partition_expression) |
| 1474 | assert dataset_reader.to_table(new_fragment).equals( |
| 1475 | dataset_reader.to_table(fragment)) |
| 1476 | assert_yields_projected(new_fragment, (0, 4)) |
| 1477 | |
| 1478 | # filter / column projection, inspected schema |
| 1479 | new_fragment = parquet_format.make_fragment( |
| 1480 | fragment.path, fragment.filesystem, |
| 1481 | partition_expression=fragment.partition_expression) |
| 1482 | assert_yields_projected(new_fragment, (0, 2), filter=ds.field('f1') < 2) |
| 1483 | |
| 1484 | # filter requiring cast / column projection, inspected schema |
| 1485 | new_fragment = parquet_format.make_fragment( |
| 1486 | fragment.path, fragment.filesystem, |
| 1487 | partition_expression=fragment.partition_expression) |
| 1488 | assert_yields_projected(new_fragment, (0, 2), |
| 1489 | columns=['f1'], filter=ds.field('f1') < 2.0) |
| 1490 | |
| 1491 | # filter on the partition column |
| 1492 | new_fragment = parquet_format.make_fragment( |
| 1493 | fragment.path, fragment.filesystem, |
| 1494 | partition_expression=fragment.partition_expression) |
| 1495 | assert_yields_projected(new_fragment, (0, 4), |
| 1496 | filter=ds.field('part') == 'a') |
| 1497 | |
| 1498 | # Fragments don't contain the partition's columns if not provided to the |
| 1499 | # `to_table(schema=...)` method. |
| 1500 | pattern = (r'No match for FieldRef.Name\(part\) in ' + |
| 1501 | fragment.physical_schema.to_string(False, False, False)) |
| 1502 | with pytest.raises(ValueError, match=pattern): |
| 1503 | new_fragment = parquet_format.make_fragment( |
| 1504 | fragment.path, fragment.filesystem, |
| 1505 | partition_expression=fragment.partition_expression) |
| 1506 | dataset_reader.to_table(new_fragment, filter=ds.field('part') == 'a') |
nothing calls this directly
no test coverage detected