(dataset, dataset_reader)
| 429 | |
| 430 | @pytest.mark.parquet |
| 431 | def test_dataset(dataset, dataset_reader): |
| 432 | assert isinstance(dataset, ds.Dataset) |
| 433 | assert isinstance(dataset.schema, pa.Schema) |
| 434 | |
| 435 | non_boolean_expr = ds.field('i64') |
| 436 | with pytest.raises(TypeError, match="must evaluate to bool"): |
| 437 | dataset.to_table(filter=non_boolean_expr) |
| 438 | |
| 439 | non_boolean_expr2 = ds.field('i64') + 1 |
| 440 | with pytest.raises(TypeError, match="must evaluate to bool"): |
| 441 | dataset.to_table(filter=non_boolean_expr2) |
| 442 | |
| 443 | expected_i64 = pa.array([0, 1, 2, 3, 4], type=pa.int64()) |
| 444 | expected_f64 = pa.array([0, 1, 2, 3, 4], type=pa.float64()) |
| 445 | |
| 446 | for batch in dataset_reader.to_batches(dataset): |
| 447 | assert isinstance(batch, pa.RecordBatch) |
| 448 | assert batch.column(0).equals(expected_i64) |
| 449 | assert batch.column(1).equals(expected_f64) |
| 450 | |
| 451 | for batch in dataset_reader.scanner(dataset).scan_batches(): |
| 452 | assert isinstance(batch, ds.TaggedRecordBatch) |
| 453 | assert isinstance(batch.fragment, ds.Fragment) |
| 454 | |
| 455 | table = dataset_reader.to_table(dataset) |
| 456 | assert isinstance(table, pa.Table) |
| 457 | assert len(table) == 10 |
| 458 | |
| 459 | condition = ds.field('i64') == 1 |
| 460 | result = dataset.to_table(use_threads=True, filter=condition) |
| 461 | # Don't rely on the scanning order |
| 462 | result = result.sort_by('group').to_pydict() |
| 463 | |
| 464 | assert result['i64'] == [1, 1] |
| 465 | assert result['f64'] == [1., 1.] |
| 466 | assert sorted(result['group']) == [1, 2] |
| 467 | assert sorted(result['key']) == ['xxx', 'yyy'] |
| 468 | |
| 469 | # Filtering on a nested field ref |
| 470 | condition = ds.field(('struct', 'b')) == '1' |
| 471 | result = dataset.to_table(use_threads=True, filter=condition) |
| 472 | result = result.sort_by('group').to_pydict() |
| 473 | |
| 474 | assert result['i64'] == [1, 4, 1, 4] |
| 475 | assert result['f64'] == [1.0, 4.0, 1.0, 4.0] |
| 476 | assert result['group'] == [1, 1, 2, 2] |
| 477 | assert result['key'] == ['xxx', 'xxx', 'yyy', 'yyy'] |
| 478 | |
| 479 | # Projecting on a nested field ref expression |
| 480 | projection = { |
| 481 | 'i64': ds.field('i64'), |
| 482 | 'f64': ds.field('f64'), |
| 483 | 'new': ds.field(('struct', 'b')) == '1', |
| 484 | } |
| 485 | result = dataset.to_table(use_threads=True, columns=projection) |
| 486 | result = result.sort_by('i64').to_pydict() |
| 487 | |
| 488 | assert list(result) == ['i64', 'f64', 'new'] |
nothing calls this directly
no test coverage detected