rdb/lib/rdb/cmd/relational.rb

# frozen_string_literal: true

# Selecting, combining and summarising tables.

Rdb.register('row', 'Make a new table where rows match logical condition',
             "row 'column condition column-or-value ..' < table") do |args|
  die "usage: row 'condition' < table" unless args.size == 1
  t = table_in
  prog = "BEGIN{OFS=\"\\t\"} #{to_awk(args[0], t.head)} {print}"
  t.with(rows: awk(prog, t.rows)).write
end
Rdb.alias_for('select', 'row')

Rdb.register('column', 'Display columns of a table in any order', 'column [ Column ... ] < tableorlist') do |args|
  t = table_in
  next t.write if args.empty?

  idx = args.map { |a| Rdb.column_index(t.head, a) }
  head = args.each_with_index.map { |a, n| idx[n] ? t.head[idx[n]] : a }
  t.with(head: head, rows: t.rows.map { |r| idx.map { |i| i ? r[i] : '' } }).write
end
Rdb.alias_for('project', 'column')

Rdb.register('sorttable', 'Sort a table by one or more columns',
             'sorttable [-nrf] [columnname ...] < tableorlist') do |args|
  t = table_in
  numeric = reverse = fold = false
  cols = []
  args.each do |a|
    if a.match?(/\A-[nrf]+\z/)
      numeric ||= a.include?('n')
      reverse ||= a.include?('r')
      fold ||= a.include?('f')
    else
      cols << t.index_of(a)
    end
  end
  cols = (0...t.head.size).to_a if cols.empty?
  key = ->(r) { cols.map { |i| numeric ? r[i].to_f : (fold ? r[i].downcase : r[i]) } }
  sorted = t.rows.each_with_index.sort_by { |r, n| [key.(r), n] }.map(&:first)
  sorted.reverse! if reverse
  t.with(rows: sorted).write
end

Rdb.register('compute', 'Calculate columns of a table', "compute 'column = expression [ ; ... ]' < tableorlist") do |args|
  die "usage: compute 'column = expression [; ...]' < table" unless args.size == 1
  t = table_in
  prog = "BEGIN{OFS=\"\\t\"} { #{to_awk(args[0], t.head)}; print }"
  t.with(rows: awk(prog, t.rows)).write
end

Rdb.register('addcol', 'Add a column to a table', 'addcol newcolumn < table') do |args|
  die 'usage: addcol newcolumn ... < table' if args.empty?
  t = table_in
  t.with(head: t.head + args, rows: t.rows.map { |r| r + [''] * args.size }).write
end

Rdb.register('rename', 'Rename a column', 'rename [ oldcolumn newcolumn ... ] < tableorlist') do |args|
  die 'usage: rename oldcolumn newcolumn ... < table' if args.empty? || args.size.odd?
  t = table_in
  head = t.head.dup
  args.each_slice(2) { |old, new| head[t.index_of(old)] = new }
  t.with(head: head).write
end

{ 'number' => 'number', 'Number' => 'Number' }.each do |cmd, title|
  Rdb.register(cmd, 'Insert a column-row number into a table or list', "#{cmd} < tableorlist") do |_args|
    t = table_in
    t.with(head: [title] + t.head, rows: t.rows.each_with_index.map { |r, i| [(i + 1).to_s] + r }).write
  end
end

Rdb.register('record', 'Find and output a record from a table', 'record number < table') do |args|
  die 'usage: record number < table' unless args.size == 1
  t = table_in
  row = t.rows[args[0].to_i - 1]
  t.with(rows: row ? [row] : []).write
end

# ---- combining tables ------------------------------------------------------

Rdb.register('jointable', 'Join two tables into one where keys match',
             'jointable [-a[1|2]] [-j[1|2] column]... [-n] [-] table1 [table2]') do |args|
  all1 = all2 = false
  k1 = k2 = nil
  files = []
  until args.empty?
    a = args.shift
    case a
    when '-a'  then all1 = all2 = true
    when '-a1' then all1 = true
    when '-a2' then all2 = true
    when '-j'  then k1 = k2 = args.shift
    when '-j1' then k1 = args.shift
    when '-j2' then k2 = args.shift
    when '-n'  then nil
    else files << a
    end
  end
  die 'usage: jointable [-a[1|2]] [-j[1|2] column] table1 [table2]' if files.empty? || files.size > 2

  files << '-' if files.size == 1
  t1, t2 = files.map { |f| open_table(f) }
  i1 = k1 ? t1.index_of(k1) : 0
  i2 = k2 ? t2.index_of(k2) : 0
  keep2 = (0...t2.head.size).to_a - [i2]
  by_key = Hash.new { |h, k| h[k] = [] }
  t2.rows.each { |r| by_key[r[i2]] << r }
  used = {}
  out = []
  t1.rows.each do |r|
    if (m = by_key[r[i1]]).any?
      used[r[i1]] = true
      m.each { |r2| out << r + keep2.map { |i| r2[i] } }
    elsif all1
      out << r + [''] * keep2.size
    end
  end
  if all2
    t2.rows.each do |r2|
      next if used[r2[i2]]

      blank = [''] * t1.head.size
      blank[i1] = r2[i2]
      out << blank + keep2.map { |i| r2[i] }
    end
  end
  Rdb::Table.new(t1.head + keep2.map { |i| t2.head[i] }, out).write
