<commit_before>require 'csv'

module Cassava
  class Document
    attr_accessor :name, :columns, :column_offset, :offset_column
    attr_accessor :rows

    def initialize opts = nil
      @opts = opts
      @name = opts[:name] or raise ArgumentError, "name not specified"
      @columns = opts[:columns]
      @column_offset = { }
      @offset_column = [ ]
      @ncols = 0
      @rows = [ ]
      @index = { }
    end

    def nrows
      @rows.size
    end

    def empty_rows!
      @rows = [ ]
      @index.clear
      @column_types = nil
      self
    end

    def append_rows! rows
      return self unless rows
      if Document === rows
        rows.columns.each { | c | add_column!(c) }
        rows = rows.rows
      end
      rows.map!{ | r | array_to_row(r) }
      row_i = @rows.size
      @rows.concat(rows)
      rows.each do | r |
        @ncols = r.size if @ncols < r.size
        r[:_row_i] = row_i
        row_i += 1
      end
      @index.clear
      @column_types = nil
      self
    end

    def parse!
      # debugger if $DEBUG
      csv = nil
      if RUBY_VERSION =~ /^1\.8/
        @rows = [ ]
        csv = CSV.open(name, "rb", @opts[:col_sep]) do | r |
          @rows << r
        end
      else
        csv_opts = { }
        csv_opts[:col_sep] = @opts[:col_sep] if @opts[:col_sep]
        csv = CSV.open(name, "rb", csv_opts)
        @rows = csv.read
      end

      @columns = @rows.shift unless @columns
      update_columns!

      row_i = 0
      @rows.map! do | r |
        @ncols = r.size if @ncols < r.size
        h = { :_row_i => (row_i += 1) }
        @column_offset.each do | c, i |
          h[c] = r[i] if c && i
        end
        h
      end
      # debugger
      self
    ensure
      csv.close if csv
    end

    def add_column! c
      c = c.to_sym
      unless i = @column_offset[c]
        @columns << c
        update_columns!
      end
      i
    end

    def update_columns!
      @ncols = @columns.size
      i = -1
      @columns.map! do | c |
        i += 1
        c = c.to_s
        next if c.empty?
        c = c.to_sym
        @column_offset[c] = i
        @offset_column[i] = c
        c
      end
      @column_types = nil
      self
    end

    def index! c
      c = @column[c] if Integer === c
      unless ind = @index[c]
        ind = { }
        @rows.each do | r |
          (ind[r[c]] ||= [ ]) << r
        end
        # pp [ :index!, c, ind.size ]
        @index[c] = ind
      end
      ind
    end

    def get c, v
      index!(c)[v]
    end

    def emit! file
      out = nil
      case file
      when '/dev/stdout', '/dev/fd/0'
        tmp_file = true
        out = $stdout
      when %r{^/dev/}
        tmp_file = true
      end
      if tmp_file
        tmp_file = "/tmp/cassava-#{$$}.csv"
        $stderr.puts "  #{file} => #{tmp_file}"
        file = tmp_file
      end

      _emit! file

      if tmp_file
        out ||= File.open(tmp, "w")
        $stderr.puts "  #{tmp_file} => #{out}"
        out.write(File.read(tmp_file))
      end
    ensure
      File.unlink(tmp_file) if tmp_file
    end

    def _emit! file
      CSV.open(file, "wb") do | out |
        a = @offset_column.map do | c |
          c && c.to_s
        end
        out << a
        @rows.each do | r |
          a = @offset_column.map do | i |
            i && r[i]
          end
          out << a
        end
      end
      self
    end

    def array_to_row a, columns = nil
      if Array === a
        columns ||= self.columns<commit_msg>Does not really require a tmp file.<commit_after>require 'csv'

module Cassava
  class Document
    attr_accessor :name, :columns, :column_offset, :offset_column
    attr_accessor :rows

    def initialize opts = nil
      @opts = opts
      @name = opts[:name] or raise ArgumentError, "name not specified"
      @columns = opts[:columns]
      @column_offset = { }
      @offset_column = [ ]
      @ncols = 0
      @rows = [ ]
      @index = { }
    end

    def nrows
      @rows.size
    end

    def empty_rows!
      @rows = [ ]
      @index.clear
      @column_types = nil
      self
    end

    def append_rows! rows
      return self unless rows
      if Document === rows
        rows.columns.each { | c | add_column!(c) }
        rows = rows.rows
      end
      rows.map!{ | r | array_to_row(r) }
      row_i = @rows.size
      @rows.concat(rows)
      rows.each do | r |
        @ncols = r.size if @ncols < r.size
        r[:_row_i] = row_i
        row_i += 1
      end
      @index.clear
      @column_types = nil
      self
    end

    def parse!
      # debugger if $DEBUG
      csv = nil
      if RUBY_VERSION =~ /^1\.8/
        @rows = [ ]
        csv = CSV.open(name, "rb", @opts[:col_sep]) do | r |
          @rows << r
        end
      else
        csv_opts = { }
        csv_opts[:col_sep] = @opts[:col_sep] if @opts[:col_sep]
        csv = CSV.open(name, "rb", csv_opts)
        @rows = csv.read
      end

      @columns = @rows.shift unless @columns
      update_columns!

      row_i = 0
      @rows.map! do | r |
        @ncols = r.size if @ncols < r.size
        h = { :_row_i => (row_i += 1) }
        @column_offset.each do | c, i |
          h[c] = r[i] if c && i
        end
        h
      end
      # debugger
      self
    ensure
      csv.close if csv
    end

    def add_column! c
      c = c.to_sym
      unless i = @column_offset[c]
        @columns << c
        update_columns!
      end
      i
    end

    def update_columns!
      @ncols = @columns.size
      i = -1
      @columns.map! do | c |
        i += 1
        c = c.to_s
        next if c.empty?
        c = c.to_sym
        @column_offset[c] = i
        @offset_column[i] = c
        c
      end
      @column_types = nil
      self
    end

    def index! c
      c = @column[c] if Integer === c
      unless ind = @index[c]
        ind = { }
        @rows.each do | r |
          (ind[r[c]] ||= [ ]) << r
        end
        # pp [ :index!, c, ind.size ]
        @index[c] = ind
      end
      ind
    end

    def get c, v
      index!(c)[v]
    end

    def emit! file
      out = nil
      case file
      when '/dev/stdout', '/dev/fd/0'
        tmp_file = true
        out = $stdout
      when %r{^/dev/}
        tmp_file = true
      end
      if tmp_file
        tmp_file = "/tmp/cassava-#{$$}.csv"
        # $stderr.puts "  #{file} => #{tmp_file}"
        file = tmp_file
      end

      _emit! file

      if tmp_file
        out ||= File.open(tmp, "w")
        # $stderr.puts "  #{tmp_file} => #{out}"
        out.write(File.read(tmp_file))
      end
    ensure
      File.unlink(tmp_file) if tmp_file
    end

    def _emit! file
      CSV.open(file, "wb") do | out |
        a = @offset_column.map do | c |
          c && c.to_s
        end
        out << a
        @rows.each do | r |
          a = @offset_column.map do | i |
            i && r[i]
          end
          out << a
        end
      end
      self
    end
    alias :emit! :_emit!

    def array_to_row a, columns = nil
      if Array === a
        columns ||= self.columns
{'fraction_non_alphanumeric': 0.11940087607743394, 'fraction_numerical': 0.002543450614667232, 'mean_word_length': 1.3680160588825694, 'pattern_counts': {'":': 0, '<': 27, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1271608', 'n_tokens_mistral': 2503, 'n_tokens_neox': 2436, 'n_words': 652}