"""IPC write statistics

"""IPC read statistics

"""

__slots__ = ()

# cdef block is in lib.pxd

def __init__(self, *, bint ensure_native_endian=True,

bint use_threads=True, list included_fields=None):

self.c_options = CIpcReadOptions.Defaults()

self.ensure_native_endian = ensure_native_endian

self.use_threads = use_threads

if included_fields is not None:

self.included_fields = included_fields

@property

def ensure_native_endian(self):

return self.c_options.ensure_native_endian

@ensure_native_endian.setter

def ensure_native_endian(self, bint value):

self.c_options.ensure_native_endian = value

@property

def use_threads(self):

return self.c_options.use_threads

@use_threads.setter

def use_threads(self, bint value):

self.c_options.use_threads = value

@property

def included_fields(self):

return self.c_options.included_fields

@included_fields.setter

def included_fields(self, list value not None):

"""

__slots__ = ()

# cdef block is in lib.pxd

def __init__(self, *, metadata_version=MetadataVersion.V5,

bint allow_64bit=False, use_legacy_format=False,

compression=None, bint use_threads=True,

bint emit_dictionary_deltas=False,

bint unify_dictionaries=False):

self.c_options = CIpcWriteOptions.Defaults()

self.allow_64bit = allow_64bit

self.use_legacy_format = use_legacy_format

self.metadata_version = metadata_version

if compression is not None:

self.compression = compression

self.use_threads = use_threads

self.emit_dictionary_deltas = emit_dictionary_deltas

self.unify_dictionaries = unify_dictionaries

@property

def allow_64bit(self):

return self.c_options.allow_64bit

@allow_64bit.setter

def allow_64bit(self, bint value):

self.c_options.allow_64bit = value

@property

def use_legacy_format(self):

return self.c_options.write_legacy_ipc_format

@use_legacy_format.setter

def use_legacy_format(self, bint value):

self.c_options.write_legacy_ipc_format = value

@property

def metadata_version(self):

return _wrap_metadata_version(self.c_options.metadata_version)

@metadata_version.setter

def metadata_version(self, value):

self.c_options.metadata_version = _unwrap_metadata_version(value)

@property

def compression(self):

if self.c_options.codec == nullptr:

return None

else:

return frombytes(self.c_options.codec.get().name())

@compression.setter

def compression(self, value):

if value is None:

self.c_options.codec.reset()

elif isinstance(value, str):

codec_type = _ensure_compression(value)

if codec_type != CCompressionType_ZSTD and codec_type != CCompressionType_LZ4_FRAME:

raise ValueError("Compression type must be lz4, zstd or None")

self.c_options.codec = shared_ptr[CCodec](GetResultValue(

CCodec.Create(codec_type)).release())

elif isinstance(value, Codec):

if value.name != "lz4" and value.name != "zstd":

raise ValueError("Compression type must be lz4, zstd or None")

self.c_options.codec = (<Codec>value).wrapped

else:

raise TypeError(

"Property `compression` must be None, str, or pyarrow.Codec")

@property

def use_threads(self):

return self.c_options.use_threads

@use_threads.setter

def use_threads(self, bint value):

self.c_options.use_threads = value

@property

def emit_dictionary_deltas(self):

return self.c_options.emit_dictionary_deltas

@emit_dictionary_deltas.setter

def emit_dictionary_deltas(self, bint value):

self.c_options.emit_dictionary_deltas = value

@property

def unify_dictionaries(self):

return self.c_options.unify_dictionaries

@unify_dictionaries.setter

def unify_dictionaries(self, bint value):

"""

def __cinit__(self):

pass

def __init__(self):

raise TypeError("Do not call {}'s constructor directly, use "

"`pyarrow.ipc.read_message` function instead."

.format(self.__class__.__name__))

@property

def type(self):

return frombytes(FormatMessageType(self.message.get().type()))

@property

def metadata(self):

return pyarrow_wrap_buffer(self.message.get().metadata())

@property

def metadata_version(self):

return _wrap_metadata_version(self.message.get().metadata_version())

@property

def body(self):

cdef shared_ptr[CBuffer] body = self.message.get().body()

if body.get() == NULL:

return None

else:

return pyarrow_wrap_buffer(body)

def equals(self, Message other):

"""

cdef c_bool result

with nogil:

result = self.message.get().Equals(deref(other.message.get()))

return result

def serialize_to(self, NativeFile sink, alignment=8, memory_pool=None):

