"""Kafka datasource for bounded data reads. This module provides a Kafka datasource implementation for Ray Data that supports bounded reads with offset-based range queries. Message keys and values are returned as raw bytes to support any serialization format (JSON, Avro, Protobuf, etc.). Users can decode them using map operations. Requires: - confluent-kafka: https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/ """ import logging import time import warnings from dataclasses import dataclass from datetime import datetime, timezone from typing import ( TYPE_CHECKING, Any, Dict, Iterable, List, Literal, Optional, Set, Tuple, Union, ) import pyarrow as pa if TYPE_CHECKING: from confluent_kafka import Consumer, TopicPartition from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption from ray.data._internal.util import _check_import from ray.data.block import Block, BlockMetadata from ray.data.context import DataContext from ray.data.datasource import Datasource, ReadTask PartitionOffsets = Dict[int, Union[int, str]] PerPartitionOffsets = Dict[str, PartitionOffsets] logger = logging.getLogger(__name__) # Mapping from kafka-python style KafkaAuthConfig fields to Confluent/librdkafka config. # TODO(youcheng): Remove this mapping and use consumer_config directly. _KAFKA_AUTH_TO_CONFLUENT: Dict[str, str] = { "security_protocol": "security.protocol", "sasl_mechanism": "sasl.mechanism", "sasl_plain_username": "sasl.username", "sasl_plain_password": "sasl.password", "sasl_kerberos_service_name": "sasl.kerberos.service.name", "sasl_kerberos_name": "sasl.kerberos.principal", "ssl_cafile": "ssl.ca.location", "ssl_certfile": "ssl.certificate.location", "ssl_keyfile": "ssl.key.location", "ssl_password": "ssl.key.password", "ssl_ciphers": "ssl.cipher.suites", "ssl_crlfile": "ssl.crl.location", # Note: ssl_check_hostname is intentionally NOT mapped due to semantics mismatch. } KAFKA_TOPIC_METADATA_TIMEOUT_S = 10 KAFKA_QUERY_OFFSET_TIMEOUT_S = 10 # Cap each consume timeout to keep responsiveness of timeout/position checks. KAFKA_CONSUME_TIMEOUT_MAX_S = 10 # 10 seconds per consume call KAFKA_MSG_SCHEMA = pa.schema( [ ("offset", pa.int64()), ("key", pa.binary()), ("value", pa.binary()), ("topic", pa.string()), ("partition", pa.int32()), ("timestamp", pa.int64()), # Kafka timestamp in milliseconds ( "timestamp_type", pa.int32(), ), # 0=TIMESTAMP_NOT_AVAILABLE, 1=TIMESTAMP_CREATE_TIME, 2=TIMESTAMP_LOG_APPEND_TIME ("headers", pa.map_(pa.string(), pa.binary())), # Message headers ] ) @dataclass class KafkaAuthConfig: """Authentication configuration for Kafka connections. Uses standard kafka-python parameter names. See kafka-python documentation for full details: https://kafka-python.readthedocs.io/ Note: Ray Data maps these options to Confluent/librdkafka config under the hood. Some options have different semantics or are unsupported by the Confluent client; see notes below for those fields. Prefer passing Confluent options directly via ``consumer_config`` where possible. security_protocol: Protocol used to communicate with brokers. Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL. Default: PLAINTEXT. sasl_mechanism: Authentication mechanism when security_protocol is configured for SASL_PLAINTEXT or SASL_SSL. Valid values are: PLAIN, GSSAPI, OAUTHBEARER, SCRAM-SHA-256, SCRAM-SHA-512. sasl_plain_username: username for sasl PLAIN and SCRAM authentication. Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms. sasl_plain_password: password for sasl PLAIN and SCRAM authentication. Required if sasl_mechanism is PLAIN or one of the SCRAM mechanisms. sasl_kerberos_name: Constructed gssapi.Name for use with sasl mechanism handshake. If provided, sasl_kerberos_service_name and sasl_kerberos_domain name are ignored. Default: None. sasl_kerberos_service_name: Service name to include in GSSAPI sasl mechanism handshake. Default: 'kafka' sasl_kerberos_domain_name: kerberos domain name to use in GSSAPI sasl mechanism handshake. Default: one of bootstrap servers. Note (Confluent): This option is not supported by Confluent/librdkafka and will be ignored when building the client configuration. Prefer specifying an explicit principal via ``sasl_kerberos_name`` or rely on defaults. sasl_oauth_token_provider: OAuthBearer token provider instance. Default: None. Note (Confluent): Not supported directly; use ``consumer_config`` with ``sasl.oauthbearer.*`` options instead. ssl_context: Pre-configured SSLContext for wrapping socket connections. If provided, all other ssl_* configurations will be ignored. Default: None. Note (Confluent): Passing an SSLContext object is not supported and will be ignored. Use ``ssl_cafile``, ``ssl_certfile``, and ``ssl_keyfile`` instead. ssl_check_hostname: Flag to configure whether ssl handshake should verify that the certificate matches the broker's hostname. Default: True. Note (Confluent): There is no 1:1 equivalent; disabling hostname verification via ``enable.ssl.certificate.verification=False`` would also disable the entire certificate chain verification. To avoid weakening security, this flag is not mapped when False. If you need to disable only hostname verification, set ``ssl.endpoint.identification.algorithm=none`` via ``consumer_config`` (if supported by your librdkafka version). ssl_cafile: Optional filename of ca file to use in certificate verification. Default: None. ssl_certfile: Optional filename of file in pem format containing the client certificate, as well as any ca certificates needed to establish the certificate's authenticity. Default: None. ssl_keyfile: Optional filename containing the client private key. Default: None. ssl_password: Optional password to be used when loading the certificate chain. Default: None. ssl_crlfile: Optional filename containing the CRL to check for certificate expiration. By default, no CRL check is done. When providing a file, only the leaf certificate will be checked against this CRL. The CRL can only be checked with Python 3.4+ or 2.7.9+. Default: None. ssl_ciphers: optionally set the available ciphers for ssl connections. It should be a string in the OpenSSL cipher list format. If no cipher can be selected (because compile-time options or other configuration forbids use of all the specified ciphers), an ssl.SSLError will be raised. See ssl.SSLContext.set_ciphers. """ # Security protocol security_protocol: Optional[str] = None # SASL configuration sasl_mechanism: Optional[str] = None sasl_plain_username: Optional[str] = None sasl_plain_password: Optional[str] = None sasl_kerberos_name: Optional[str] = None sasl_kerberos_service_name: Optional[str] = None sasl_kerberos_domain_name: Optional[str] = None sasl_oauth_token_provider: Optional[Any] = None # SSL configuration ssl_context: Optional[Any] = None ssl_check_hostname: Optional[bool] = None ssl_cafile: Optional[str] = None ssl_certfile: Optional[str] = None ssl_keyfile: Optional[str] = None ssl_password: Optional[str] = None ssl_ciphers: Optional[str] = None ssl_crlfile: Optional[str] = None def _handle_deprecated_configs(kafka_auth_config: KafkaAuthConfig) -> None: # Handle special fields with warnings if kafka_auth_config.ssl_context is not None: logger.warning( "ssl_context is not supported by Confluent. Skipping. " "Use KafkaAuthConfig fields ssl_cafile, ssl_certfile, ssl_keyfile instead." ) if kafka_auth_config.sasl_oauth_token_provider is not None: logger.warning( "sasl_oauth_token_provider is not supported by Confluent. Skipping. " "Use consumer_config with sasl.oauthbearer.* options instead." ) if kafka_auth_config.sasl_kerberos_domain_name is not None: logger.warning( "sasl_kerberos_domain_name is not supported by Confluent and will be ignored. " "Set sasl_kerberos_name (principal) or rely on defaults." ) if kafka_auth_config.ssl_check_hostname is False: logger.warning( "ssl_check_hostname=False cannot be mapped safely to Confluent; " "setting enable.ssl.certificate.verification=False would disable all certificate verification. " "Ignoring ssl_check_hostname. If you need to disable only hostname verification, " "configure the client directly via consumer_config (e.g., ssl.endpoint.identification.algorithm=none)." ) def _add_authentication_to_config( config: Dict[str, Any], kafka_auth_config: Optional[KafkaAuthConfig] ) -> None: """Map KafkaAuthConfig (kafka-python style) into Confluent/librdkafka config. Special cases: - ssl_context: unsupported; warn and ignore - sasl_oauth_token_provider: unsupported; warn and ignore - sasl_kerberos_domain_name: unsupported; warn and ignore - ssl_check_hostname: not mapped due to semantics; if False, warn and ignore """ if not kafka_auth_config: return warnings.warn( "kafka_auth_config (kafka-python style) is deprecated and will be removed in a future release. " "Please provide Confluent/librdkafka options via consumer_config instead.", DeprecationWarning, stacklevel=2, ) _handle_deprecated_configs(kafka_auth_config) # Map directly compatible fields for key, confluent_key in _KAFKA_AUTH_TO_CONFLUENT.items(): val = getattr(kafka_auth_config, key, None) if val is not None: config[confluent_key] = val def _build_confluent_config( bootstrap_servers: List[str], kafka_auth_config: Optional[KafkaAuthConfig] = None, extra: Optional[Dict[str, Any]] = None, user_config: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Build Confluent config with bootstrap servers and auth. Args: bootstrap_servers: List of Kafka broker addresses. kafka_auth_config: Authentication configuration (kafka-python style). Deprecated; prefer consumer_config with Confluent keys. Mutually exclusive with consumer_config. extra: Additional config options. user_config: User-provided config options. Returns: Confluent configuration dict. """ config: Dict[str, Any] = { "bootstrap.servers": ",".join(bootstrap_servers), } # Map kafka-python-style auth if provided _add_authentication_to_config(config, kafka_auth_config) if extra: config.update(extra) if user_config: if ( "bootstrap.servers" in user_config and user_config["bootstrap.servers"] != config["bootstrap.servers"] ): logger.warning( "Ignoring 'bootstrap.servers' from consumer_config; use bootstrap_servers parameter instead." ) for k, v in user_config.items(): if k == "bootstrap.servers": continue config[k] = v return config def _build_consumer_config_for_read( bootstrap_servers: List[str], kafka_auth_config: Optional[KafkaAuthConfig] = None, consumer_config: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Build Consumer config for reading messages (Confluent).""" return _build_confluent_config( bootstrap_servers, extra={ "enable.auto.commit": False, # Confluent requires a group.id even when using manual assign. "group.id": "ray-data-kafka-reader", }, user_config=consumer_config, kafka_auth_config=kafka_auth_config, ) def _datetime_to_ms(dt: datetime) -> int: """Convert a datetime to milliseconds since epoch (UTC). If the datetime has no timezone info (i.e., ``tzinfo is None``), it is assumed to be UTC. Timezone-aware datetimes are converted to UTC automatically via ``datetime.timestamp()``. Args: dt: A datetime object, with or without timezone info. Returns: Milliseconds since Unix epoch. """ if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) return int(dt.timestamp() * 1000) def _validate_offsets( start_offset: Union[int, datetime, Literal["earliest"], PerPartitionOffsets], end_offset: Union[int, datetime, Literal["latest"], PerPartitionOffsets], ) -> None: if isinstance(start_offset, dict): for topic, partition_map in start_offset.items(): if not isinstance(partition_map, dict): raise ValueError( f"start_offset[{topic!r}] must be a dict mapping " f"partition_id (int) to offset (int or str)." ) for partition_id, offset in partition_map.items(): if not isinstance(partition_id, int): raise ValueError( f"start_offset[{topic!r}] keys must be integers " f"(partition IDs), got {type(partition_id).__name__!r}." ) if isinstance(offset, str) and offset == "latest": raise ValueError( f"start_offset[{topic!r}][{partition_id}] cannot be 'latest'." ) else: if isinstance(start_offset, int) and isinstance(end_offset, int): if start_offset > end_offset: raise ValueError("start_offset must be less than end_offset") if isinstance(start_offset, datetime) and isinstance(end_offset, datetime): if _datetime_to_ms(start_offset) > _datetime_to_ms(end_offset): raise ValueError("start_offset must be less than end_offset") if isinstance(start_offset, str) and start_offset == "latest": raise ValueError("start_offset cannot be 'latest'") if isinstance(end_offset, dict): for topic, partition_map in end_offset.items(): if not isinstance(partition_map, dict): raise ValueError( f"end_offset[{topic!r}] must be a dict mapping " f"partition_id (int) to offset (int or str)." ) for partition_id, offset in partition_map.items(): if not isinstance(partition_id, int): raise ValueError( f"end_offset[{topic!r}] keys must be integers " f"(partition IDs), got {type(partition_id).__name__!r}." ) if isinstance(offset, str) and offset == "earliest": raise ValueError( f"end_offset[{topic!r}][{partition_id}] cannot be 'earliest'." ) else: if isinstance(end_offset, str) and end_offset == "earliest": raise ValueError("end_offset cannot be 'earliest'") def _resolve_datetime_to_offset( consumer: "Consumer", topic_partition: "TopicPartition", dt: datetime, fallback_offset: int, ) -> int: """Resolve a datetime to an integer offset via offsets_for_times. Returns fallback_offset when offsets_for_times returns empty or offset < 0 (e.g., datetime in future or no messages at that time). """ from confluent_kafka import TopicPartition timestamp_ms = _datetime_to_ms(dt) tp_with_ts = TopicPartition( topic_partition.topic, topic_partition.partition, timestamp_ms ) result = consumer.offsets_for_times( [tp_with_ts], timeout=KAFKA_QUERY_OFFSET_TIMEOUT_S ) if result and result[0].offset >= 0: return result[0].offset return fallback_offset def _resolve_offsets( consumer: "Consumer", topic_partition: "TopicPartition", start_offset: Union[int, datetime, Literal["earliest"]], end_offset: Union[int, datetime, Literal["latest"]], ) -> Tuple[int, int]: """Resolve start and end offsets to actual integer offsets. Handles int offsets, "earliest"/"latest" strings, and datetime objects. For datetime objects, uses ``consumer.offsets_for_times()`` to find the earliest offset whose timestamp is >= the given datetime. Args: consumer: Confluent Kafka consumer instance. topic_partition: TopicPartition to resolve offsets for. start_offset: Start offset (int, datetime, or "earliest"). end_offset: End offset (int, datetime, or "latest"). Returns: Tuple of (resolved_start_offset, resolved_end_offset). """ # TODO(youcheng): add retry logic for this call. low, high = consumer.get_watermark_offsets( topic_partition, timeout=KAFKA_QUERY_OFFSET_TIMEOUT_S ) earliest_offset = low latest_offset = high # Keep original values for error messages original_start = start_offset original_end = end_offset if start_offset == "earliest" or start_offset is None: start_offset = earliest_offset elif isinstance(start_offset, datetime): # fallback to latest_offset if the start_offset is in the future, so the read range is empty (start == end). start_offset = _resolve_datetime_to_offset( consumer, topic_partition, start_offset, latest_offset ) if end_offset == "latest" or end_offset is None: end_offset = latest_offset elif isinstance(end_offset, datetime): end_offset = _resolve_datetime_to_offset( consumer, topic_partition, end_offset, latest_offset ) # Clamp end_offset to the high watermark so we never try to read beyond # what is currently available. This prevents the read loop from polling # indefinitely when a user-supplied integer end_offset exceeds the number # of messages in the partition. end_offset = min(end_offset, latest_offset) if isinstance(start_offset, int) and start_offset < earliest_offset: logger.warning( f"start_offset ({start_offset}) is below the earliest available offset " f"({earliest_offset}) for partition {topic_partition.partition} in topic " f"{topic_partition.topic} (data may have been deleted by Kafka retention). " f"Falling back to earliest available offset ({earliest_offset})." ) start_offset = earliest_offset if start_offset > end_offset: start_str = ( f"{original_start}" if original_start == start_offset else f"{original_start} (resolved to {start_offset})" ) end_str = ( f"{original_end}" if original_end == end_offset else f"{original_end} (resolved to {end_offset})" ) raise ValueError( f"start_offset ({start_str}) > end_offset ({end_str}) " f"for partition {topic_partition.partition} in topic {topic_partition.topic}" ) return start_offset, end_offset class KafkaDatasource(Datasource): """Kafka datasource for reading from Kafka topics with bounded reads.""" # Batch size for incremental block yielding BATCH_SIZE_FOR_YIELD = 1000 def __init__( self, topics: Union[str, List[str]], bootstrap_servers: Union[str, List[str]], start_offset: Union[int, datetime, Literal["earliest"], PerPartitionOffsets], end_offset: Union[int, datetime, Literal["latest"], PerPartitionOffsets], kafka_auth_config: Optional[KafkaAuthConfig] = None, consumer_config: Optional[Dict[str, Any]] = None, timeout_ms: Optional[int] = None, ): """Initialize Kafka datasource. Args: topics: Kafka topic name(s) to read from. bootstrap_servers: Kafka broker addresses (string or list of strings). start_offset: Starting position. Can be: - int: Offset number - datetime: Read from the first message at or after this time. datetimes with no timezone info are treated as UTC. - str: "earliest" end_offset: Ending position (exclusive). Can be: - int: Offset number - datetime: Read up to (but not including) the first message at or after this time. datetimes with no timezone info are treated as UTC. - str: "latest" kafka_auth_config: Authentication configuration (kafka-python style). Deprecated; prefer consumer_config with Confluent keys. Mutually exclusive with consumer_config. consumer_config: Confluent/librdkafka consumer configuration dict. Keys and values are passed through to the underlying client. The `bootstrap.servers` option is derived from `bootstrap_servers` and cannot be overridden here. timeout_ms: Optional timeout in milliseconds for every read task to poll until reaching end_offset. If None (default), no task-level timeout is applied and each read task will poll until it reaches end_offset. If set, the read task will stop polling after the timeout and return the messages it has read so far. Raises: ValueError: If required configuration is missing. ImportError: If confluent-kafka is not installed. """ _check_import(self, module="confluent_kafka", package="confluent-kafka") if not topics: raise ValueError("topics cannot be empty") if not bootstrap_servers: raise ValueError("bootstrap_servers cannot be empty") if timeout_ms is not None and timeout_ms <= 0: raise ValueError("timeout_ms must be positive") _validate_offsets(start_offset, end_offset) # Validate bootstrap_servers format if isinstance(bootstrap_servers, str): if not bootstrap_servers or ":" not in bootstrap_servers: raise ValueError( f"Invalid bootstrap_servers format: {bootstrap_servers}. " "Expected 'host:port' or list of 'host:port' strings." ) elif isinstance(bootstrap_servers, list): if not bootstrap_servers: raise ValueError("bootstrap_servers cannot be empty list") for server in bootstrap_servers: if not isinstance(server, str) or ":" not in server: raise ValueError( f"Invalid bootstrap_servers format: {server}. " "Expected 'host:port' string." ) # Disallow specifying both config styles at once to avoid ambiguity. if kafka_auth_config is not None and consumer_config is not None: raise ValueError( "Provide only one of kafka_auth_config (deprecated) or consumer_config, not both." ) self._topics = topics if isinstance(topics, list) else [topics] self._bootstrap_servers = ( bootstrap_servers if isinstance(bootstrap_servers, list) else [bootstrap_servers] ) self._start_offset = start_offset self._end_offset = end_offset self._kafka_auth_config = kafka_auth_config self._consumer_config = consumer_config self._timeout_ms = timeout_ms self._target_max_block_size = DataContext.get_current().target_max_block_size def estimate_inmemory_data_size(self) -> Optional[int]: """Return an estimate of the in-memory data size, or None if unknown.""" return None def get_read_tasks( self, parallelism: int, per_task_row_limit: Optional[int] = None, data_context: Optional["DataContext"] = None, ) -> List[ReadTask]: """Create read tasks for Kafka partitions. Creates one read task per partition. Each task reads from a single partition of a single topic. Args: parallelism: This argument is deprecated. per_task_row_limit: Maximum number of rows per read task. data_context: The data context to use to get read tasks. This is not used by this datasource. Returns: List of ReadTask objects, one per partition. """ from confluent_kafka import Consumer consumer_config = _build_consumer_config_for_read( self._bootstrap_servers, self._kafka_auth_config, self._consumer_config ) discovery_consumer = Consumer(consumer_config) try: metadata = discovery_consumer.list_topics( timeout=KAFKA_TOPIC_METADATA_TIMEOUT_S ) topic_partitions: List[Tuple[str, int]] = [] for topic in self._topics: if topic not in metadata.topics: raise ValueError( f"Topic {topic} has no partitions or doesn't exist" ) topic_meta = metadata.topics[topic] if not topic_meta.partitions: raise ValueError( f"Topic {topic} has no partitions or doesn't exist" ) for partition_id in topic_meta.partitions.keys(): topic_partitions.append((topic, partition_id)) finally: discovery_consumer.close() bootstrap_servers = self._bootstrap_servers start_offset = self._start_offset end_offset = self._end_offset timeout_ms = self._timeout_ms target_max_block_size = self._target_max_block_size # Validate that any partitions referenced in a per-partition dict # actually exist on the broker. Check once per topic before the loop. actual_partition_ids: Dict[str, Set[int]] = {} for topic_name, partition_id in topic_partitions: actual_partition_ids.setdefault(topic_name, set()).add(partition_id) for param_name, offset in ( ("start_offset", start_offset), ("end_offset", end_offset), ): if isinstance(offset, dict): for topic, partition_map in offset.items(): existing_partitions = actual_partition_ids.get(topic, set()) for pid in partition_map: if pid not in existing_partitions: raise ValueError( f"{param_name} references partition {pid} in topic " f"{topic!r}, but that partition does not exist." ) tasks = [] for topic_name, partition_id in topic_partitions: def create_kafka_read_fn( topic_name: str = topic_name, partition_id: int = partition_id, bootstrap_servers: List[str] = bootstrap_servers, start_offset: Optional[ Union[int, datetime, Literal["earliest"]] ] = start_offset, end_offset: Optional[ Union[int, datetime, Literal["latest"]] ] = end_offset, kafka_auth_config: Optional[KafkaAuthConfig] = self._kafka_auth_config, user_consumer_config: Optional[Dict[str, Any]] = self._consumer_config, timeout_ms: Optional[int] = timeout_ms, target_max_block_size: int = target_max_block_size, ): """Create a Kafka read function with captured variables.""" def kafka_read_fn() -> Iterable[Block]: """Read function for a single Kafka partition using confluent-kafka.""" from confluent_kafka import ( Consumer, KafkaError, KafkaException, TopicPartition, ) built_consumer_config = _build_consumer_config_for_read( bootstrap_servers, kafka_auth_config, user_consumer_config ) consumer = Consumer(built_consumer_config) try: topic_partition = TopicPartition(topic_name, partition_id) resolved_start, resolved_end = _resolve_offsets( consumer, topic_partition, start_offset, end_offset ) records = [] output_buffer = BlockOutputBuffer( OutputBlockSizeOption.of( target_max_block_size=target_max_block_size ) ) if resolved_start < resolved_end: start_time = time.perf_counter() timeout_seconds = ( timeout_ms / 1000.0 if timeout_ms is not None else None ) tp_with_offset = TopicPartition( topic_name, partition_id, resolved_start ) consumer.assign([tp_with_offset]) next_offset = resolved_start partition_done = False while not partition_done: if next_offset >= resolved_end: break if timeout_seconds is not None: elapsed_time_s = time.perf_counter() - start_time if elapsed_time_s >= timeout_seconds: logger.warning( f"Kafka read task timed out after {timeout_ms}ms while reading partition {partition_id} of topic {topic_name}; " f"end_offset {resolved_end} was not reached. Returning {len(records)} messages collected in this read task so far." ) break remaining_timeout_s = ( timeout_seconds - elapsed_time_s ) consume_timeout_s = min( remaining_timeout_s, KAFKA_CONSUME_TIMEOUT_MAX_S ) else: consume_timeout_s = KAFKA_CONSUME_TIMEOUT_MAX_S remaining = resolved_end - next_offset batch_size = min( remaining, KafkaDatasource.BATCH_SIZE_FOR_YIELD, ) msgs = consumer.consume( num_messages=batch_size, timeout=consume_timeout_s, ) if not msgs: continue for msg in msgs: if msg.error(): # In confluent-kafka, errors are delivered # as messages. Only skip partition EOF # events, raise others. err = msg.error() if err.code() == KafkaError._PARTITION_EOF: continue raise KafkaException(err) # Stop once we reached the end offset # (exclusive). if msg.offset() >= resolved_end: partition_done = True break ts_type, ts_ms = msg.timestamp() headers_list = msg.headers() or [] headers_dict = dict(headers_list) records.append( { "offset": msg.offset(), "key": msg.key(), "value": msg.value(), "topic": msg.topic(), "partition": msg.partition(), "timestamp": ts_ms, "timestamp_type": ts_type, "headers": headers_dict, } ) next_offset = msg.offset() + 1 if ( len(records) >= KafkaDatasource.BATCH_SIZE_FOR_YIELD ): table = pa.Table.from_pylist(records) output_buffer.add_block(table) yield from output_buffer.iter_ready_blocks() records = [] # Yield any remaining records if records: table = pa.Table.from_pylist(records) output_buffer.add_block(table) output_buffer.finalize() yield from output_buffer.iter_ready_blocks() finally: consumer.close() return kafka_read_fn # TODO: We could output the offset range for every partition after the read is done. metadata = BlockMetadata( num_rows=None, size_bytes=None, input_files=[f"kafka://{topic_name}/{partition_id}"], exec_stats=None, ) effective_start = ( start_offset.get(topic_name, {}).get(partition_id, "earliest") if isinstance(start_offset, dict) else start_offset ) effective_end = ( end_offset.get(topic_name, {}).get(partition_id, "latest") if isinstance(end_offset, dict) else end_offset ) kafka_read_fn = create_kafka_read_fn( topic_name, partition_id, start_offset=effective_start, end_offset=effective_end, ) # Create read task task = ReadTask( read_fn=kafka_read_fn, metadata=metadata, schema=KAFKA_MSG_SCHEMA, per_task_row_limit=per_task_row_limit, ) tasks.append(task) return tasks