<commit_before># encoding: utf-8

module Ktl
  class Reassigner
    def initialize(type, zk_client, options={})
      @type = type
      @zk_client = zk_client
      @json_max_size = options[:json_max_size] || JSON_MAX_SIZE
    end

    def reassignment_in_progress?
      partitions = @zk_client.partitions_being_reassigned
      partitions.size > 0
    end

    def overflow?
      overflow_znodes = @zk_client.get_children(overflow_base_path)
      overflow_znodes.size > 0
    rescue ZkClient::Exception::ZkNoNodeException
      false
    end

    def load_overflow
      overflow = Scala::Collection::Map.empty
      overflow_nodes = @zk_client.get_children(overflow_base_path)
      overflow_nodes.foreach do |index|
        overflow_json = @zk_client.read_data(overflow_path(index)).first
        data = parse_reassignment_json(overflow_json)
        overflow = overflow.send('++', data)
      end
      overflow
    end

    def execute(reassignment)
      json = reassignment_json(reassignment)
      reassignments = split(reassignment, json.bytesize)
      json = reassignment_json(reassignments.next)
      @zk_client.reassign_partitions(json)
      manage_overflow(reassignments)
    end

    private

    JSON_MAX_SIZE = 1024**3

    def manage_overflow(reassignments)
      if reassignments.has_next?
        write_overflow(reassignments)
      else<commit_msg>Let’s use megabytes and not gibibytes<commit_after># encoding: utf-8

module Ktl
  class Reassigner
    def initialize(type, zk_client, options={})
      @type = type
      @zk_client = zk_client
      @json_max_size = options[:json_max_size] || JSON_MAX_SIZE
    end

    def reassignment_in_progress?
      partitions = @zk_client.partitions_being_reassigned
      partitions.size > 0
    end

    def overflow?
      overflow_znodes = @zk_client.get_children(overflow_base_path)
      overflow_znodes.size > 0
    rescue ZkClient::Exception::ZkNoNodeException
      false
    end

    def load_overflow
      overflow = Scala::Collection::Map.empty
      overflow_nodes = @zk_client.get_children(overflow_base_path)
      overflow_nodes.foreach do |index|
        overflow_json = @zk_client.read_data(overflow_path(index)).first
        data = parse_reassignment_json(overflow_json)
        overflow = overflow.send('++', data)
      end
      overflow
    end

    def execute(reassignment)
      json = reassignment_json(reassignment)
      reassignments = split(reassignment, json.bytesize)
      json = reassignment_json(reassignments.next)
      @zk_client.reassign_partitions(json)
      manage_overflow(reassignments)
    end

    private

    JSON_MAX_SIZE = 1000**2

    def manage_overflow(reassignments)
      if reassignments.has_next?
        write_overflow(reassignments)
      else
{'fraction_non_alphanumeric': 0.07606344628695025, 'fraction_numerical': 0.005767844268204758, 'mean_word_length': 3.230182926829268, 'pattern_counts': {'":': 0, '<': 3, '<?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': '2058967', 'n_tokens_mistral': 878, 'n_tokens_neox': 851, 'n_words': 156}