"""

cdef:

int64_t output_length = 0

COutputStream* out

CIpcWriteOptions options

options.alignment = alignment

out = sink.get_output_stream().get()

with nogil:

check_status(self.message.get()

.SerializeTo(out, options, &output_length))

def serialize(self, alignment=8, memory_pool=None):

"""

stream = BufferOutputStream(memory_pool)

self.serialize_to(stream, alignment=alignment, memory_pool=memory_pool)

return stream.getvalue()

def __repr__(self):

if self.message == nullptr:

return """pyarrow.Message(uninitialized)"""

metadata_len = self.metadata.size

body = self.body

body_len = 0 if body is None else body.size

return """pyarrow.Message

"""

cdef:

unique_ptr[CMessageReader] reader

def __cinit__(self):

pass

def __init__(self):

raise TypeError("Do not call {}'s constructor directly, use "

"`pyarrow.ipc.MessageReader.open_stream` function "

"instead.".format(self.__class__.__name__))

@staticmethod

def open_stream(source):

"""

cdef:

MessageReader result = MessageReader.__new__(MessageReader)

shared_ptr[CInputStream] in_stream

unique_ptr[CMessageReader] reader

_get_input_stream(source, &in_stream)

with nogil:

reader = CMessageReader.Open(in_stream)

result.reader.reset(reader.release())

return result

def __iter__(self):

while True:

yield self.read_next_message()

def read_next_message(self):

"""

cdef Message result = Message.__new__(Message)

with nogil:

result.message = move(GetResultValue(self.reader.get()

.ReadNextMessage()))

if result.message.get() == NULL:

raise StopIteration

"""The base RecordBatchWriter wrapper.

# cdef block is in lib.pxd

def write(self, table_or_batch):

"""

if isinstance(table_or_batch, RecordBatch):

self.write_batch(table_or_batch)

elif isinstance(table_or_batch, Table):

self.write_table(table_or_batch)

else:

raise ValueError(type(table_or_batch))

def write_batch(self, RecordBatch batch, custom_metadata=None):

"""

metadata = ensure_metadata(custom_metadata, allow_none=True)

c_meta = pyarrow_unwrap_metadata(metadata)

with nogil:

check_status(self.writer.get()

.WriteRecordBatch(deref(batch.batch), c_meta))

def write_table(self, Table table, max_chunksize=None):

"""

cdef:

# max_chunksize must be > 0 to have any impact

int64_t c_max_chunksize = -1

if max_chunksize is not None:

c_max_chunksize = max_chunksize

with nogil:

check_status(self.writer.get().WriteTable(table.table[0],

c_max_chunksize))

def close(self):

"""

with nogil:

check_status(self.writer.get().Close())

def __enter__(self):

return self

def __exit__(self, exc_type, exc_val, exc_tb):

self.close()

@property

def stats(self):

"""

if not self.writer:

raise ValueError("Operation on closed writer")

cdef:

CIpcWriteOptions options

bint closed

def __cinit__(self):

pass

def __dealloc__(self):

pass

@property

def _use_legacy_format(self):

# For testing (see test_ipc.py)

return self.options.write_legacy_ipc_format

@property

def _metadata_version(self):

# For testing (see test_ipc.py)

return _wrap_metadata_version(self.options.metadata_version)

def _open(self, sink, Schema schema not None,

IpcWriteOptions options=IpcWriteOptions()):

cdef:

shared_ptr[COutputStream] c_sink

self.options = options.c_options

get_writer(sink, &c_sink)

with nogil:

self.writer = GetResultValue(

MakeStreamWriter(c_sink, schema.sp_schema,

"""

table = self.read_all()

"""Base class for reading stream of record batches.

# cdef block is in lib.pxd

def __iter__(self):

while True:

try:

yield self.read_next_batch()

except StopIteration:

return

@property

def schema(self):

"""

cdef shared_ptr[CSchema] c_schema

with nogil:

c_schema = self.reader.get().schema()

return pyarrow_wrap_schema(c_schema)

def read_next_batch(self):

"""

cdef shared_ptr[CRecordBatch] batch

with nogil:

check_status(self.reader.get().ReadNext(&batch))

if batch.get() == NULL:

raise StopIteration

return pyarrow_wrap_batch(batch)

def read_next_batch_with_custom_metadata(self):

"""

cdef:

CRecordBatchWithMetadata batch_with_metadata

with nogil:

batch_with_metadata = GetResultValue(self.reader.get().ReadNext())

if batch_with_metadata.batch.get() == NULL:

raise StopIteration

return _wrap_record_batch_with_metadata(batch_with_metadata)

def iter_batches_with_custom_metadata(self):

"""

while True:

try:

yield self.read_next_batch_with_custom_metadata()

except StopIteration:

return

def read_all(self):

"""

cdef shared_ptr[CTable] table

with nogil:

check_status(self.reader.get().ToTable().Value(&table))

return pyarrow_wrap_table(table)

read_pandas = _ReadPandasMixin.read_pandas

def close(self):

"""

with nogil:

check_status(self.reader.get().Close())

def __enter__(self):

return self

def __exit__(self, exc_type, exc_val, exc_tb):

self.close()

def _export_to_c(self, out_ptr):

"""

cdef:

void* c_ptr = _as_c_pointer(out_ptr)

with nogil:

check_status(ExportRecordBatchReader(

self.reader, <ArrowArrayStream*> c_ptr))

@staticmethod

def _import_from_c(in_ptr):

"""

cdef:

void* c_ptr = _as_c_pointer(in_ptr)

shared_ptr[CRecordBatchReader] c_reader

RecordBatchReader self

with nogil:

c_reader = GetResultValue(ImportRecordBatchReader(

<ArrowArrayStream*> c_ptr))

self = RecordBatchReader.__new__(RecordBatchReader)

self.reader = c_reader

return self

@staticmethod

def from_batches(Schema schema not None, batches):

"""

cdef:

shared_ptr[CSchema] c_schema

shared_ptr[CRecordBatchReader] c_reader

RecordBatchReader self

c_schema = pyarrow_unwrap_schema(schema)

c_reader = GetResultValue(CPyRecordBatchReader.Make(

c_schema, batches))

self = RecordBatchReader.__new__(RecordBatchReader)

self.reader = c_reader

cdef:

shared_ptr[CInputStream] in_stream

CIpcReadOptions options

CRecordBatchStreamReader* stream_reader

def __cinit__(self):

pass

def _open(self, source, IpcReadOptions options=IpcReadOptions(),

MemoryPool memory_pool=None):

self.options = options.c_options

self.options.memory_pool = maybe_unbox_memory_pool(memory_pool)

_get_input_stream(source, &self.in_stream)

with nogil:

self.reader = GetResultValue(CRecordBatchStreamReader.Open(

self.in_stream, self.options))

self.stream_reader = <CRecordBatchStreamReader*> self.reader.get()

@property

def stats(self):

"""

if not self.reader:

raise ValueError("Operation on closed reader")

def _open(self, sink, Schema schema not None,

IpcWriteOptions options=IpcWriteOptions()):

cdef:

shared_ptr[COutputStream] c_sink

self.options = options.c_options

get_writer(sink, &c_sink)

with nogil:

self.writer = GetResultValue(

"""RecordBatch with its custom metadata

cdef:

shared_ptr[CRecordBatchFileReader] reader

shared_ptr[CRandomAccessFile] file

CIpcReadOptions options

cdef readonly:

Schema schema

def __cinit__(self):

pass

def _open(self, source, footer_offset=None,

IpcReadOptions options=IpcReadOptions(),

MemoryPool memory_pool=None):

self.options = options.c_options

self.options.memory_pool = maybe_unbox_memory_pool(memory_pool)

try:

source = as_buffer(source)

except TypeError:

pass

get_reader(source, False, &self.file)

cdef int64_t offset = 0

if footer_offset is not None:

offset = footer_offset

with nogil:

if offset != 0:

self.reader = GetResultValue(

CRecordBatchFileReader.Open2(self.file.get(), offset,

self.options))

else:

self.reader = GetResultValue(

CRecordBatchFileReader.Open(self.file.get(),

self.options))

self.schema = pyarrow_wrap_schema(self.reader.get().schema())

@property

def num_record_batches(self):

"""

return self.reader.get().num_record_batches()

def get_batch(self, int i):

"""

cdef shared_ptr[CRecordBatch] batch

if i < 0 or i >= self.num_record_batches:

raise ValueError('Batch number {0} out of range'.format(i))

with nogil:

batch = GetResultValue(self.reader.get().ReadRecordBatch(i))

return pyarrow_wrap_batch(batch)

# TODO(wesm): ARROW-503: Function was renamed. Remove after a period of

# time has passed

get_record_batch = get_batch

def get_batch_with_custom_metadata(self, int i):

"""

cdef:

CRecordBatchWithMetadata batch_with_metadata

if i < 0 or i >= self.num_record_batches: