kafka

KafkaReader

KafkaReader(url, metadata, start=None)

Bases: StreamReader

A connection object to read data from Kafka.

Parameters:

Name	Type	Description	Default
url	str	URL of Kafka broker to connect to.	required
metadata	dict[str, Channel]]	Dictionary of channel metadata for request.	required
start	int | None	GPS start time of stream in nanoseconds, defaults to "now".	None

Source code in arrakis/kafka.py

def __init__(
    self,
    url: str,
    metadata: dict[str, Channel],
    start: int | None = None,
):
    self.url = url
    self.metadata = metadata
    self.start = start

    # track stream -> index -> channel
    self._name_lookup: dict[str, dict[int, Channel]] = defaultdict(defaultdict)
    # filters for partitions
    self._partition_filter_id_sets = defaultdict(set)

    for channel in self.metadata.values():
        partition_id = channel.partition_id
        assert partition_id is not None
        partition_index = channel.partition_index
        assert partition_index is not None
        stream_id = self._id(partition_id)
        self._name_lookup[stream_id][partition_index] = channel
        self._partition_filter_id_sets[partition_id].add(partition_index)

    # pre-compute Arrow arrays for filtering to avoid repeated
    # array creation
    self._partition_filter_id_arrays = {
        partition_id: pyarrow.array(channel_id, type=pyarrow.int32())
        for partition_id, channel_id in self._partition_filter_id_sets.items()
    }

    # Optimize IPC stream reader creation/destruction overhead
    self._ipc_options = pyarrow.ipc.IpcReadOptions(use_threads=False)

    # create Kafka consumer
    consumer_settings = {
        "bootstrap.servers": url,
        "group.id": generate_groupid(),
        "message.max.bytes": 10_000_000,  # 10 MB
        "queued.min.messages": 100,
        "queued.max.messages.kbytes": 4096,  # 4 MB
        "enable.auto.commit": False,
    }
    # librdkafka internal debugging, e.g. "cgrp,topic,metadata"
    # (add "fetch" only for short sessions: it logs per request).
    # Diagnostic switch — not intended as a standing setting.
    kafka_debug = os.getenv("ARRAKIS_KAFKA_DEBUG")
    if kafka_debug:
        consumer_settings["debug"] = kafka_debug
    # route librdkafka logs through the arrakis logger so they
    # respect the application's logging configuration
    self._consumer = Consumer(consumer_settings, logger=logger)
    logger.debug("Kafka consumer: %s", consumer_settings)
    # topics requested but not present on the broker; their
    # channels are served as gaps until the topics appear
    self._missing_topics: set[str] = set()
    self._next_missing_check = 0.0
    self._topic_to_partition: dict[str, str] = {}
    for partition_id in self._partition_filter_id_sets:
        channel = next(
            ch for ch in self.metadata.values() if ch.partition_id == partition_id
        )
        t = topic_name(partition_id, channel.replay_id)
        self._topic_to_partition[t] = partition_id
    self._topics = list(self._topic_to_partition)
    logger.debug("kafka topics: %s", self._topics)

enter

enter()

Assign the consumer to the requested topics.

Partitions are assigned directly rather than subscribed via a consumer group: every reader is a single-member group with no committed offsets, so group membership buys nothing and costs the coordinator dependency (slow or failing joins starve the reader entirely) plus rebalance churn when any requested topic does not exist. Missing topics are excluded — their channels are served as gaps by the muxer — and re-checked periodically so data starts flowing if a topic appears.

Source code in arrakis/kafka.py

def enter(self) -> None:
    """Assign the consumer to the requested topics.

    Partitions are assigned directly rather than subscribed via a
    consumer group: every reader is a single-member group with no
    committed offsets, so group membership buys nothing and costs
    the coordinator dependency (slow or failing joins starve the
    reader entirely) plus rebalance churn when any requested topic
    does not exist.  Missing topics are excluded — their channels
    are served as gaps by the muxer — and re-checked periodically
    so data starts flowing if a topic appears.
    """
    logger.debug("initiating kafka assignments...")
    now = time_as_ns(gpsnow())
    existing = self._existing_topics()
    missing = set(self._topics) - set(existing)
    if missing:
        logger.warning(
            "%d requested topic(s) do not exist on the broker and "
            "will be served as gaps until they appear: %s",
            len(missing),
            sorted(missing),
        )
    self._missing_topics = missing
    self._next_missing_check = time.monotonic() + self.MISSING_TOPIC_RECHECK

    # if start time is specified, point the assignment at the
    # data requested
    if self.start and self.start <= now:
        if not existing:
            return
        # convert to UNIX time in ms
        offset_time = int(gps2unix(self.start // Time.SECONDS) * 1000)
        # get offsets corresponding to times
        partitions = [
            TopicPartition(topic, partition=0, offset=offset_time)
            for topic in existing
        ]
        partitions = self._offsets_for_times(partitions)
        # An unresolved offset (no message at/after the requested
        # time yet — e.g. a fresh topic whose data is still being
        # written) must not decay to auto.offset.reset (latest),
        # which would silently skip data: start from the beginning
        # instead; the muxer drops blocks older than requested.
        for tp in partitions:
            if tp.offset < 0:
                logger.debug(
                    "no offset at requested time for %s; "
                    "starting from the beginning",
                    tp.topic,
                )
                tp.offset = OFFSET_BEGINNING
        self._consumer.assign(partitions)

    # FIXME: start times in the future are being handled as if a
    # start time was not specified
    elif existing:
        self._consumer.assign(
            [
                TopicPartition(topic, partition=0, offset=OFFSET_END)
                for topic in existing
            ]
        )

generate_groupid

generate_groupid()

Generate a random Kafka group ID.

Source code in arrakis/kafka.py

def generate_groupid() -> str:
    """Generate a random Kafka group ID."""
    return random_alphanum(16)

random_alphanum

random_alphanum(n)

Generate a random alpha-numeric sequence of N characters.

Source code in arrakis/kafka.py

def random_alphanum(n: int) -> str:
    """Generate a random alpha-numeric sequence of N characters."""
    alphanum = string.ascii_uppercase + string.digits
    return "".join(random.SystemRandom().choice(alphanum) for _ in range(n))