(dataset, dataset_reader)
| 567 | |
| 568 | @pytest.mark.parquet |
| 569 | def test_scanner(dataset, dataset_reader): |
| 570 | scanner = dataset_reader.scanner( |
| 571 | dataset, memory_pool=pa.default_memory_pool()) |
| 572 | assert isinstance(scanner, ds.Scanner) |
| 573 | |
| 574 | with pytest.raises(pa.ArrowInvalid): |
| 575 | dataset_reader.scanner(dataset, columns=['unknown']) |
| 576 | |
| 577 | scanner = dataset_reader.scanner(dataset, columns=['i64'], |
| 578 | memory_pool=pa.default_memory_pool()) |
| 579 | assert scanner.dataset_schema == dataset.schema |
| 580 | assert scanner.projected_schema == pa.schema([("i64", pa.int64())]) |
| 581 | |
| 582 | assert isinstance(scanner, ds.Scanner) |
| 583 | table = scanner.to_table() |
| 584 | for batch in scanner.to_batches(): |
| 585 | assert batch.schema == scanner.projected_schema |
| 586 | assert batch.num_columns == 1 |
| 587 | assert table == scanner.to_reader().read_all() |
| 588 | |
| 589 | assert table.schema == scanner.projected_schema |
| 590 | for i in range(table.num_rows): |
| 591 | indices = pa.array([i]) |
| 592 | assert table.take(indices) == scanner.take(indices) |
| 593 | with pytest.raises(pa.ArrowIndexError): |
| 594 | scanner.take(pa.array([table.num_rows])) |
| 595 | |
| 596 | assert table.num_rows == scanner.count_rows() |
| 597 | |
| 598 | scanner = dataset_reader.scanner(dataset, columns=['__filename', |
| 599 | '__fragment_index', |
| 600 | '__batch_index', |
| 601 | '__last_in_fragment'], |
| 602 | memory_pool=pa.default_memory_pool()) |
| 603 | table = scanner.to_table() |
| 604 | expected_names = ['__filename', '__fragment_index', |
| 605 | '__batch_index', '__last_in_fragment'] |
| 606 | assert table.column_names == expected_names |
| 607 | |
| 608 | sorted_table = table.sort_by('__fragment_index') |
| 609 | assert sorted_table['__filename'].to_pylist() == ( |
| 610 | ['subdir/1/xxx/file0.parquet'] * 5 + |
| 611 | ['subdir/2/yyy/file1.parquet'] * 5) |
| 612 | assert sorted_table['__fragment_index'].to_pylist() == ([0] * 5 + [1] * 5) |
| 613 | assert sorted_table['__batch_index'].to_pylist() == [0] * 10 |
| 614 | assert sorted_table['__last_in_fragment'].to_pylist() == [True] * 10 |
| 615 | |
| 616 | |
| 617 | @pytest.mark.parquet |
nothing calls this directly
no test coverage detected