| 658 | |
| 659 | @pytest.mark.parquet |
| 660 | def test_take(dataset, dataset_reader): |
| 661 | fragment = next(dataset.get_fragments()) |
| 662 | for indices in [[1, 3], pa.array([1, 3])]: |
| 663 | expected = dataset_reader.to_table(fragment).take(indices) |
| 664 | assert dataset_reader.take(fragment, indices) == expected |
| 665 | with pytest.raises(IndexError): |
| 666 | dataset_reader.take(fragment, pa.array([5])) |
| 667 | |
| 668 | for indices in [[1, 7], pa.array([1, 7])]: |
| 669 | assert dataset_reader.take( |
| 670 | dataset, indices) == dataset_reader.to_table(dataset).take(indices) |
| 671 | with pytest.raises(IndexError): |
| 672 | dataset_reader.take(dataset, pa.array([10])) |
| 673 | |
| 674 | |
| 675 | @pytest.mark.parquet |