# 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.
# Arrow file and stream reader/writer classes, and other messaging tools
import os
import pyarrow as pa
from pyarrow.lib import (IpcWriteOptions, ReadStats, WriteStats, # noqa
Message, MessageReader,
RecordBatchReader, _ReadPandasMixin,
MetadataVersion,
read_message, read_record_batch, read_schema,
read_tensor, write_tensor,
get_record_batch_size, get_tensor_size)
import pyarrow.lib as lib
class RecordBatchStreamReader(lib._RecordBatchStreamReader):
"""
Reader for the Arrow streaming binary format.
Parameters
----------
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
Either an in-memory buffer, or a readable file object.
"""
def __init__(self, source):
self._open(source)
_ipc_writer_class_doc = """\
Parameters
----------
sink : str, pyarrow.NativeFile, or file-like Python object
Either a file path, or a writable file object.
schema : pyarrow.Schema
The Arrow schema for data to be written to the file.
options : pyarrow.ipc.IpcWriteOptions
Options for IPC serialization.
If None, default values will be used: the legacy format will not
be used unless overridden by setting the environment variable
ARROW_PRE_0_15_IPC_FORMAT=1, and the V5 metadata version will be
used unless overridden by setting the environment variable
ARROW_PRE_1_0_METADATA_VERSION=1.
use_legacy_format : bool, default None
Deprecated in favor of setting options. Cannot be provided with
options.
If None, False will be used unless this default is overridden by
setting the environment variable ARROW_PRE_0_15_IPC_FORMAT=1"""
class RecordBatchStreamWriter(lib._RecordBatchStreamWriter):
__doc__ = """Writer for the Arrow streaming binary format
{}""".format(_ipc_writer_class_doc)
def __init__(self, sink, schema, *, use_legacy_format=None, options=None):
options = _get_legacy_format_default(use_legacy_format, options)
self._open(sink, schema, options=options)
class RecordBatchFileReader(lib._RecordBatchFileReader):
"""
Class for reading Arrow record batch data from the Arrow binary file format
Parameters
----------
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
Either an in-memory buffer, or a readable file object
footer_offset : int, default None
If the file is embedded in some larger file, this is the byte offset to
the very end of the file data
"""
def __init__(self, source, footer_offset=None):
self._open(source, footer_offset=footer_offset)
class RecordBatchFileWriter(lib._RecordBatchFileWriter):
__doc__ = """Writer to create the Arrow binary file format
{}""".format(_ipc_writer_class_doc)
def __init__(self, sink, schema, *, use_legacy_format=None, options=None):
options = _get_legacy_format_default(use_legacy_format, options)
self._open(sink, schema, options=options)
def _get_legacy_format_default(use_legacy_format, options):
if use_legacy_format is not None and options is not None:
raise ValueError(
"Can provide at most one of options and use_legacy_format")
elif options:
if not isinstance(options, IpcWriteOptions):
raise TypeError("expected IpcWriteOptions, got {}"
.format(type(options)))
return options
metadata_version = MetadataVersion.V5
if use_legacy_format is None:
use_legacy_format = \
bool(int(os.environ.get('ARROW_PRE_0_15_IPC_FORMAT', '0')))
if bool(int(os.environ.get('ARROW_PRE_1_0_METADATA_VERSION', '0'))):
metadata_version = MetadataVersion.V4
return IpcWriteOptions(use_legacy_format=use_legacy_format,
metadata_version=metadata_version)
def new_stream(sink, schema, *, use_legacy_format=None, options=None):
return RecordBatchStreamWriter(sink, schema,
use_legacy_format=use_legacy_format,
options=options)
new_stream.__doc__ = """\
Create an Arrow columnar IPC stream writer instance
{}""".format(_ipc_writer_class_doc)
[docs]def open_stream(source):
"""
Create reader for Arrow streaming format.
Parameters
----------
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
Either an in-memory buffer, or a readable file object.
Returns
-------
reader : RecordBatchStreamReader
"""
return RecordBatchStreamReader(source)
def new_file(sink, schema, *, use_legacy_format=None, options=None):
return RecordBatchFileWriter(sink, schema,
use_legacy_format=use_legacy_format,
options=options)
new_file.__doc__ = """\
Create an Arrow columnar IPC file writer instance
{}""".format(_ipc_writer_class_doc)
[docs]def open_file(source, footer_offset=None):
"""
Create reader for Arrow file format.
Parameters
----------
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
Either an in-memory buffer, or a readable file object.
footer_offset : int, default None
If the file is embedded in some larger file, this is the byte offset to
the very end of the file data.
Returns
-------
reader : RecordBatchFileReader
"""
return RecordBatchFileReader(source, footer_offset=footer_offset)
def serialize_pandas(df, *, nthreads=None, preserve_index=None):
"""
Serialize a pandas DataFrame into a buffer protocol compatible object.
Parameters
----------
df : pandas.DataFrame
nthreads : int, default None
Number of threads to use for conversion to Arrow, default all CPUs.
preserve_index : bool, default None
The default of None will store the index as a column, except for
RangeIndex which is stored as metadata only. If True, always
preserve the pandas index data as a column. If False, no index
information is saved and the result will have a default RangeIndex.
Returns
-------
buf : buffer
An object compatible with the buffer protocol.
"""
batch = pa.RecordBatch.from_pandas(df, nthreads=nthreads,
preserve_index=preserve_index)
sink = pa.BufferOutputStream()
with pa.RecordBatchStreamWriter(sink, batch.schema) as writer:
writer.write_batch(batch)
return sink.getvalue()
def deserialize_pandas(buf, *, use_threads=True):
"""Deserialize a buffer protocol compatible object into a pandas DataFrame.
Parameters
----------
buf : buffer
An object compatible with the buffer protocol.
use_threads : bool, default True
Whether to parallelize the conversion using multiple threads.
Returns
-------
df : pandas.DataFrame
"""
buffer_reader = pa.BufferReader(buf)
with pa.RecordBatchStreamReader(buffer_reader) as reader:
table = reader.read_all()
return table.to_pandas(use_threads=use_threads)
