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