Repository: ray-project/ray
File: python/ray/data/_internal/datasource/kafka_datasource.py (Python)

Code before:
        consumer: 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).
    """
    earliest_offset = consumer.beginning_offsets([topic_partition])[topic_partition]
    latest_offset = consumer.end_offsets([topic_partition])[topic_partition]

    # 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):
        timestamp_ms = datetime_to_ms(start_offset)
        result = consumer.offsets_for_times({topic_partition: timestamp_ms})
        offset_and_ts = result.get(topic_partition)
        if offset_and_ts is None:
            # No messages with timestamp >= start_offset; nothing to read
            start_offset = latest_offset
        else:
            start_offset = offset_and_ts.offset

    if end_offset == "latest" or end_offset is None:
        end_offset = latest_offset
    elif isinstance(end_offset, datetime):
        timestamp_ms = datetime_to_ms(end_offset)
        result = consumer.offsets_for_times({topic_partition: timestamp_ms})
        offset_and_ts = result.get(topic_partition)
        if offset_and_ts is None:
            # No messages with timestamp >= end_offset; read everything
            end_offset = latest_offset
        else:
            end_offset = offset_and_ts.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(

Review comment (refactor):
Nit: This code is duplicated with lines below. Consider abstracting it to avoid duplication and improve readability

Code after:
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return int(dt.timestamp() * 1000)


def _resolve_datetime_offset(
    consumer: "KafkaConsumer",
    topic_partition: "TopicPartition",
    dt: datetime,
    fallback_offset: int,
) -> int:
    """Resolve a datetime to a concrete Kafka offset.

    Uses ``consumer.offsets_for_times()`` to find the earliest offset whose
    timestamp is >= the given datetime. If no such offset exists, returns
    ``fallback_offset``.

    Args:
        consumer: Kafka consumer instance.
        topic_partition: TopicPartition to resolve the offset for.
        dt: The datetime to resolve.
        fallback_offset: Offset to return if no messages match the timestamp.

    Returns:
        The resolved integer offset.
    """
    timestamp_ms: int = _datetime_to_ms(dt)
    offsets: Dict = consumer.offsets_for_times({topic_partition: timestamp_ms})
    offset_and_ts: Optional[Tuple[int, int]] = offsets.get(topic_partition)
    if offset_and_ts is None:
        return fallback_offset
    return offset_and_ts.offset


def _resolve_offsets(
    consumer: "KafkaConsumer",
    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: 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").