Module enrgdaq.utils.arrow_ipc

Arrow IPC utilities for zero-copy serialization/deserialization.

This module provides functions to write and read PyArrow tables directly to/from shared memory buffers without intermediate copies.

Functions

def get_table_ipc_size(table: pyarrow.lib.Table) ‑> int

Expand source code

def get_table_ipc_size(table: pa.Table) -> int:
    """
    Calculate the size needed to serialize a PyArrow table in IPC format.

    This includes schema overhead, so the actual size may vary slightly.
    We add a buffer for safety.

    Args:
        table: The PyArrow table to measure.

    Returns:
        Estimated size in bytes.
    """
    # Use a mock stream to calculate size
    sink = pa.BufferOutputStream()
    with pa.ipc.new_stream(sink, table.schema) as writer:
        writer.write_table(table)
    return sink.getvalue().size

Calculate the size needed to serialize a PyArrow table in IPC format.

This includes schema overhead, so the actual size may vary slightly. We add a buffer for safety.

Args

-----=

table

The PyArrow table to measure.

Returns -----= Estimated size in bytes.

def ipc_bytes_to_table(data: bytes) ‑> pyarrow.lib.Table

Expand source code

def ipc_bytes_to_table(data: bytes) -> pa.Table:
    """
    Deserialize a PyArrow table from IPC bytes.

    This is a fallback for when zero-copy is not available.

    Args:
        data: Bytes containing the IPC-serialized table.

    Returns:
        The deserialized PyArrow Table.
    """
    reader = pa.ipc.open_stream(data)
    return reader.read_all()

Deserialize a PyArrow table from IPC bytes.

This is a fallback for when zero-copy is not available.

Args

-----=

data

Bytes containing the IPC-serialized table.

Returns -----= The deserialized PyArrow Table.

def read_table_from_address(address: int, size: int, base: object) ‑> pyarrow.lib.Table

Expand source code

def read_table_from_address(address: int, size: int, base: object) -> pa.Table:
    """
    Read a PyArrow table from a memory address using zero-copy.

    This is useful when you have the raw memory address (e.g., from
    get_slot_address) rather than a memoryview.

    Args:
        address: Memory address of the IPC data.
        size: Number of bytes of valid data.
        base: Python object to keep alive (prevents garbage collection
              of the underlying memory).

    Returns:
        A PyArrow Table that references the memory directly.
    """
    arrow_buf = pa.foreign_buffer(address, size, base=base)
    reader = pa.ipc.open_stream(arrow_buf)
    return reader.read_all()

Read a PyArrow table from a memory address using zero-copy.

This is useful when you have the raw memory address (e.g., from get_slot_address) rather than a memoryview.

Args

-----=

address

Memory address of the IPC data.

size

Number of bytes of valid data.

base

Python object to keep alive (prevents garbage collection of the underlying memory).

Returns -----= A PyArrow Table that references the memory directly.

def read_table_from_buffer(buffer: memoryview, size: int) ‑> pyarrow.lib.Table

Expand source code

def read_table_from_buffer(buffer: memoryview, size: int) -> pa.Table:
    """
    Read a PyArrow table from a buffer using zero-copy.

    This uses pa.foreign_buffer to wrap the shared memory, allowing
    PyArrow to read directly from the memory without copying.

    Args:
        buffer: A memoryview containing the IPC data.
        size: Number of bytes of valid data in the buffer.

    Returns:
        A PyArrow Table that references the shared memory directly.

    Important:
        The returned table references the shared memory buffer. The buffer
        must remain valid for as long as the table is in use.
    """
    # Get the memory address of the buffer
    # For a memoryview backed by shared memory, this gives us the shared address
    c_array = (ctypes.c_char * len(buffer)).from_buffer(buffer)
    address = ctypes.addressof(c_array)

    # Create an Arrow buffer that references the shared memory
    # The `buffer` parameter keeps the Python memoryview alive
    arrow_buf = pa.foreign_buffer(address, size, base=buffer)

    # Create a buffer reader and read the IPC stream
    reader = pa.ipc.open_stream(arrow_buf)
    table = reader.read_all()

    return table

Read a PyArrow table from a buffer using zero-copy.

This uses pa.foreign_buffer to wrap the shared memory, allowing PyArrow to read directly from the memory without copying.

Args

-----=

buffer

A memoryview containing the IPC data.

size

Number of bytes of valid data in the buffer.

Returns -----= A PyArrow Table that references the shared memory directly.

Important -----= The returned table references the shared memory buffer. The buffer must remain valid for as long as the table is in use.

