<commit_before># frozen_string_literal: true

require "kafka/protocol/member_assignment"

module Kafka

  # A consumer group partition assignment strategy that assigns partitions to
  # consumers in a round-robin fashion.
  class RoundRobinAssignmentStrategy
    def initialize(cluster:)
      @cluster = cluster
    end

    # Assign the topic partitions to the group members.
    #
    # @param members [Array<String>] member ids
    # @param topics [Array<String>] topics
    # @return [Hash<String, Protocol::MemberAssignment>] a hash mapping member
    #   ids to assignments.
    def assign(members:, topics:)
      group_assignment = {}

      members.each do |member_id|
        group_assignment[member_id] = Protocol::MemberAssignment.new
      end

      topic_partitions = topics.flat_map do |topic|
        begin
          partitions = @cluster.partitions_for(topic).map(&:partition_id)
        rescue UnknownTopicOrPartition
          raise UnknownTopicOrPartition, "unknown topic #{topic}"
        end
        Array.new(partitions.count) { topic }.zip(partitions)
      end

      partitions_per_member = [1, (topic_partitions.count / members.count)].max
      partitions_for_each_member = topic_partitions.sort.each_slice(partitions_per_member).to_a

      if partitions_for_each_member.count > members.count
        leftovers = partitions_for_each_member.pop
        partitions_for_each_member[0] = partitions_for_each_member[0] + leftovers
      end

      members.zip(partitions_for_each_member).each do |member_id, member_partitions|
        unless member_partitions.nil?
          member_partitions.group_by(&:first).each do |topic, values|
            partition_ids = values.map { |(topic, partition_id)| partition_id }
            group_assignment[member_id].assign(topic, partition_ids)
          end
        end
      end

      group_assignment
    rescue Kafka::LeaderNotAvailable
      sleep 1
      retry
    end
  end
end
<commit_msg>Simplify logic<commit_after># frozen_string_literal: true

require "kafka/protocol/member_assignment"

module Kafka

  # A consumer group partition assignment strategy that assigns partitions to
  # consumers in a round-robin fashion.
  class RoundRobinAssignmentStrategy
    def initialize(cluster:)
      @cluster = cluster
    end

    # Assign the topic partitions to the group members.
    #
    # @param members [Array<String>] member ids
    # @param topics [Array<String>] topics
    # @return [Hash<String, Protocol::MemberAssignment>] a hash mapping member
    #   ids to assignments.
    def assign(members:, topics:)
      group_assignment = {}

      members.each do |member_id|
        group_assignment[member_id] = Protocol::MemberAssignment.new
      end

      topic_partitions = topics.flat_map do |topic|
        begin
          partitions = @cluster.partitions_for(topic).map(&:partition_id)
        rescue UnknownTopicOrPartition
          raise UnknownTopicOrPartition, "unknown topic #{topic}"
        end
        Array.new(partitions.count) { topic }.zip(partitions)
      end

      partitions_per_member = topic_partitions.group_by.with_index do |_, index|
        index % members.count
      end.values

      members.zip(partitions_per_member).each do |member_id, member_partitions|
        unless member_partitions.nil?
          member_partitions.each do |topic, partition|
            group_assignment[member_id].assign(topic, [partition])
          end
        end
      end

      group_assignment
    rescue Kafka::LeaderNotAvailable
      sleep 1
      retry
    end
  end
end

{'fraction_non_alphanumeric': 0.07664335664335664, 'fraction_numerical': 0.0013986013986013986, 'mean_word_length': 3.0134680134680134, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1551424', 'n_tokens_mistral': 1032, 'n_tokens_neox': 996, 'n_words': 246}