2017-01-23 09:10:18 -05:00
|
|
|
# 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
|
|
|
|
|
|
2020-07-11 18:12:35 -05:00
|
|
|
import os
|
2019-05-27 18:01:27 +02:00
|
|
|
|
2017-05-16 13:14:18 -04:00
|
|
|
import pyarrow as pa
|
2017-07-15 16:51:51 -04:00
|
|
|
|
2022-04-19 15:31:41 +02:00
|
|
|
from pyarrow.lib import (IpcReadOptions, IpcWriteOptions, ReadStats, WriteStats, # noqa
|
2020-12-14 18:12:34 -06:00
|
|
|
Message, MessageReader,
|
2020-10-01 12:00:12 +02:00
|
|
|
RecordBatchReader, _ReadPandasMixin,
|
2025-05-20 01:32:04 +02:00
|
|
|
MetadataVersion, Alignment,
|
2017-08-24 23:04:58 -04:00
|
|
|
read_message, read_record_batch, read_schema,
|
2017-07-15 16:51:51 -04:00
|
|
|
read_tensor, write_tensor,
|
|
|
|
|
get_record_batch_size, get_tensor_size)
|
2017-05-13 15:44:43 -04:00
|
|
|
import pyarrow.lib as lib
|
2017-01-23 09:10:18 -05:00
|
|
|
|
|
|
|
|
|
2020-10-01 12:00:12 +02:00
|
|
|
class RecordBatchStreamReader(lib._RecordBatchStreamReader):
|
2017-01-23 09:10:18 -05:00
|
|
|
"""
|
2020-03-25 14:44:38 +01:00
|
|
|
Reader for the Arrow streaming binary format.
|
2017-01-23 09:10:18 -05:00
|
|
|
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
2018-07-24 15:23:06 -04:00
|
|
|
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
|
2020-03-25 14:44:38 +01:00
|
|
|
Either an in-memory buffer, or a readable file object.
|
2022-06-20 10:39:42 -05:00
|
|
|
If you want to use memory map use MemoryMappedFile as source.
|
2022-04-19 15:31:41 +02:00
|
|
|
options : pyarrow.ipc.IpcReadOptions
|
|
|
|
|
Options for IPC deserialization.
|
|
|
|
|
If None, default values will be used.
|
|
|
|
|
memory_pool : MemoryPool, default None
|
|
|
|
|
If None, default memory pool is used.
|
2017-01-23 09:10:18 -05:00
|
|
|
"""
|
2019-10-18 13:58:26 -04:00
|
|
|
|
2022-04-19 15:31:41 +02:00
|
|
|
def __init__(self, source, *, options=None, memory_pool=None):
|
|
|
|
|
options = _ensure_default_ipc_read_options(options)
|
|
|
|
|
self._open(source, options=options, memory_pool=memory_pool)
|
2017-01-23 09:10:18 -05:00
|
|
|
|
|
|
|
|
|
2019-09-18 13:47:38 -05:00
|
|
|
_ipc_writer_class_doc = """\
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
|
|
|
|
sink : str, pyarrow.NativeFile, or file-like Python object
|
2020-03-25 14:44:38 +01:00
|
|
|
Either a file path, or a writable file object.
|
2019-09-18 13:47:38 -05:00
|
|
|
schema : pyarrow.Schema
|
2020-03-25 14:44:38 +01:00
|
|
|
The Arrow schema for data to be written to the file.
|
2020-07-11 18:12:35 -05:00
|
|
|
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
|
2022-10-19 23:41:24 -08:00
|
|
|
ARROW_PRE_1_0_METADATA_VERSION=1."""
|
2019-09-18 13:47:38 -05:00
|
|
|
|
|
|
|
|
|
2025-05-12 11:16:31 +02:00
|
|
|
_ipc_file_writer_class_doc = (
|
|
|
|
|
_ipc_writer_class_doc
|
|
|
|
|
+ "\n"
|
|
|
|
|
+ """\
|
|
|
|
|
metadata : dict | pyarrow.KeyValueMetadata, optional
|
|
|
|
|
Key/value pairs (both must be bytes-like) that will be stored
|
|
|
|
|
in the file footer and are retrievable via
|
|
|
|
|
pyarrow.ipc.open_file(...).metadata."""
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
2019-02-14 18:16:56 +01:00
|
|
|
class RecordBatchStreamWriter(lib._RecordBatchStreamWriter):
|
2025-05-12 20:34:49 +08:00
|
|
|
__doc__ = f"""Writer for the Arrow streaming binary format
|
2017-01-23 09:10:18 -05:00
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
{_ipc_writer_class_doc}"""
|
2019-09-18 13:47:38 -05:00
|
|
|
|
2025-04-16 10:49:50 +02:00
|
|
|
def __init__(self, sink, schema, *, options=None):
|
|
|
|
|
options = _get_legacy_format_default(options)
|
2020-07-11 18:12:35 -05:00
|
|
|
self._open(sink, schema, options=options)
|
2017-01-23 09:10:18 -05:00
|
|
|
|
|
|
|
|
|
2020-10-01 12:00:12 +02:00
|
|
|
class RecordBatchFileReader(lib._RecordBatchFileReader):
|
2017-01-23 09:10:18 -05:00
|
|
|
"""
|
|
|
|
|
Class for reading Arrow record batch data from the Arrow binary file format
|
|
|
|
|
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
2018-07-24 15:23:06 -04:00
|
|
|
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
|
2022-06-20 10:39:42 -05:00
|
|
|
Either an in-memory buffer, or a readable file object.
|
|
|
|
|
If you want to use memory map use MemoryMappedFile as source.
|
2017-01-23 09:10:18 -05:00
|
|
|
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
|
2022-04-19 15:31:41 +02:00
|
|
|
options : pyarrow.ipc.IpcReadOptions
|
|
|
|
|
Options for IPC serialization.
|
|
|
|
|
If None, default values will be used.
|
|
|
|
|
memory_pool : MemoryPool, default None
|
|
|
|
|
If None, default memory pool is used.
|
2017-01-23 09:10:18 -05:00
|
|
|
"""
|
2020-05-19 09:33:39 +02:00
|
|
|
|
2022-04-19 15:31:41 +02:00
|
|
|
def __init__(self, source, footer_offset=None, *, options=None,
|
|
|
|
|
memory_pool=None):
|
|
|
|
|
options = _ensure_default_ipc_read_options(options)
|
|
|
|
|
self._open(source, footer_offset=footer_offset,
|
|
|
|
|
options=options, memory_pool=memory_pool)
|
2017-01-23 09:10:18 -05:00
|
|
|
|
|
|
|
|
|
2017-05-14 08:55:26 -04:00
|
|
|
class RecordBatchFileWriter(lib._RecordBatchFileWriter):
|
2017-01-23 09:10:18 -05:00
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
__doc__ = f"""Writer to create the Arrow binary file format
|
2019-09-18 13:47:38 -05:00
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
{_ipc_file_writer_class_doc}"""
|
2019-09-18 13:47:38 -05:00
|
|
|
|
2025-05-12 11:16:31 +02:00
|
|
|
def __init__(self, sink, schema, *, options=None, metadata=None):
|
2025-04-16 10:49:50 +02:00
|
|
|
options = _get_legacy_format_default(options)
|
2025-05-12 11:16:31 +02:00
|
|
|
self._open(sink, schema, options=options, metadata=metadata)
|
2020-07-11 18:12:35 -05:00
|
|
|
|
2019-09-18 13:47:38 -05:00
|
|
|
|
2025-04-16 10:49:50 +02:00
|
|
|
def _get_legacy_format_default(options):
|
|
|
|
|
if options:
|
2020-07-13 12:49:07 -05:00
|
|
|
if not isinstance(options, IpcWriteOptions):
|
2025-05-12 20:34:49 +08:00
|
|
|
raise TypeError(f"expected IpcWriteOptions, got {type(options)}")
|
2020-07-11 18:12:35 -05:00
|
|
|
return options
|
2019-09-18 13:47:38 -05:00
|
|
|
|
2020-07-11 18:12:35 -05:00
|
|
|
metadata_version = MetadataVersion.V5
|
2025-04-16 10:49:50 +02:00
|
|
|
use_legacy_format = \
|
|
|
|
|
bool(int(os.environ.get('ARROW_PRE_0_15_IPC_FORMAT', '0')))
|
2020-07-11 18:12:35 -05:00
|
|
|
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)
|
2017-05-14 08:55:26 -04:00
|
|
|
|
|
|
|
|
|
2022-04-19 15:31:41 +02:00
|
|
|
def _ensure_default_ipc_read_options(options):
|
|
|
|
|
if options and not isinstance(options, IpcReadOptions):
|
2025-05-12 20:34:49 +08:00
|
|
|
raise TypeError(f"expected IpcReadOptions, got {type(options)}")
|
2022-04-19 15:31:41 +02:00
|
|
|
return options or IpcReadOptions()
|
|
|
|
|
|
|
|
|
|
|
2025-04-16 10:49:50 +02:00
|
|
|
def new_stream(sink, schema, *, options=None):
|
2020-04-06 20:58:35 -05:00
|
|
|
return RecordBatchStreamWriter(sink, schema,
|
2020-07-11 18:12:35 -05:00
|
|
|
options=options)
|
2020-04-06 20:58:35 -05:00
|
|
|
|
|
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
new_stream.__doc__ = f"""\
|
2020-04-06 20:58:35 -05:00
|
|
|
Create an Arrow columnar IPC stream writer instance
|
|
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
{_ipc_writer_class_doc}
|
2022-06-12 07:11:27 -07:00
|
|
|
|
|
|
|
|
Returns
|
|
|
|
|
-------
|
|
|
|
|
writer : RecordBatchStreamWriter
|
2023-01-06 14:21:27 -09:00
|
|
|
A writer for the given sink
|
2025-05-12 20:34:49 +08:00
|
|
|
"""
|
2020-04-06 20:58:35 -05:00
|
|
|
|
|
|
|
|
|
2022-04-19 15:31:41 +02:00
|
|
|
def open_stream(source, *, options=None, memory_pool=None):
|
2017-05-14 08:55:26 -04:00
|
|
|
"""
|
2020-03-25 14:44:38 +01:00
|
|
|
Create reader for Arrow streaming format.
|
2017-05-14 08:55:26 -04:00
|
|
|
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
2018-07-24 15:23:06 -04:00
|
|
|
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
|
2020-03-25 14:44:38 +01:00
|
|
|
Either an in-memory buffer, or a readable file object.
|
2022-04-19 15:31:41 +02:00
|
|
|
options : pyarrow.ipc.IpcReadOptions
|
|
|
|
|
Options for IPC serialization.
|
|
|
|
|
If None, default values will be used.
|
|
|
|
|
memory_pool : MemoryPool, default None
|
|
|
|
|
If None, default memory pool is used.
|
2023-01-06 14:21:27 -09:00
|
|
|
|
2017-05-14 08:55:26 -04:00
|
|
|
Returns
|
|
|
|
|
-------
|
|
|
|
|
reader : RecordBatchStreamReader
|
2023-01-06 14:21:27 -09:00
|
|
|
A reader for the given source
|
2017-05-14 08:55:26 -04:00
|
|
|
"""
|
2022-04-19 15:31:41 +02:00
|
|
|
return RecordBatchStreamReader(source, options=options,
|
|
|
|
|
memory_pool=memory_pool)
|
2017-05-14 08:55:26 -04:00
|
|
|
|
|
|
|
|
|
2025-05-12 11:16:31 +02:00
|
|
|
def new_file(sink, schema, *, options=None, metadata=None):
|
|
|
|
|
return RecordBatchFileWriter(sink, schema, options=options, metadata=metadata)
|
2020-04-06 20:58:35 -05:00
|
|
|
|
|
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
new_file.__doc__ = f"""\
|
2020-04-06 20:58:35 -05:00
|
|
|
Create an Arrow columnar IPC file writer instance
|
|
|
|
|
|
2025-05-12 20:34:49 +08:00
|
|
|
{_ipc_file_writer_class_doc}
|
2022-06-12 07:11:27 -07:00
|
|
|
|
|
|
|
|
Returns
|
|
|
|
|
-------
|
|
|
|
|
writer : RecordBatchFileWriter
|
2023-01-06 14:21:27 -09:00
|
|
|
A writer for the given sink
|
2025-05-12 20:34:49 +08:00
|
|
|
"""
|
2020-04-06 20:58:35 -05:00
|
|
|
|
|
|
|
|
|
2022-04-19 15:31:41 +02:00
|
|
|
def open_file(source, footer_offset=None, *, options=None, memory_pool=None):
|
2017-05-14 08:55:26 -04:00
|
|
|
"""
|
2020-03-25 14:44:38 +01:00
|
|
|
Create reader for Arrow file format.
|
2017-05-14 08:55:26 -04:00
|
|
|
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
2018-07-24 15:23:06 -04:00
|
|
|
source : bytes/buffer-like, pyarrow.NativeFile, or file-like Python object
|
2020-03-25 14:44:38 +01:00
|
|
|
Either an in-memory buffer, or a readable file object.
|
2017-05-14 08:55:26 -04:00
|
|
|
footer_offset : int, default None
|
|
|
|
|
If the file is embedded in some larger file, this is the byte offset to
|
2020-03-25 14:44:38 +01:00
|
|
|
the very end of the file data.
|
2022-04-19 15:31:41 +02:00
|
|
|
options : pyarrow.ipc.IpcReadOptions
|
|
|
|
|
Options for IPC serialization.
|
|
|
|
|
If None, default values will be used.
|
|
|
|
|
memory_pool : MemoryPool, default None
|
|
|
|
|
If None, default memory pool is used.
|
2023-01-06 14:21:27 -09:00
|
|
|
|
2017-05-14 08:55:26 -04:00
|
|
|
Returns
|
|
|
|
|
-------
|
|
|
|
|
reader : RecordBatchFileReader
|
2023-01-06 14:21:27 -09:00
|
|
|
A reader for the given source
|
2017-05-14 08:55:26 -04:00
|
|
|
"""
|
2022-04-19 15:31:41 +02:00
|
|
|
return RecordBatchFileReader(
|
|
|
|
|
source, footer_offset=footer_offset,
|
|
|
|
|
options=options, memory_pool=memory_pool)
|
2017-05-16 13:14:18 -04:00
|
|
|
|
|
|
|
|
|
2020-07-13 12:49:07 -05:00
|
|
|
def serialize_pandas(df, *, nthreads=None, preserve_index=None):
|
2020-03-25 14:44:38 +01:00
|
|
|
"""
|
|
|
|
|
Serialize a pandas DataFrame into a buffer protocol compatible object.
|
2017-05-16 13:14:18 -04:00
|
|
|
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
|
|
|
|
df : pandas.DataFrame
|
2017-10-08 11:45:06 -04:00
|
|
|
nthreads : int, default None
|
2020-03-25 14:44:38 +01:00
|
|
|
Number of threads to use for conversion to Arrow, default all CPUs.
|
|
|
|
|
preserve_index : bool, default None
|
2019-06-25 20:57:02 -07:00
|
|
|
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.
|
2017-05-16 13:14:18 -04:00
|
|
|
|
|
|
|
|
Returns
|
|
|
|
|
-------
|
|
|
|
|
buf : buffer
|
2020-03-25 14:44:38 +01:00
|
|
|
An object compatible with the buffer protocol.
|
2017-05-16 13:14:18 -04:00
|
|
|
"""
|
2017-10-09 21:02:48 -04:00
|
|
|
batch = pa.RecordBatch.from_pandas(df, nthreads=nthreads,
|
|
|
|
|
preserve_index=preserve_index)
|
2017-08-21 22:30:20 -04:00
|
|
|
sink = pa.BufferOutputStream()
|
2019-10-18 13:58:26 -04:00
|
|
|
with pa.RecordBatchStreamWriter(sink, batch.schema) as writer:
|
|
|
|
|
writer.write_batch(batch)
|
2018-09-04 17:24:23 -04:00
|
|
|
return sink.getvalue()
|
2017-05-16 13:14:18 -04:00
|
|
|
|
|
|
|
|
|
2020-07-13 12:49:07 -05:00
|
|
|
def deserialize_pandas(buf, *, use_threads=True):
|
2017-05-16 13:14:18 -04:00
|
|
|
"""Deserialize a buffer protocol compatible object into a pandas DataFrame.
|
|
|
|
|
|
|
|
|
|
Parameters
|
|
|
|
|
----------
|
|
|
|
|
buf : buffer
|
2020-03-25 14:44:38 +01:00
|
|
|
An object compatible with the buffer protocol.
|
2021-10-04 11:44:40 +02:00
|
|
|
use_threads : bool, default True
|
2020-03-25 14:44:38 +01:00
|
|
|
Whether to parallelize the conversion using multiple threads.
|
2017-05-16 13:14:18 -04:00
|
|
|
|
|
|
|
|
Returns
|
|
|
|
|
-------
|
|
|
|
|
df : pandas.DataFrame
|
2023-01-06 14:21:27 -09:00
|
|
|
The buffer deserialized as pandas DataFrame
|
2017-05-16 13:14:18 -04:00
|
|
|
"""
|
|
|
|
|
buffer_reader = pa.BufferReader(buf)
|
2019-10-18 13:58:26 -04:00
|
|
|
with pa.RecordBatchStreamReader(buffer_reader) as reader:
|
|
|
|
|
table = reader.read_all()
|
2018-05-24 21:28:15 -04:00
|
|
|
return table.to_pandas(use_threads=use_threads)
|
