| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import pytest |
| |
| import pyarrow as pa |
| from pyarrow import fs |
| |
| try: |
| import pyarrow.parquet as pq |
| from pyarrow.tests.parquet.common import (_read_table, _test_dataframe, |
| _test_table, _range_integers) |
| except ImportError: |
| pq = None |
| |
| |
| try: |
| import pandas as pd |
| import pandas.testing as tm |
| |
| except ImportError: |
| pd = tm = None |
| |
| |
| # Marks all of the tests in this module |
| # Ignore these with pytest ... -m 'not parquet' |
| pytestmark = pytest.mark.parquet |
| |
| |
| @pytest.mark.pandas |
| def test_parquet_incremental_file_build(tempdir): |
| df = _test_dataframe(100) |
| df['unique_id'] = 0 |
| |
| arrow_table = pa.Table.from_pandas(df, preserve_index=False) |
| out = pa.BufferOutputStream() |
| |
| writer = pq.ParquetWriter(out, arrow_table.schema, version='2.6') |
| |
| frames = [] |
| for i in range(10): |
| df['unique_id'] = i |
| arrow_table = pa.Table.from_pandas(df, preserve_index=False) |
| writer.write_table(arrow_table) |
| |
| frames.append(df.copy()) |
| |
| writer.close() |
| |
| buf = out.getvalue() |
| result = _read_table(pa.BufferReader(buf)) |
| |
| expected = pd.concat(frames, ignore_index=True) |
| tm.assert_frame_equal(result.to_pandas(), expected) |
| |
| |
| def test_validate_schema_write_table(tempdir): |
| # ARROW-2926 |
| simple_fields = [ |
| pa.field('POS', pa.uint32()), |
| pa.field('desc', pa.string()) |
| ] |
| |
| simple_schema = pa.schema(simple_fields) |
| |
| # simple_table schema does not match simple_schema |
| simple_from_array = [pa.array([1]), pa.array(['bla'])] |
| simple_table = pa.Table.from_arrays(simple_from_array, ['POS', 'desc']) |
| |
| path = tempdir / 'simple_validate_schema.parquet' |
| |
| with pq.ParquetWriter(path, simple_schema, |
| version='2.6', |
| compression='snappy', flavor='spark') as w: |
| with pytest.raises(ValueError): |
| w.write_table(simple_table) |
| |
| |
| def test_parquet_invalid_writer(tempdir): |
| # avoid segfaults with invalid construction |
| with pytest.raises(TypeError): |
| some_schema = pa.schema([pa.field("x", pa.int32())]) |
| pq.ParquetWriter(None, some_schema) |
| |
| with pytest.raises(TypeError): |
| pq.ParquetWriter(tempdir / "some_path", None) |
| |
| |
| @pytest.mark.pandas |
| def test_parquet_writer_context_obj(tempdir): |
| df = _test_dataframe(100) |
| df['unique_id'] = 0 |
| |
| arrow_table = pa.Table.from_pandas(df, preserve_index=False) |
| out = pa.BufferOutputStream() |
| |
| with pq.ParquetWriter(out, arrow_table.schema, version='2.6') as writer: |
| |
| frames = [] |
| for i in range(10): |
| df['unique_id'] = i |
| arrow_table = pa.Table.from_pandas(df, preserve_index=False) |
| writer.write_table(arrow_table) |
| |
| frames.append(df.copy()) |
| |
| buf = out.getvalue() |
| result = _read_table(pa.BufferReader(buf)) |
| |
| expected = pd.concat(frames, ignore_index=True) |
| tm.assert_frame_equal(result.to_pandas(), expected) |
| |
| |
| @pytest.mark.pandas |
| def test_parquet_writer_context_obj_with_exception(tempdir): |
| df = _test_dataframe(100) |
| df['unique_id'] = 0 |
| |
| arrow_table = pa.Table.from_pandas(df, preserve_index=False) |
| out = pa.BufferOutputStream() |
| error_text = 'Artificial Error' |
| |
| try: |
| with pq.ParquetWriter(out, |
| arrow_table.schema, |
| version='2.6') as writer: |
| |
| frames = [] |
| for i in range(10): |
| df['unique_id'] = i |
| arrow_table = pa.Table.from_pandas(df, preserve_index=False) |
| writer.write_table(arrow_table) |
| frames.append(df.copy()) |
| if i == 5: |
| raise ValueError(error_text) |
| except Exception as e: |
| assert str(e) == error_text |
| |
| buf = out.getvalue() |
| result = _read_table(pa.BufferReader(buf)) |
| |
| expected = pd.concat(frames, ignore_index=True) |
| tm.assert_frame_equal(result.to_pandas(), expected) |
| |
| |
| @pytest.mark.pandas |
| @pytest.mark.parametrize("filesystem", [ |
| None, |
| fs.LocalFileSystem(), |
| ]) |
| def test_parquet_writer_write_wrappers(tempdir, filesystem): |
| df = _test_dataframe(100) |
| table = pa.Table.from_pandas(df, preserve_index=False) |
| batch = pa.RecordBatch.from_pandas(df, preserve_index=False) |
| path_table = str(tempdir / 'data_table.parquet') |
| path_batch = str(tempdir / 'data_batch.parquet') |
| |
| with pq.ParquetWriter( |
| path_table, table.schema, filesystem=filesystem, version='2.6' |
| ) as writer: |
| writer.write_table(table) |
| |
| result = _read_table(path_table).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| with pq.ParquetWriter( |
| path_batch, table.schema, filesystem=filesystem, version='2.6' |
| ) as writer: |
| writer.write_batch(batch) |
| |
| result = _read_table(path_batch).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| with pq.ParquetWriter( |
| path_table, table.schema, filesystem=filesystem, version='2.6' |
| ) as writer: |
| writer.write(table) |
| |
| result = _read_table(path_table).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| with pq.ParquetWriter( |
| path_batch, table.schema, filesystem=filesystem, version='2.6' |
| ) as writer: |
| writer.write(batch) |
| |
| result = _read_table(path_batch).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| |
| @pytest.mark.large_memory |
| @pytest.mark.pandas |
| def test_parquet_writer_chunk_size(tempdir): |
| default_chunk_size = 1024 * 1024 |
| abs_max_chunk_size = 64 * 1024 * 1024 |
| |
| def check_chunk_size(data_size, chunk_size, expect_num_chunks): |
| table = pa.Table.from_arrays([ |
| _range_integers(data_size, 'b') |
| ], names=['x']) |
| if chunk_size is None: |
| pq.write_table(table, tempdir / 'test.parquet') |
| else: |
| pq.write_table(table, tempdir / 'test.parquet', row_group_size=chunk_size) |
| metadata = pq.read_metadata(tempdir / 'test.parquet') |
| expected_chunk_size = default_chunk_size if chunk_size is None else chunk_size |
| assert metadata.num_row_groups == expect_num_chunks |
| latched_chunk_size = min(expected_chunk_size, abs_max_chunk_size) |
| # First chunks should be full size |
| for chunk_idx in range(expect_num_chunks - 1): |
| assert metadata.row_group(chunk_idx).num_rows == latched_chunk_size |
| # Last chunk may be smaller |
| remainder = data_size - (expected_chunk_size * (expect_num_chunks - 1)) |
| if remainder == 0: |
| assert metadata.row_group( |
| expect_num_chunks - 1).num_rows == latched_chunk_size |
| else: |
| assert metadata.row_group(expect_num_chunks - 1).num_rows == remainder |
| |
| check_chunk_size(default_chunk_size * 2, default_chunk_size - 100, 3) |
| check_chunk_size(default_chunk_size * 2, default_chunk_size, 2) |
| check_chunk_size(default_chunk_size * 2, default_chunk_size + 100, 2) |
| check_chunk_size(default_chunk_size + 100, default_chunk_size + 100, 1) |
| # Even though the chunk size requested is large enough it will be capped |
| # by the absolute max chunk size |
| check_chunk_size(abs_max_chunk_size * 2, abs_max_chunk_size * 2, 2) |
| |
| # These tests don't pass a chunk_size to write_table and so the chunk size |
| # should be default_chunk_size |
| check_chunk_size(default_chunk_size, None, 1) |
| check_chunk_size(default_chunk_size + 1, None, 2) |
| |
| |
| @pytest.mark.pandas |
| @pytest.mark.parametrize("filesystem", [ |
| None, |
| fs.LocalFileSystem(), |
| ]) |
| def test_parquet_writer_filesystem_local(tempdir, filesystem): |
| df = _test_dataframe(100) |
| table = pa.Table.from_pandas(df, preserve_index=False) |
| path = str(tempdir / 'data.parquet') |
| |
| with pq.ParquetWriter( |
| path, table.schema, filesystem=filesystem, version='2.6' |
| ) as writer: |
| writer.write_table(table) |
| |
| result = _read_table(path).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| |
| @pytest.mark.pandas |
| @pytest.mark.s3 |
| def test_parquet_writer_filesystem_s3(s3_example_fs): |
| df = _test_dataframe(100) |
| table = pa.Table.from_pandas(df, preserve_index=False) |
| |
| fs, uri, path = s3_example_fs |
| |
| with pq.ParquetWriter( |
| path, table.schema, filesystem=fs, version='2.6' |
| ) as writer: |
| writer.write_table(table) |
| |
| result = _read_table(uri).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| |
| @pytest.mark.pandas |
| @pytest.mark.s3 |
| def test_parquet_writer_filesystem_s3_uri(s3_example_fs): |
| df = _test_dataframe(100) |
| table = pa.Table.from_pandas(df, preserve_index=False) |
| |
| fs, uri, path = s3_example_fs |
| |
| with pq.ParquetWriter(uri, table.schema, version='2.6') as writer: |
| writer.write_table(table) |
| |
| result = _read_table(path, filesystem=fs).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| |
| @pytest.mark.pandas |
| @pytest.mark.s3 |
| def test_parquet_writer_filesystem_s3fs(s3_example_s3fs): |
| df = _test_dataframe(100) |
| table = pa.Table.from_pandas(df, preserve_index=False) |
| |
| fs, directory = s3_example_s3fs |
| path = directory + "/test.parquet" |
| |
| with pq.ParquetWriter( |
| path, table.schema, filesystem=fs, version='2.6' |
| ) as writer: |
| writer.write_table(table) |
| |
| result = _read_table(path, filesystem=fs).to_pandas() |
| tm.assert_frame_equal(result, df) |
| |
| |
| @pytest.mark.numpy |
| def test_parquet_writer_filesystem_buffer_raises(): |
| table = _test_table(100) |
| filesystem = fs.LocalFileSystem() |
| |
| # Should raise ValueError when filesystem is passed with file-like object |
| with pytest.raises(ValueError, match="specified path is file-like"): |
| pq.ParquetWriter( |
| pa.BufferOutputStream(), table.schema, filesystem=filesystem |
| ) |
| |
| |
| def test_parquet_writer_store_schema(tempdir): |
| table = pa.table({'a': [1, 2, 3]}) |
| |
| # default -> write schema information |
| path1 = tempdir / 'test_with_schema.parquet' |
| with pq.ParquetWriter(path1, table.schema) as writer: |
| writer.write_table(table) |
| |
| meta = pq.read_metadata(path1) |
| assert b'ARROW:schema' in meta.metadata |
| assert meta.metadata[b'ARROW:schema'] |
| |
| # disable adding schema information |
| path2 = tempdir / 'test_without_schema.parquet' |
| with pq.ParquetWriter(path2, table.schema, store_schema=False) as writer: |
| writer.write_table(table) |
| |
| meta = pq.read_metadata(path2) |
| assert meta.metadata is None |
| |
| |
| def test_parquet_writer_append_key_value_metadata(tempdir): |
| table = pa.Table.from_arrays([pa.array([], type='int32')], ['f0']) |
| path = tempdir / 'metadata.parquet' |
| |
| with pq.ParquetWriter(path, table.schema) as writer: |
| writer.write_table(table) |
| writer.add_key_value_metadata({'key1': '1', 'key2': 'x'}) |
| writer.add_key_value_metadata({'key2': '2', 'key3': '3'}) |
| reader = pq.ParquetFile(path) |
| metadata = reader.metadata.metadata |
| assert metadata[b'key1'] == b'1' |
| assert metadata[b'key2'] == b'2' |
| assert metadata[b'key3'] == b'3' |
| |
| |
| def test_parquet_content_defined_chunking(tempdir): |
| table = pa.table({'a': range(100_000)}) |
| |
| # use PLAIN encoding because we compare the overall size of the row groups |
| # which would vary depending on the encoding making the assertions wrong |
| pq.write_table(table, tempdir / 'unchunked.parquet', |
| use_dictionary=False, |
| column_encoding="PLAIN") |
| pq.write_table(table, tempdir / 'chunked-default.parquet', |
| use_dictionary=False, |
| column_encoding="PLAIN", |
| use_content_defined_chunking=True) |
| pq.write_table(table, tempdir / 'chunked-custom.parquet', |
| use_dictionary=False, |
| column_encoding="PLAIN", |
| use_content_defined_chunking={"min_chunk_size": 32_768, |
| "max_chunk_size": 65_536}) |
| |
| # the data must be the same |
| unchunked = pq.read_table(tempdir / 'unchunked.parquet') |
| chunked_default = pq.read_table(tempdir / 'chunked-default.parquet') |
| chunked_custom = pq.read_table(tempdir / 'chunked-custom.parquet') |
| assert unchunked.equals(chunked_default) |
| assert unchunked.equals(chunked_custom) |
| |
| # number of row groups and their sizes are not affected by content defined chunking |
| unchunked_metadata = pq.read_metadata(tempdir / 'unchunked.parquet') |
| chunked_default_metadata = pq.read_metadata(tempdir / 'chunked-default.parquet') |
| chunked_custom_metadata = pq.read_metadata(tempdir / 'chunked-custom.parquet') |
| |
| assert unchunked_metadata.num_row_groups == chunked_default_metadata.num_row_groups |
| assert unchunked_metadata.num_row_groups == chunked_custom_metadata.num_row_groups |
| |
| for i in range(unchunked_metadata.num_row_groups): |
| rg_unchunked = unchunked_metadata.row_group(i) |
| rg_chunked_default = chunked_default_metadata.row_group(i) |
| rg_chunked_custom = chunked_custom_metadata.row_group(i) |
| assert rg_unchunked.num_rows == rg_chunked_default.num_rows |
| assert rg_unchunked.num_rows == rg_chunked_custom.num_rows |
| # since PageReader is not exposed we cannot inspect the page sizes |
| # so just check that the total byte size is different |
| assert rg_unchunked.total_byte_size < rg_chunked_default.total_byte_size |
| assert rg_unchunked.total_byte_size < rg_chunked_custom.total_byte_size |
| assert rg_chunked_default.total_byte_size < rg_chunked_custom.total_byte_size |
| |
| |
| def test_parquet_content_defined_chunking_parameters(tempdir): |
| table = pa.table({'a': range(100)}) |
| path = tempdir / 'chunked-invalid.parquet' |
| |
| # it raises OSError, not ideal but this is how parquet exceptions are handled |
| # currently |
| msg = "max_chunk_size must be greater than min_chunk_size" |
| with pytest.raises(Exception, match=msg): |
| cdc_options = {"min_chunk_size": 65_536, "max_chunk_size": 32_768} |
| pq.write_table(table, path, use_content_defined_chunking=cdc_options) |
| |
| cases = [ |
| ( |
| {"min_chunk_size": 64 * 1024, "unknown_option": True}, |
| "Unknown options in 'use_content_defined_chunking': {'unknown_option'}" |
| ), |
| ( |
| {"min_chunk_size": 64 * 1024}, |
| "Missing options in 'use_content_defined_chunking': {'max_chunk_size'}" |
| ), |
| ( |
| {"max_chunk_size": 64 * 1024}, |
| "Missing options in 'use_content_defined_chunking': {'min_chunk_size'}" |
| ) |
| ] |
| for cdc_options, msg in cases: |
| with pytest.raises(ValueError, match=msg): |
| pq.write_table(table, path, use_content_defined_chunking=cdc_options) |
| |
| # using the default parametrization |
| pq.write_table(table, path, use_content_defined_chunking=True) |
| |
| # using min_chunk_size and max_chunk_size |
| cdc_options = {"min_chunk_size": 32_768, "max_chunk_size": 65_536} |
| pq.write_table(table, path, use_content_defined_chunking=cdc_options) |
| |
| # using min_chunk_size, max_chunk_size and norm_level |
| cdc_options = {"min_chunk_size": 32_768, "max_chunk_size": 65_536, "norm_level": 1} |
| pq.write_table(table, path, use_content_defined_chunking=cdc_options) |
| |
| |
| @pytest.mark.parametrize("time_type, time_unit", [ |
| (pa.time32, "s"), |
| (pa.time32, "ms"), |
| (pa.time64, "us"), |
| (pa.time64, "ns"), |
| ]) |
| @pytest.mark.parametrize("utc_flag_val", [False, True]) |
| def test_arrow_writer_props_time_adjusted_to_utc( |
| tempdir, |
| utc_flag_val, |
| time_type, |
| time_unit, |
| ): |
| # GH-47441 |
| filename = tempdir / "time_adjusted_to_utc.parquet" |
| |
| time_values = [0, 123, 10_000, 86_399] |
| |
| table = pa.table({ |
| "time_col": pa.array(time_values, type=time_type(time_unit)), |
| }) |
| |
| schema = pa.schema([ |
| ("time_col", time_type(time_unit)), |
| ]) |
| |
| with pq.ParquetWriter( |
| where=filename, |
| schema=schema, |
| write_time_adjusted_to_utc=utc_flag_val, |
| ) as writer: |
| writer.write_table(table) |
| |
| result = pq.read_table(filename, schema=schema) |
| |
| result.validate(full=True) |
| |
| assert result.equals(table) |
| |
| |
| @pytest.mark.parametrize( |
| "max_rows_per_page", |
| [1, 10, 100, 1_000, None], |
| ) |
| def test_writer_props_max_rows_per_page(tempdir, max_rows_per_page): |
| # GH-48096 |
| filename = tempdir / "max_rows_per_page.parquet" |
| |
| table = pa.table({ |
| "x": pa.array([1, 2, 3, 4, 5, 6, 7], type=pa.int8()), |
| "y": pa.array([11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0], type=pa.float16()), |
| }) |
| |
| schema = pa.schema([ |
| ("x", pa.int8()), |
| ("y", pa.float16()), |
| ]) |
| |
| with pq.ParquetWriter( |
| where=filename, |
| schema=schema, |
| max_rows_per_page=max_rows_per_page, |
| ) as writer: |
| writer.write_table(table) |
| |
| result = pq.read_table(filename, schema=schema) |
| |
| result.validate(full=True) |
| |
| assert result.equals(table) |
| |
| |
| def test_writer_props_max_rows_per_page_file_size(tempdir): |
| # GH-48096 |
| table = pa.table({ |
| "x": pa.array(range(1_000_000)) |
| }) |
| |
| local = fs.LocalFileSystem() |
| file_infos = [] |
| |
| for max_rows in (1_000, 10_000): |
| path = f"{tempdir}/max_rows_per_page_{max_rows}.parquet" |
| |
| with pq.ParquetWriter( |
| where=path, |
| schema=table.schema, |
| max_rows_per_page=max_rows, |
| ) as writer: |
| writer.write_table(table) |
| |
| file_infos.append(local.get_file_info(path)) |
| |
| # A smaller maximum rows parameter should produce a larger file |
| assert file_infos[0].size > file_infos[1].size |
