publish

Publisher API.

Publisher

Publisher(publisher_id, url=None, replay_id=None)

Publish timeseries data to Arrakis.

Parameters:

Name	Type	Description	Default
publisher_id	str	Publisher ID string.	required
url	str | None	Initial Flight URL to connect to. Will be automatically determined if not specified.	None
replay_id	str | None	If set, publish under this replay context. Channels will be written to replay-namespaced topics.	None

Source code in arrakis/publish.py

def __init__(
    self,
    publisher_id: str,
    url: str | None = None,
    replay_id: str | None = None,
):
    """Initialize Publisher.

    Parameters
    ----------
    publisher_id : str
        Publisher ID string.
    url : str | None
        Initial Flight URL to connect to.  Will be automatically
        determined if not specified.
    replay_id : str | None
        If set, publish under this replay context.  Channels will
        be written to replay-namespaced topics.

    """
    if not HAS_KAFKA:
        msg = (
            "Publishing requires confluent-kafka to be installed. "
            "It can be installed through pip or conda."
        )
        raise ImportError(msg)

    self.publisher_id = publisher_id
    self.replay_id = replay_id
    self.initial_url = parse_url(url)

    self.channels: dict[str, Channel] = {}
    self.stride: int | None = None

    self._producer: Producer
    self._validator = RequestValidator()
    self._last_published_ns: int | None = None

close

close()

Exit publication context manager.

Source code in arrakis/publish.py

def close(self) -> None:
    """Exit publication context manager."""
    logger.info("closing kafka producer...")
    with contextlib.suppress(Exception):
        self._producer.flush()

enter

enter()

Enter publication context manager

Source code in arrakis/publish.py

def enter(self) -> None:
    """Enter publication context manager"""
    # get connection properties
    producer_info: dict[str, str] = {}
    descriptor = create_descriptor(
        RequestType.Publish,
        publisher_id=self.publisher_id,
        replay_id=self.replay_id,
        validator=self._validator,
    )
    with connect(self.initial_url) as client:
        flight_info = client.get_flight_info(descriptor)
        with MultiFlightReader(flight_info.endpoints, client) as stream:
            for data in stream.unpack():
                kv_pairs = data["properties"]
                producer_info.update(dict(kv_pairs))
    logger.info("producer info: %s", producer_info)

    # set up producer
    self._producer = Producer(
        {
            "message.max.bytes": 10_000_000,  # 10 MB
            "enable.idempotence": True,
            **producer_info,
        }
    )

publish

publish(block, timeout=constants.DEFAULT_TIMEOUT)

Publish timeseries data

Parameters:

Name	Type	Description	Default
block	SeriesBlock	A data block with all channels to publish.	required
timeout	timedelta	The maximum time to wait to publish before timing out. Default is 2 seconds.	DEFAULT_TIMEOUT

Source code in arrakis/publish.py

def publish(
    self,
    block: SeriesBlock,
    timeout: timedelta = constants.DEFAULT_TIMEOUT,
) -> None:
    """Publish timeseries data

    Parameters
    ----------
    block : SeriesBlock
        A data block with all channels to publish.
    timeout : timedelta, optional
        The maximum time to wait to publish before timing out.
        Default is 2 seconds.

    """
    if not hasattr(self, "_producer") or not self._producer:
        msg = (
            "publication interface not initialized, "
            "please use context manager when publishing."
        )
        raise RuntimeError(msg)

    # check for attempt to publish invalid channels
    # FIXME: warning for missing channels
    for name, channel in block.channels.items():
        if channel != self.channels[name]:
            msg = (
                f"channel metadata mismatch for '{name}': "
                f"got {channel!r}, expected {self.channels[name]!r}"
            )
            raise ValueError(msg)

    # check for block duration not matching expected stride
    if block.duration_ns != self.stride:
        msg = (
            "block to publish does not match expected publisher stride "
            f"got {block.duration_ns:_} ns, expected {self.stride:_} ns."
        )
        raise ValueError(msg)

    # check for non-monotonic or duplicate timestamps
    if (
        self._last_published_ns is not None
        and block.time_ns <= self._last_published_ns
    ):
        msg = (
            "block timestamp is not monotonically increasing: got "
            f"{block.time_ns:_} ns, last published {self._last_published_ns:_} ns."
        )
        raise ValueError(msg)

    # publish data for each data type
    for partition_id, batch in block.to_row_batches(self.channels):
        topic = topic_name(partition_id, self.replay_id)
        logger.debug("publishing to topic %s: %s", topic, batch)
        self._producer.produce(topic=topic, value=serialize_batch(batch))
        self._last_published_ns = block.time_ns
        self._producer.flush()

