<commit_before># encoding: utf-8

module Ktl
  class ShufflePlan
    def initialize(zk_client, options = {})
      @zk_client = zk_client
      @options = options
    end

    def generate
      topics = @zk_client.all_topics
      if (filter = @options[:filter])
        topics = topics.filter { |t| !!t.match(filter) }
      end
      topics_partitions = ScalaEnumerable.new(@zk_client.partitions_for_topics(topics))
      topics_partitions = topics_partitions.sort_by(&:first)
      replica_assignments = @zk_client.replica_assignment_for_topics(topics)
      brokers = select_brokers
      reassignment_plan = Scala::Collection::Map.empty
      topics_partitions.each do |tp|
        topic, partitions = tp.elements
        nr_replicas = replica_assignments.apply(Kafka::TopicAndPartition.new(topic, 0)).size
        assignment = assign_replicas_to_brokers(topic, brokers, partitions.size, nr_replicas)
        assignment.each do |pr|
          partition, replicas = pr.elements
          topic_partition = Kafka::TopicAndPartition.new(topic, partition)
          current_assignment = replica_assignments.apply(topic_partition)
          unless current_assignment == replicas
            reassignment_plan += Scala::Tuple.new(topic_partition, replicas)
          end
        end
      end
      reassignment_plan
    end

    private

    def select_brokers
      brokers = @options[:brokers] ? Array(@options[:brokers]) : ScalaEnumerable.new(@zk_client.broker_ids).to_a
      brokers -= Array(@options[:blacklist]) if @options[:blacklist]
      brokers
    end

    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
      brokers = Scala::Collection::JavaConversions.as_scala_iterable(brokers.map { |x| x.to_java(:int) }).to_list
      Kafka::Admin.assign_replicas_to_brokers(brokers, partition_count, replica_count)
    rescue Kafka::Admin::AdminOperationException => e
      raise ArgumentError, sprintf('%s (%s)', e.message, e.class.name), e.backtrace
    end
  end

  class RendezvousShufflePlan < ShufflePlan
    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
      if replica_count > brokers.size
        raise ArgumentError, sprintf('replication factor: %i larger than available brokers: %i', replica_count, brokers.size)
      end
      result = []
      partition_count.times do |partition|
        sorted = brokers.sort_by do |broker|
          key = [partition, topic, broker].pack('l<a*l<')
          Java::OrgJrubyUtil::MurmurHash.hash32(key.to_java_bytes, 0, key.bytesize, SEED)
        end
        selected = sorted.take(replica_count)
        result.push(Scala::Tuple.new(partition, Scala::Collection::JavaConversions.as_scala_iterable(selected).to_list))
      end
      result
    end

    private

    SEED = 1683520333
  end
end
<commit_msg>Ensure that broker IDs are integers<commit_after># encoding: utf-8

module Ktl
  class ShufflePlan
    def initialize(zk_client, options = {})
      @zk_client = zk_client
      @options = options
    end

    def generate
      topics = @zk_client.all_topics
      if (filter = @options[:filter])
        topics = topics.filter { |t| !!t.match(filter) }
      end
      topics_partitions = ScalaEnumerable.new(@zk_client.partitions_for_topics(topics))
      topics_partitions = topics_partitions.sort_by(&:first)
      replica_assignments = @zk_client.replica_assignment_for_topics(topics)
      brokers = select_brokers
      reassignment_plan = Scala::Collection::Map.empty
      topics_partitions.each do |tp|
        topic, partitions = tp.elements
        nr_replicas = replica_assignments.apply(Kafka::TopicAndPartition.new(topic, 0)).size
        assignment = assign_replicas_to_brokers(topic, brokers, partitions.size, nr_replicas)
        assignment.each do |pr|
          partition, replicas = pr.elements
          topic_partition = Kafka::TopicAndPartition.new(topic, partition)
          current_assignment = replica_assignments.apply(topic_partition)
          unless current_assignment == replicas
            reassignment_plan += Scala::Tuple.new(topic_partition, replicas)
          end
        end
      end
      reassignment_plan
    end

    private

    def select_brokers
      brokers = @options[:brokers] ? Array(@options[:brokers]).map(&:to_i) : ScalaEnumerable.new(@zk_client.broker_ids).to_a
      brokers -= Array(@options[:blacklist]).map(&:to_i) if @options[:blacklist]
      brokers
    end

    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
      brokers = Scala::Collection::JavaConversions.as_scala_iterable(brokers.map { |x| x.to_java(:int) }).to_list
      Kafka::Admin.assign_replicas_to_brokers(brokers, partition_count, replica_count)
    rescue Kafka::Admin::AdminOperationException => e
      raise ArgumentError, sprintf('%s (%s)', e.message, e.class.name), e.backtrace
    end
  end

  class RendezvousShufflePlan < ShufflePlan
    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
      if replica_count > brokers.size
        raise ArgumentError, sprintf('replication factor: %i larger than available brokers: %i', replica_count, brokers.size)
      end
      result = []
      partition_count.times do |partition|
        sorted = brokers.sort_by do |broker|
          key = [partition, topic, broker].pack('l<a*l<')
          Java::OrgJrubyUtil::MurmurHash.hash32(key.to_java_bytes, 0, key.bytesize, SEED)
        end
        selected = sorted.take(replica_count)
        result.push(Scala::Tuple.new(partition, Scala::Collection::JavaConversions.as_scala_iterable(selected).to_list))
      end
      result
    end

    private

    SEED = 1683520333
  end
end

{'fraction_non_alphanumeric': 0.09789473684210526, 'fraction_numerical': 0.005263157894736842, 'mean_word_length': 3.433125972006221, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, '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': '202495', 'n_tokens_mistral': 1811, 'n_tokens_neox': 1761, 'n_words': 333}