end

Rdb.register('union', 'Concatenate tables', 'union tableorlist ... [-] [< tableorlist]') do |args|
  args = ['-'] if args.empty?
  tabs = args.map { |a| open_table(a) }
  tabs.first.with(rows: tabs.flat_map(&:rows)).write
end

Rdb.register('intersect', 'Write table of rows that are in both input tables', 'intersect table1 table2') do |args|
  die 'usage: intersect table1 table2' unless args.size == 2
  a, b = args.map { |x| open_table(x) }
  seen = b.rows.to_h { |r| [r, true] }
  a.with(rows: a.rows.select { |r| seen[r] }.uniq).write
end

Rdb.register('difference', 'Output table of rows that are in only one table', 'difference table1 table2') do |args|
  die 'usage: difference table1 table2' unless args.size == 2
  a, b = args.map { |x| open_table(x) }
  sb = b.rows.to_h { |r| [r, true] }
  a.with(rows: a.rows.reject { |r| sb[r] }.uniq).write
end

Rdb.register('paste.rdb', 'Paste together two or more tables', 'paste.rdb table1 table2 ...') do |args|
  die 'usage: paste.rdb table1 table2 ...' if args.size < 2
  tabs = args.map { |a| open_table(a) }
  n = tabs.map { |t| t.rows.size }.max
  head = tabs.flat_map(&:head)
  rows = (0...n).map { |i| tabs.flat_map { |t| t.rows[i] || [''] * t.head.size } }
  Rdb::Table.new(head, rows).write
end

Rdb.register('splittable', 'Divide a table horizontally', 'splittable [-n] table') do |args|
  n = 50
  n = args.shift[1..].to_i if args.first&.match?(/\A-\d+\z/)
  die 'usage: splittable [-n] table' unless args.size == 1
  t = open_table(args[0])
  suffix = ('aa'..'zz').each
  t.rows.each_slice(n) do |chunk|
    File.write("#{args[0]}#{suffix.next}", t.with(rows: chunk).to_s)
  end
end

# ---- aggregates ------------------------------------------------------------
# total, mean, maximum, minimum, width, precision and datatype are one program
# that decides what to compute from the name it was called by.

STATS = {
  'total' => 'Sum a column',
  'mean' => 'Display the mean of a column',
  'maximum' => 'Display the maximum value in a column',
  'minimum' => 'Display the minimum of each column selected',
  'width' => 'Display the width of each column',
  'precision' => 'Display the precision of a column',
  'datatype' => 'Display the data type of each column selected'
}.freeze

STATS.each do |cmd, desc|
  Rdb.register(cmd, desc, "#{cmd} [ -l ] [ column ... ] < table") do |args|
    list = args.first == '-l'
    args.shift if list
    t = table_in
    sel = args.empty? ? (0...t.head.size).to_a : args.map { |a| t.index_of(a) }
    result = Array.new(t.head.size, '')
    sel.each do |i|
      vals = t.rows.map { |r| r[i] }
      filled = vals.reject { |v| v.strip.empty? }
      nums = filled.all? { |v| numeric?(v) } ? vals.map { |v| v.strip.to_f } : nil
      dec = filled.map { |v| decimals(v) }.max.to_i
      result[i] =
        case name
        when 'width'     then (vals.map(&:size) + [t.head[i].size]).max.to_s
        when 'precision' then nums ? dec.to_s : ''
        when 'datatype'  then filled.empty? || !nums ? '-1' : dec.to_s
        when 'total'     then nums ? fmt(nums.sum, dec) : ''
        when 'mean'      then nums && !nums.empty? ? fmt(nums.sum / nums.size, [dec, 2].max) : ''
        when 'maximum'   then nums && !nums.empty? ? fmt(nums.max, dec) : ''
        when 'minimum'   then nums && !nums.empty? ? fmt(nums.min, dec) : ''
        end
    end
    if list
      $stdout.write t.to_s
      puts Rdb::Table.dash_row(t.head)
      puts result.join("\t")
    else
      Rdb::Table.new(t.head, [result]).write
    end
  end
end