register

register(channels=None)

register channels for publication

For most publishers, channels are not specified when registering and this method will query the server for the allowable channels for this publisher and register them internally.

For publishers allowed to register their own channels ("dynamic" publishers), all channels they expect to publish should be provided as argument. Any channel that is new, or whose core metadata (sample rate or data type) has changed relative to the server's current registration for this publisher, will be partitioned dynamically via the server, which will respond with the required Kafka partition information. If the server does not permit dynamic partitioning for this publisher, a :class:pyarrow.flight.FlightUnauthorizedError is raised with a server-provided explanatory message.

Source code in arrakis/publish.py

def register(self, channels: Iterable[Channel] | None = None):
    """register channels for publication

    For most publishers, channels are not specified when
    registering and this method will query the server for the
    allowable channels for this publisher and register them
    internally.

    For publishers allowed to register their own channels
    ("dynamic" publishers), all channels they expect to publish
    should be provided as argument.  Any channel that is new, or
    whose core metadata (sample rate or data type) has changed
    relative to the server's current registration for this
    publisher, will be partitioned dynamically via the server,
    which will respond with the required Kafka partition
    information.  If the server does not permit dynamic
    partitioning for this publisher, a
    :class:`pyarrow.flight.FlightUnauthorizedError` is raised
    with a server-provided explanatory message.

    """
    requested = list(channels) if channels else []

    self._load_registered_channels()

    if requested:
        to_partition = [
            channel
            for channel in requested
            if channel.name not in self.channels
            or channel != self.channels[channel.name]
        ]
        if to_partition:
            # exhaust the generator so the server-side exchange
            # completes; any FlightUnauthorizedError (e.g.
            # PublisherNotDynamicError / PublisherUnknownError)
            # will propagate with the server's message.
            list(self._partition_channels(to_partition))
            # re-query to pick up the newly-registered channels
            # along with any server-supplied fields (stride,
            # max_latency) that are not returned by the partition
            # exchange itself.
            self._load_registered_channels()

        for channel in requested:
            if channel.name not in self.channels:
                msg = f"channel {channel.name} was not properly registered."
                raise ValueError(msg)

    if not self.channels:
        # FIXME: more informative error message here
        msg = f"unknown publisher ID '{self.publisher_id}'."
        raise ValueError(msg)

    for channel in self.channels.values():
        assert channel.partition_id is not None, (
            f"Channel {channel} is missing partition_id."
        )
        assert channel.partition_index is not None, (
            f"Channel {channel} is missing partition_index."
        )

    return self

serialize_batch

serialize_batch(batch)

Serialize a record batch to bytes.

Parameters:

Name	Type	Description	Default
batch	RecordBatch	The batch to serialize.	required

Returns:

Type	Description
bytes	The serialized buffer.

Source code in arrakis/publish.py

def serialize_batch(batch: pyarrow.RecordBatch):
    """Serialize a record batch to bytes.

    Parameters
    ----------
    batch : pyarrow.RecordBatch
        The batch to serialize.

    Returns
    -------
    bytes
        The serialized buffer.

    """
    sink = pyarrow.BufferOutputStream()
    with pyarrow.ipc.new_stream(sink, batch.schema) as writer:
        writer.write_batch(batch)
    return sink.getvalue()