<commit_before>require 'rbbt/util/chain_methods'
require 'yaml'
module TSV
  extend ChainMethods
  self.chain_prefix = :tsv

  NIL_YAML = "--- \n"

  attr_accessor :unnamed, :serializer_module, :entity_options

  def entity_options
    options = namespace ? {:namespace => namespace, :organism => namespace} : {}
    if @entity_options
      options.merge(@entity_options)
    else
      options
    end
  end

  def with_unnamed
    saved_unnamed = @unnamed 
    @unnamed = true
    res = yield
    @unnamed = saved_unnamed
    res
  end

  def with_monitor(value = true)
    saved_monitor = @monitor
    @monitor = value
    res = yield
    @monitor = saved_monitor
    res
  end

  def self.extended(data)
    setup_chains(data)

    if not data.respond_to? :write
      class << data
        attr_accessor :writable

        def close
        end

        def read(force = false)
          @writable = false
          self
        end

        def write(force = false)
          @writable = true
          self
        end

        def write?
          @writable
        end
      end
    end

    if not data.respond_to? :serialized_get
      class << data
        alias serialized_get tsv_clean_get_brackets
        alias serialized_set tsv_clean_set_brackets
      end
    end
  end

  KEY_PREFIX = "__tsv_hash_"

  ENTRIES = []
  ENTRY_KEYS = []

  #{{{ Chained Methods
  def tsv_empty?
    length == 0
  end

  def tsv_get_brackets(key)
    value = serialized_get(key)
    return value if value.nil? or @unnamed or fields.nil?

    case type
    when :double, :list
      NamedArray.setup value, fields, key, entity_options
    when :flat, :single
      value = value.dup if value.frozen?

      value = Misc.prepare_entity(value, fields.first, entity_options)
    end
    value
  end

  def tsv_set_brackets(key,value)
    serialized_set(key, value)
  end

  def tsv_keys
    keys = tsv_clean_keys - ENTRY_KEYS
    return keys if @unnamed or key_field.nil?

    Misc.prepare_entity(keys, key_field, entity_options.merge(:dup_array => true))
  end

  def tsv_values
    values = values_at(*keys)
    return values if @unnamed or fields.nil?

    case type
    when :double, :list
      values.each{|value| NamedArray.setup value, fields, nil, entity_options}
    when :flat, :single
      values = values.collect{|v| Misc.prepare_entity(v, fields.first, entity_options)}
    end
      
    values
  end

  def tsv_each
    fields = self.fields

    serializer = self.serializer
    serializer_module = SERIALIZER_ALIAS[serializer] unless serializer.nil?
    tsv_clean_each do |key, value|
      next if ENTRY_KEYS.include? key

      # TODO Update this to be more efficient
      value = serializer_module.load(value) unless serializer.nil?

      # Annotated with Entity and NamedArray
      if not @unnamed
        if not fields.nil? 
          case type
          when :double, :list
            NamedArray.setup value, fields, key, entity_options if Array === value
          when :flat, :single
            Misc.prepare_entity(value, fields.first, entity_options)
          end
        end
        key = Misc.prepare_entity(key, key_field, entity_options)
      end

      yield key, value if block_given?
      [key, value]
    end
  end

  def tsv_collect
    serializer = self.serializer
    serializer_module = SERIALIZER_ALIAS[serializer] unless serializer.nil?
    tsv_clean_collect do |key, value|
      next if ENTRY_KEYS.include? key

      # TODO Update this to be more efficient
      value = serializer_module.load(value) unless serializer.nil?

      # Annotated with Entity and NamedArray
      if not @unnamed
        if not fields.nil? 
          case type
          when :double, :list
            NamedArray.setup value, fields, key, entity_options if Array === value 
          when :flat, :single
            value = Misc.prepare_entity(value, fields.first, entity_options)
          end
        end
        key = Misc.prepare_entity(key, key_field, entity_options)
      end


      if block_given?
        yield key, value
      else
        [key, value]
      end
    end
  end

  def tsv_size
    keys.length
  end

  def tsv_length
    keys.length
  end

  def tsv_values_at(*keys)
    keys.collect do |key|
      self[key]
    end
  end

  #{{{ Sorting

  def tsv_sort_by(field = nil, just_keys = false, &block)
    field = :all if field.nil?
    if field == :all
      elems = collect
    else
      elems = []
      case type
      when :single
        through :key, field do |key, field|
          elems << [key, field]
        end
      when :list, :flat
        through :key, field do |key, fields|
          elems << [key, fields.first]
        end
      when :double
        through :key, field do |key, fields|<commit_msg>Add chunked_values_at<commit_after>require 'rbbt/util/chain_methods'
require 'yaml'
module TSV
  extend ChainMethods
  self.chain_prefix = :tsv

  NIL_YAML = "--- \n"

  attr_accessor :unnamed, :serializer_module, :entity_options

  def entity_options
    options = namespace ? {:namespace => namespace, :organism => namespace} : {}
    if @entity_options
      options.merge(@entity_options)
    else
      options
    end
  end

  def with_unnamed
    saved_unnamed = @unnamed 
    @unnamed = true
    res = yield
    @unnamed = saved_unnamed
    res
  end

  def with_monitor(value = true)
    saved_monitor = @monitor
    @monitor = value
    res = yield
    @monitor = saved_monitor
    res
  end

  def self.extended(data)
    setup_chains(data)

    if not data.respond_to? :write
      class << data
        attr_accessor :writable

        def close
        end

        def read(force = false)
          @writable = false
          self
        end

        def write(force = false)
          @writable = true
          self
        end

        def write?
          @writable
        end
      end
    end

    if not data.respond_to? :serialized_get
      class << data
        alias serialized_get tsv_clean_get_brackets
        alias serialized_set tsv_clean_set_brackets
      end
    end
  end

  KEY_PREFIX = "__tsv_hash_"

  ENTRIES = []
  ENTRY_KEYS = []

  #{{{ Chained Methods
  def tsv_empty?
    length == 0
  end

  def tsv_get_brackets(key)
    value = serialized_get(key)
    return value if value.nil? or @unnamed or fields.nil?

    case type
    when :double, :list
      NamedArray.setup value, fields, key, entity_options
    when :flat, :single
      value = value.dup if value.frozen?

      value = Misc.prepare_entity(value, fields.first, entity_options)
    end
    value
  end

  def tsv_set_brackets(key,value)
    serialized_set(key, value)
  end

  def tsv_keys
    keys = tsv_clean_keys - ENTRY_KEYS
    return keys if @unnamed or key_field.nil?

    Misc.prepare_entity(keys, key_field, entity_options.merge(:dup_array => true))
  end

  def tsv_values
    values = values_at(*keys)
    return values if @unnamed or fields.nil?

    case type
    when :double, :list
      values.each{|value| NamedArray.setup value, fields, nil, entity_options}
    when :flat, :single
      values = values.collect{|v| Misc.prepare_entity(v, fields.first, entity_options)}
    end
      
    values
  end

  def tsv_each
    fields = self.fields

    serializer = self.serializer
    serializer_module = SERIALIZER_ALIAS[serializer] unless serializer.nil?
    tsv_clean_each do |key, value|
      next if ENTRY_KEYS.include? key

      # TODO Update this to be more efficient
      value = serializer_module.load(value) unless serializer.nil?

      # Annotated with Entity and NamedArray
      if not @unnamed
        if not fields.nil? 
          case type
          when :double, :list
            NamedArray.setup value, fields, key, entity_options if Array === value
          when :flat, :single
            Misc.prepare_entity(value, fields.first, entity_options)
          end
        end
        key = Misc.prepare_entity(key, key_field, entity_options)
      end

      yield key, value if block_given?
      [key, value]
    end
  end

  def tsv_collect
    serializer = self.serializer
    serializer_module = SERIALIZER_ALIAS[serializer] unless serializer.nil?
    tsv_clean_collect do |key, value|
      next if ENTRY_KEYS.include? key

      # TODO Update this to be more efficient
      value = serializer_module.load(value) unless serializer.nil?

      # Annotated with Entity and NamedArray
      if not @unnamed
        if not fields.nil? 
          case type
          when :double, :list
            NamedArray.setup value, fields, key, entity_options if Array === value 
          when :flat, :single
            value = Misc.prepare_entity(value, fields.first, entity_options)
          end
        end
        key = Misc.prepare_entity(key, key_field, entity_options)
      end


      if block_given?
        yield key, value
      else
        [key, value]
      end
    end
  end

  def tsv_size
    keys.length
  end

  def tsv_length
    keys.length
  end

  def tsv_values_at(*keys)
    keys.collect do |key|
      self[key]
    end
  end

  def chunked_values_at(keys, max = 5000)
    Misc.divide(keys, (keys.length.to_f / max).ceil).inject(nil) do |acc,c|
      new = self.values_at(*c)
      if acc.nil?
        acc = new
      else
        acc.concat(new)
      end
    end
  end



  #{{{ Sorting

  def tsv_sort_by(field = nil, just_keys = false, &block)
    field = :all if field.nil?
    if field == :all
      elems = collect
    else
      elems = []
      case type
      when :single
        through :key, field do |key, field|
          elems << [key, field]
        end
      when :list, :flat
        through :key, field do |key, fields|
          elems << [key, fields.first]
        end
      when :double
        through :key, field do |key, fields|
{'fraction_non_alphanumeric': 0.07728432873915263, 'fraction_numerical': 0.0006125574272588055, 'mean_word_length': 2.2034009156311316, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6403980', 'n_tokens_mistral': 3179, 'n_tokens_neox': 3048, 'n_words': 870}