(tempdir, dataset_reader)
| 3292 | |
| 3293 | @pytest.mark.parquet |
| 3294 | def test_specified_schema(tempdir, dataset_reader): |
| 3295 | table = pa.table({'a': [1, 2, 3], 'b': [.1, .2, .3]}) |
| 3296 | pq.write_table(table, tempdir / "data.parquet") |
| 3297 | |
| 3298 | def _check_dataset(schema, expected, expected_schema=None): |
| 3299 | dataset = ds.dataset(str(tempdir / "data.parquet"), schema=schema) |
| 3300 | if expected_schema is not None: |
| 3301 | assert dataset.schema.equals(expected_schema) |
| 3302 | else: |
| 3303 | assert dataset.schema.equals(schema) |
| 3304 | result = dataset_reader.to_table(dataset) |
| 3305 | assert result.equals(expected) |
| 3306 | |
| 3307 | # no schema specified |
| 3308 | schema = None |
| 3309 | expected = table |
| 3310 | _check_dataset(schema, expected, expected_schema=table.schema) |
| 3311 | |
| 3312 | # identical schema specified |
| 3313 | schema = table.schema |
| 3314 | expected = table |
| 3315 | _check_dataset(schema, expected) |
| 3316 | |
| 3317 | # Specifying schema with change column order |
| 3318 | schema = pa.schema([('b', 'float64'), ('a', 'int64')]) |
| 3319 | expected = pa.table([[.1, .2, .3], [1, 2, 3]], names=['b', 'a']) |
| 3320 | _check_dataset(schema, expected) |
| 3321 | |
| 3322 | # Specifying schema with missing column |
| 3323 | schema = pa.schema([('a', 'int64')]) |
| 3324 | expected = pa.table([[1, 2, 3]], names=['a']) |
| 3325 | _check_dataset(schema, expected) |
| 3326 | |
| 3327 | # Specifying schema with additional column |
| 3328 | schema = pa.schema([('a', 'int64'), ('c', 'int32')]) |
| 3329 | expected = pa.table([[1, 2, 3], |
| 3330 | pa.array([None, None, None], type='int32')], |
| 3331 | names=['a', 'c']) |
| 3332 | _check_dataset(schema, expected) |
| 3333 | |
| 3334 | # Specifying with differing field types |
| 3335 | schema = pa.schema([('a', 'int32'), ('b', 'float64')]) |
| 3336 | dataset = ds.dataset(str(tempdir / "data.parquet"), schema=schema) |
| 3337 | expected = pa.table([table['a'].cast('int32'), |
| 3338 | table['b']], |
| 3339 | names=['a', 'b']) |
| 3340 | _check_dataset(schema, expected) |
| 3341 | |
| 3342 | # Specifying with incompatible schema |
| 3343 | schema = pa.schema([('a', pa.list_(pa.int32())), ('b', 'float64')]) |
| 3344 | dataset = ds.dataset(str(tempdir / "data.parquet"), schema=schema) |
| 3345 | assert dataset.schema.equals(schema) |
| 3346 | with pytest.raises(NotImplementedError, |
| 3347 | match='Unsupported cast from int64 to list'): |
| 3348 | dataset_reader.to_table(dataset) |
| 3349 | |
| 3350 | |
| 3351 | @pytest.mark.parquet |
nothing calls this directly
no test coverage detected