def table_to_ipc_bytes(table: pyarrow.lib.Table) ‑> bytes

Expand source code

def table_to_ipc_bytes(table: pa.Table) -> bytes:
    """
    Serialize a PyArrow table to bytes using IPC format.

    This is a fallback for when zero-copy is not available.
    Creates a copy of the data.

    Args:
        table: The PyArrow table to serialize.

    Returns:
        Bytes containing the IPC-serialized table.
    """
    sink = pa.BufferOutputStream()
    with pa.ipc.new_stream(sink, table.schema) as writer:
        writer.write_table(table)
    return sink.getvalue().to_pybytes()

Serialize a PyArrow table to bytes using IPC format.

This is a fallback for when zero-copy is not available. Creates a copy of the data.

Args

-----=

table

The PyArrow table to serialize.

Returns -----= Bytes containing the IPC-serialized table.

def try_zero_copy_pyarrow(table: pyarrow.lib.Table, store_config: Any, tag: Any) ‑> tuple['RingBufferHandle | None', bool]

Expand source code

def try_zero_copy_pyarrow(
    table: pa.Table,
    store_config: Any,
    tag: Any,
) -> tuple["RingBufferHandle | None", bool]:
    """
    Try to write a PyArrow table to the ring buffer for zero-copy transfer.

    Args:
        table: The PyArrow table to transfer.
        store_config: The store config to copy to the new message.
        tag: The tag to copy to the new message.

    Returns:
        A tuple of (RingBufferHandle or None, success: bool).
        If successful, returns the handle; otherwise returns None.
    """
    try:
        from enrgdaq.daq.models import RingBufferHandle
        from enrgdaq.utils.shared_ring_buffer import get_global_ring_buffer

        ring_buffer = get_global_ring_buffer()

        # Calculate required size
        table_size = get_table_ipc_size(table)

        # Try to allocate a slot
        slot = ring_buffer.allocate(table_size)
        if slot is None:
            logger.debug(
                f"Ring buffer allocation failed (size={table_size}), "
                "falling back to pickle-based SHM"
            )
            return None, False

        try:
            # Write table directly to slot
            bytes_written = write_table_to_buffer(table, slot.view)

            # Mark slot as ready for reading and track bytes
            ring_buffer.mark_ready(slot.index, bytes_written)

            # Create the handle
            handle = RingBufferHandle(
                buffer_name=ring_buffer.name,
                slot_index=slot.index,
                data_size=bytes_written,
                is_pyarrow=True,
                total_size=ring_buffer.total_size,
                slot_size=ring_buffer.slot_size,
            )

            return handle, True
        except Exception as e:
            # If writing fails, release the slot and fall back
            logger.warning(f"Zero-copy write failed: {e}, falling back")
            slot.release()
            return None, False

    except Exception as e:
        logger.debug(f"Zero-copy setup failed: {e}, falling back to pickle")
        return None, False

Try to write a PyArrow table to the ring buffer for zero-copy transfer.

Args

-----=

table

The PyArrow table to transfer.

store_config

The store config to copy to the new message.

tag

The tag to copy to the new message.

Returns -----= A tuple of (RingBufferHandle or None, success: bool). If successful, returns the handle; otherwise returns None.

def write_table_to_buffer(table: pyarrow.lib.Table, buffer: memoryview) ‑> int

Expand source code

def write_table_to_buffer(table: pa.Table, buffer: memoryview) -> int:
    """
    Write a PyArrow table to a buffer using Arrow IPC streaming format.

    The table is serialized directly into the provided buffer memory.

    Args:
        table: The PyArrow table to write.
        buffer: A memoryview providing the destination memory.

    Returns:
        Number of bytes written.

    Raises:
        ValueError: If the buffer is too small.
    """
    # Create a fixed-size buffer writer that writes directly to the memoryview
    # We need to wrap the memoryview in an Arrow buffer first
    arrow_buf = pa.py_buffer(buffer)
    output_stream = pa.FixedSizeBufferWriter(arrow_buf)

    with pa.ipc.new_stream(output_stream, table.schema) as writer:
        writer.write_table(table)

    bytes_written = output_stream.tell()
    output_stream.close()

    return bytes_written

Write a PyArrow table to a buffer using Arrow IPC streaming format.

The table is serialized directly into the provided buffer memory.

Args

-----=

table

The PyArrow table to write.

buffer

A memoryview providing the destination memory.

Returns -----= Number of bytes written.

Raises

-----=

ValueError

If the buffer is too small.