Rdb.register('widest', 'Output the width of the widest entries in a table', 'widest < table') do |_args|
  t = table_in
  res = t.head.each_index.map { |i| (t.rows.map { |r| r[i].size } + [t.head[i].size]).max.to_s }
  Rdb::Table.new(t.head, [res]).write
end

Rdb.register('subtotal', 'Output subtotals of columns in a table',
             'subtotal [ -l ] [ break-column column ... ] < table') do |args|
  list = args.first == '-l'
  args.shift if list
  t = table_in
  brk = args.empty? ? 0 : t.index_of(args.shift)
  sel = args.empty? ? (0...t.head.size).to_a - [brk] : args.map { |a| t.index_of(a) }
  out = []
  group = []
  flush = lambda do
    next if group.empty?

    sub = Array.new(t.head.size, '')
    sub[brk] = group.first[brk]
    sel.each do |i|
      vals = group.map { |r| r[i] }.reject { |v| v.strip.empty? }
      next if vals.empty? || !vals.all? { |v| numeric?(v) }

      sub[i] = fmt(vals.sum { |v| v.strip.to_f }, vals.map { |v| decimals(v) }.max)
    end
    if list
      out.concat(group)
      out << t.head.map { |h| '-' * [h.size, 1].max }
    end
    out << sub
    group = []
  end
  t.rows.each do |r|
    flush.() if group.any? && group.first[brk] != r[brk]
    group << r
  end
  flush.()
  t.with(rows: out).write
end

# ---- describing tables ------------------------------------------------------

Rdb.register('schema', "Print a table's schema", 'schema tableorlist ...') do |args|
  die 'usage: schema tableorlist ...' if args.empty?
  rows = args.flat_map do |f|
    open_table(f).head.each_with_index.map { |h, i| [File.basename(f), (i + 1).to_s, h] }
  end
  Rdb::Table.new(%w[Table Field Name], rows).write
end

Rdb.register('dbdict', 'Print a data-base dictionary', 'dbdict < table') do |_args|
  t = table_in
  Rdb::Table.new(%w[field name], t.head.each_with_index.map { |h, i| [(i + 1).to_s, h] }).write
end

Rdb.register('tableorlist', 'Report whether a file has table or list format',
             'tableorlist [tableorlist ...] [< tableorlist]') do |args|
  texts = args.empty? ? [$stdin.read] : args.map { |a| File.read(a) }
  status = 0
  texts.each do |tx|
    if tx.start_with?("\n")
      puts 'list'
      status = 1
    else
      puts 'table'
    end
  end
  exit status
end

Rdb.register('check.rdb', 'Report any rows in which columns do not match head line', 'check.rdb < tableorlist') do |_args|
  lines = $stdin.read.lines.map { |l| l.chomp("\n") }
  expected = lines.first.to_s.count("\t") + 1
  bad = 0
  lines.each_with_index do |l, n|
    next if n.zero? || l.count("\t") + 1 == expected

    bad += 1
    warn "check.rdb: Columns (#{l.count("\t") + 1}) do not equal headline columns (#{expected}) in row #{n - 1} (line #{n + 1})."
    warn l
  end
  exit(bad.zero? ? 0 : 1)
end

# ---- head and dash lines -------------------------------------------------

Rdb.register('headon', 'Add an /rdb header to a table', 'headon [column ...] < table') do |args|
  rows = $stdin.each_line.map { |l| l.chomp("\n").split("\t", -1) }
  n = rows.map(&:size).max || 0
  head = args.dup
  head << (head.size + 1).to_s while head.size < n
  Rdb::Table.new(head, rows.map { |r| r.fill('', r.size...head.size) }).write
end

Rdb.register('headoff', 'Remove an /rdb head from both table and list file', 'headoff < tableorlist') do |_args|
  text = $stdin.read
  lines = text.lines
  $stdout.write lines.drop(text.start_with?("\n") ? 1 : 2).join
end

Rdb.register('insertdash', 'Insert dash line as second line in table', 'insertdash < table-without-dashline') do |_args|
  lines = $stdin.read.lines
  head = lines.shift.to_s.chomp("\n")
  puts head
  puts head.gsub(/[^\t]/, '-')
  $stdout.write lines.join
end

Rdb.register('display', 'Write table or list file to standard output', 'display < tableorlist') do |_args|
  $stdout.write $stdin.read
end

Rdb.register('see', 'Display nonprinting as well as printing characters', 'see < tableorlist') do |_args|
  $stdin.read.each_char do |c|
    o = c.ord
    $stdout.write(
      if c == "\n" then "$\n"
      elsif o < 32 then "^#{(o + 64).chr}"
      elsif o == 127 then '^?'
      else c
      end
    )
  end
end