| 635 | |
| 636 | @pytest.mark.parquet |
| 637 | def test_head(dataset, dataset_reader): |
| 638 | result = dataset_reader.head(dataset, 0) |
| 639 | assert result == pa.Table.from_batches([], schema=dataset.schema) |
| 640 | |
| 641 | result = dataset_reader.head(dataset, 1, columns=['i64']).to_pydict() |
| 642 | assert result == {'i64': [0]} |
| 643 | |
| 644 | result = dataset_reader.head(dataset, 2, columns=['i64'], |
| 645 | filter=ds.field('i64') > 1).to_pydict() |
| 646 | assert result == {'i64': [2, 3]} |
| 647 | |
| 648 | result = dataset_reader.head(dataset, 1024, columns=['i64']).to_pydict() |
| 649 | assert result == {'i64': list(range(5)) * 2} |
| 650 | |
| 651 | fragment = next(dataset.get_fragments()) |
| 652 | result = fragment.head(1, columns=['i64']).to_pydict() |
| 653 | assert result == {'i64': [0]} |
| 654 | |
| 655 | result = fragment.head(1024, columns=['i64']).to_pydict() |
| 656 | assert result == {'i64': list(range(5))} |
| 657 | |
| 658 | |
| 659 | @pytest.mark.parquet |