You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by we...@apache.org on 2017/10/23 13:24:00 UTC
[arrow] branch master updated: ARROW-571: [Python] Add unit test
for incremental Parquet file building, improve docs
This is an automated email from the ASF dual-hosted git repository.
wesm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 4eb38a2 ARROW-571: [Python] Add unit test for incremental Parquet file building, improve docs
4eb38a2 is described below
commit 4eb38a227d100165c59ed784c1aad09eea41236f
Author: Wes McKinney <we...@twosigma.com>
AuthorDate: Mon Oct 23 09:23:55 2017 -0400
ARROW-571: [Python] Add unit test for incremental Parquet file building, improve docs
This was actually already documented. I improved the docstring for `ParquetWriter` and added a unit test to validate the user API. Added ParquetWriter to the API listing also
Author: Wes McKinney <we...@twosigma.com>
Closes #1218 from wesm/ARROW-571 and squashes the following commits:
e4f44b47 [Wes McKinney] Use py27-compatible class __doc__ setting
bb16d9b5 [Wes McKinney] Add ParquetWriter to API listing
4cda8ab2 [Wes McKinney] Add unit test for incremental Parquet file building
---
python/doc/source/api.rst | 1 +
python/pyarrow/parquet.py | 138 +++++++++++++++++------------
python/pyarrow/tests/test_parquet.py | 29 ++++++
python/pyarrow/tests/test_serialization.py | 44 +++++----
4 files changed, 134 insertions(+), 78 deletions(-)
diff --git a/python/doc/source/api.rst b/python/doc/source/api.rst
index 26ccb98..6bceba3 100644
--- a/python/doc/source/api.rst
+++ b/python/doc/source/api.rst
@@ -312,6 +312,7 @@ Apache Parquet
ParquetDataset
ParquetFile
+ ParquetWriter
read_table
read_metadata
read_pandas
diff --git a/python/pyarrow/parquet.py b/python/pyarrow/parquet.py
index 1584b84..b6a7b12 100644
--- a/python/pyarrow/parquet.py
+++ b/python/pyarrow/parquet.py
@@ -202,17 +202,47 @@ def _sanitize_table(table, new_schema, flavor):
return table
+_parquet_writer_arg_docs = """version : {"1.0", "2.0"}, default "1.0"
+ The Parquet format version, defaults to 1.0
+use_dictionary : bool or list
+ Specify if we should use dictionary encoding in general or only for
+ some columns.
+use_deprecated_int96_timestamps : boolean, default None
+ Write nanosecond resolution timestamps to INT96 Parquet
+ format. Defaults to False unless enabled by flavor argument
+coerce_timestamps : string, default None
+ Cast timestamps a particular resolution.
+ Valid values: {None, 'ms', 'us'}
+compression : str or dict
+ Specify the compression codec, either on a general basis or per-column.
+flavor : {'spark'}, default None
+ Sanitize schema or set other compatibility options for compatibility"""
+
+
class ParquetWriter(object):
- """
- Parameters
- ----------
- where
- schema
- flavor : {'spark', ...}
- Set options for compatibility with a particular reader
- """
- def __init__(self, where, schema, flavor=None, **options):
+ __doc__ = """
+Class for incrementally building a Parquet file for Arrow tables
+
+Parameters
+----------
+where : path or file-like object
+schema : arrow Schema
+{0}
+""".format(_parquet_writer_arg_docs)
+
+ def __init__(self, where, schema, flavor=None,
+ version='1.0',
+ use_dictionary=True,
+ compression='snappy',
+ use_deprecated_int96_timestamps=None, **options):
+ if use_deprecated_int96_timestamps is None:
+ # Use int96 timestamps for Spark
+ if flavor is not None and 'spark' in flavor:
+ use_deprecated_int96_timestamps = True
+ else:
+ use_deprecated_int96_timestamps = False
+
self.flavor = flavor
if flavor is not None:
schema, self.schema_changed = _sanitize_schema(schema, flavor)
@@ -220,15 +250,29 @@ class ParquetWriter(object):
self.schema_changed = False
self.schema = schema
- self.writer = _parquet.ParquetWriter(where, schema, **options)
+ self.writer = _parquet.ParquetWriter(
+ where, schema,
+ version=version,
+ compression=compression,
+ use_dictionary=use_dictionary,
+ use_deprecated_int96_timestamps=use_deprecated_int96_timestamps,
+ **options)
+ self.is_open = True
+
+ def __del__(self):
+ if self.is_open:
+ self.close()
def write_table(self, table, row_group_size=None):
if self.schema_changed:
table = _sanitize_table(table, self.schema, self.flavor)
+ assert self.is_open
self.writer.write_table(table, row_group_size=row_group_size)
def close(self):
- self.writer.close()
+ if self.is_open:
+ self.writer.close()
+ self.is_open = False
def _get_pandas_index_columns(keyvalues):
@@ -857,52 +901,19 @@ def write_table(table, where, row_group_size=None, version='1.0',
use_deprecated_int96_timestamps=None,
coerce_timestamps=None,
flavor=None, **kwargs):
- """
- Write a Table to Parquet format
-
- Parameters
- ----------
- table : pyarrow.Table
- where: string or pyarrow.io.NativeFile
- row_group_size : int, default None
- The maximum number of rows in each Parquet RowGroup. As a default,
- we will write a single RowGroup per file.
- version : {"1.0", "2.0"}, default "1.0"
- The Parquet format version, defaults to 1.0
- use_dictionary : bool or list
- Specify if we should use dictionary encoding in general or only for
- some columns.
- use_deprecated_int96_timestamps : boolean, default None
- Write nanosecond resolution timestamps to INT96 Parquet
- format. Defaults to False unless enabled by flavor argument
- coerce_timestamps : string, default None
- Cast timestamps a particular resolution.
- Valid values: {None, 'ms', 'us'}
- compression : str or dict
- Specify the compression codec, either on a general basis or per-column.
- flavor : {'spark'}, default None
- Sanitize schema or set other compatibility options for compatibility
- """
- row_group_size = kwargs.get('chunk_size', row_group_size)
-
- if use_deprecated_int96_timestamps is None:
- # Use int96 timestamps for Spark
- if flavor is not None and 'spark' in flavor:
- use_deprecated_int96_timestamps = True
- else:
- use_deprecated_int96_timestamps = False
-
- options = dict(
- use_dictionary=use_dictionary,
- compression=compression,
- version=version,
- use_deprecated_int96_timestamps=use_deprecated_int96_timestamps,
- coerce_timestamps=coerce_timestamps)
+ row_group_size = kwargs.pop('chunk_size', row_group_size)
writer = None
try:
- writer = ParquetWriter(where, table.schema, flavor=flavor,
- **options)
+ writer = ParquetWriter(
+ where, table.schema,
+ version=version,
+ flavor=flavor,
+ use_dictionary=use_dictionary,
+ coerce_timestamps=coerce_timestamps,
+ compression=compression,
+ use_deprecated_int96_timestamps=use_deprecated_int96_timestamps,
+ **kwargs)
writer.write_table(table, row_group_size=row_group_size)
except:
if writer is not None:
@@ -917,6 +928,17 @@ def write_table(table, where, row_group_size=None, version='1.0',
writer.close()
+write_table.__doc__ = """
+Write a Table to Parquet format
+
+Parameters
+----------
+table : pyarrow.Table
+where: string or pyarrow.io.NativeFile
+{0}
+""".format(_parquet_writer_arg_docs)
+
+
def write_to_dataset(table, root_path, partition_cols=None,
filesystem=None, preserve_index=True, **kwargs):
"""
@@ -1013,12 +1035,10 @@ def write_metadata(schema, where, version='1.0',
Cast timestamps a particular resolution.
Valid values: {None, 'ms', 'us'}
"""
- options = dict(
- version=version,
+ writer = ParquetWriter(
+ where, schema, version=version,
use_deprecated_int96_timestamps=use_deprecated_int96_timestamps,
- coerce_timestamps=coerce_timestamps
- )
- writer = ParquetWriter(where, schema, **options)
+ coerce_timestamps=coerce_timestamps)
writer.close()
diff --git a/python/pyarrow/tests/test_parquet.py b/python/pyarrow/tests/test_parquet.py
index deb4b3f..09184cc 100644
--- a/python/pyarrow/tests/test_parquet.py
+++ b/python/pyarrow/tests/test_parquet.py
@@ -301,6 +301,35 @@ def test_pandas_parquet_native_file_roundtrip(tmpdir):
@parquet
+def test_parquet_incremental_file_build(tmpdir):
+ import pyarrow.parquet as pq
+
+ 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.0')
+
+ 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.get_result()
+ result = _read_table(pa.BufferReader(buf))
+
+ expected = pd.concat(frames, ignore_index=True)
+ tm.assert_frame_equal(result.to_pandas(), expected)
+
+
+@parquet
def test_read_pandas_column_subset(tmpdir):
import pyarrow.parquet as pq
diff --git a/python/pyarrow/tests/test_serialization.py b/python/pyarrow/tests/test_serialization.py
index 9321ebc..67798ac 100644
--- a/python/pyarrow/tests/test_serialization.py
+++ b/python/pyarrow/tests/test_serialization.py
@@ -266,30 +266,36 @@ def test_default_dict_serialization(large_memory_map):
def test_numpy_serialization(large_memory_map):
with pa.memory_map(large_memory_map, mode="r+") as mmap:
- for t in ["bool", "int8", "uint8", "int16", "uint16", "int32", "uint32",
- "float16", "float32", "float64"]:
+ for t in ["bool", "int8", "uint8", "int16", "uint16", "int32",
+ "uint32", "float16", "float32", "float64"]:
obj = np.random.randint(0, 10, size=(100, 100)).astype(t)
serialization_roundtrip(obj, mmap)
def test_datetime_serialization(large_memory_map):
- data = [# Principia Mathematica published
- datetime.datetime(year=1687, month=7, day=5),
- # Some random date
- datetime.datetime(year=1911, month=6, day=3, hour=4,
- minute=55, second=44),
- # End of WWI
- datetime.datetime(year=1918, month=11, day=11),
- # Beginning of UNIX time
- datetime.datetime(year=1970, month=1, day=1),
- # The Berlin wall falls
- datetime.datetime(year=1989, month=11, day=9),
- # Another random date
- datetime.datetime(year=2011, month=6, day=3, hour=4,
- minute=0, second=3),
- # Another random date
- datetime.datetime(year=1970, month=1, day=3, hour=4,
- minute=0, second=0)]
+ data = [
+ # Principia Mathematica published
+ datetime.datetime(year=1687, month=7, day=5),
+
+ # Some random date
+ datetime.datetime(year=1911, month=6, day=3, hour=4,
+ minute=55, second=44),
+ # End of WWI
+ datetime.datetime(year=1918, month=11, day=11),
+
+ # Beginning of UNIX time
+ datetime.datetime(year=1970, month=1, day=1),
+
+ # The Berlin wall falls
+ datetime.datetime(year=1989, month=11, day=9),
+
+ # Another random date
+ datetime.datetime(year=2011, month=6, day=3, hour=4,
+ minute=0, second=3),
+ # Another random date
+ datetime.datetime(year=1970, month=1, day=3, hour=4,
+ minute=0, second=0)
+ ]
with pa.memory_map(large_memory_map, mode="r+") as mmap:
for d in data:
serialization_roundtrip(d, mmap)
--
To stop receiving notification emails like this one, please contact
['"commits@arrow.apache.org" <co...@arrow.apache.org>'].