red_amber 0.1.5 → 0.1.8

data/lib/red-amber.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require_relative 'red_amber'

data/lib/red_amber/data_frame.rb

@@ -6,18 +6,14 @@ module RedAmber
   class DataFrame
     # mix-in
     include DataFrameDisplayable
-    include DataFrameHelper
     include DataFrameIndexable
     include DataFrameSelectable
-    include DataFrameObservationOperation
     include DataFrameVariableOperation
+    include Helper
 
     def initialize(*args)
       @variables = @keys = @vectors = @types = @data_types = nil
-      # bug in gobject-introspection: ruby-gnome/ruby-gnome#1472
-      #  [Arrow::Table] == [nil] shows ArgumentError
-      #  temporary use yoda condition to workaround
-      if args.empty? || args == [[]] || args == [{}] || [nil] == args
+      if args.empty? || args[0] == [] || args[0] == {} || args[0].nil?
         # DataFrame.new, DataFrame.new([]), DataFrame.new({}), DataFrame.new(nil)
         #   returns empty DataFrame
         @table = Arrow::Table.new({}, [])
@@ -35,6 +31,7 @@ module RedAmber
             raise DataFrameTypeError, "invalid argument: #{arg}"
           end
       end
+      name_unnamed_keys
     end
 
     def self.load(path, options = {})
@@ -44,7 +41,7 @@ module RedAmber
     attr_reader :table
 
     def to_arrow
-      table
+      @table
     end
 
     def save(output, options = {})
@@ -101,10 +98,10 @@ module RedAmber
       @vectors || @vectors = init_instance_vars(:vectors)
     end
 
-    def indexes
-      0...size
+    def indices
+      (0...size).to_a
     end
-    alias_method :indices, :indexes
+    alias_method :indexes, :indices
 
     def to_h
       variables.transform_values(&:to_a)
@@ -130,9 +127,27 @@ module RedAmber
     end
 
     def to_rover
+      require 'rover'
       Rover::DataFrame.new(to_h)
     end
 
+    def to_iruby
+      require 'iruby'
+      return ['text/plain', '(empty DataFrame)'] if empty?
+
+      if ENV.fetch('RED_AMBER_OUTPUT_MODE', 'Table') == 'TDR'
+        size <= 5 ? ['text/plain', tdr_str(tally: 0)] : ['text/plain', tdr_str]
+      else
+        ['text/html', html_table]
+      end
+    end
+
+    def group(*group_keys, &block)
+      g = Group.new(self, group_keys)
+      g = g.summarize(&block) if block
+      g
+    end
+
     private
 
     # initialize @variable, @keys, @vectors and return one of them
@@ -148,5 +163,42 @@ module RedAmber
       @variables, @keys, @vectors = ary
       ary[%i[variables keys vectors].index(var)]
     end
+
+    def html_table
+      reduced = size > 8 ? self[0..4, -4..-1] : self
+
+      converted = reduced.assign do
+        vectors.select.with_object({}) do |vector, assigner|
+          if vector.has_nil?
+            assigner[vector.key] = vector.to_a.map do |e|
+              e = e.nil? ? '<i>(nil)</i>' : e.to_s # nil
+              e = '""' if e.empty? # empty string
+              e.sub(/(\s+)/, '"\1"') # blank spaces
+            end
+          end
+        end
+      end
+
+      html = IRuby::HTML.table(converted.to_h, maxrows: 8, maxcols: 15)
+      "#{self.class} <#{size} x #{n_keys} vector#{pl(n_keys)}> #{html}"
+    end
+
+    def name_unnamed_keys
+      return unless @table[:'']
+
+      # We can't use #keys because it causes mismatch of @table and @keys
+      keys = @table.schema.fields.map { |f| f.name.to_sym }
+      unnamed = (:unnamed1..).find { |e| !keys.include?(e) }
+      fields =
+        @table.schema.fields.map do |field|
+          if field.name.empty?
+            Arrow::Field.new(unnamed, field.data_type)
+          else
+            field
+          end
+        end
+      schema = Arrow::Schema.new(fields)
+      @table = Arrow::Table.new(schema, @table.columns)
+    end
   end
 end

data/lib/red_amber/data_frame_displayable.rb

@@ -5,8 +5,12 @@ require 'stringio'
 module RedAmber
   # mix-ins for the class DataFrame
   module DataFrameDisplayable
+    INDEX_KEY = :index_key_for_format_table
+
     def to_s
-      @table.to_s
+      return '' if empty?
+
+      format_table(width: 80)
     end
 
     # def describe() end
@@ -14,7 +18,11 @@ module RedAmber
     # def summary() end
 
     def inspect
-      "#<#{shape_str(with_id: true)}>\n#{dataframe_info(3)}"
+      if ENV.fetch('RED_AMBER_OUTPUT_MODE', 'Table') == 'TDR'
+        "#<#{shape_str(with_id: true)}>\n#{dataframe_info(3)}"
+      else
+        "#<#{shape_str(with_id: true)}>\n#{self}"
+      end
     end
 
     # - limit: max num of Vectors to show
@@ -30,10 +38,6 @@ module RedAmber
 
     private # =====
 
-    def pl(num)
-      num > 1 ? 's' : ''
-    end
-
     def shape_str(with_id: false)
       shape_info = empty? ? '(empty)' : "#{size} x #{n_keys} Vector#{pl(n_keys)}"
       id = with_id ? format(', 0x%016x', object_id) : ''
@@ -81,12 +85,12 @@ module RedAmber
     end
 
     def make_header_format(levels, headers, quoted_keys)
-      # find longest word to adjust column width
+      # find longest word to adjust width
       w_idx = n_keys.to_s.size
       w_key = [quoted_keys.map(&:size).max, headers[:key].size].max
       w_type = [types.map(&:size).max, headers[:type].size].max
-      w_row = [levels.map { |l| l.to_s.size }.max, headers[:levels].size].max
-      "%-#{w_idx}s %-#{w_key}s %-#{w_type}s %#{w_row}s %s\n"
+      w_level = [levels.map { |l| l.to_s.size }.max, headers[:levels].size].max
+      "%-#{w_idx}s %-#{w_key}s %-#{w_type}s %#{w_level}s %s\n"
     end
 
     def type_group(data_type)
@@ -128,5 +132,78 @@ module RedAmber
       a << "#{n_nil} nil#{pl(n_nil)}" unless n_nil.zero?
       a
     end
+
+    def format_table(width: 80)
+      head = 5
+      tail = 3
+      n_digit = 1
+
+      original = self
+      indices = size > head + tail ? [*0...head, *(size - tail)...size] : [*0...size]
+      df = slice(indices).assign do
+        assigner = { INDEX_KEY => indices.map { |i| (i + 1).to_s } }
+        vectors.each_with_object(assigner) do |v, a|
+          a[v.key] = v.to_a.map do |e|
+            if e.nil?
+              '(nil)'
+            elsif v.float?
+              e.round(n_digit).to_s
+            elsif v.string?
+              e
+            else
+              e.to_s
+            end
+          end
+        end
+      end
+
+      df = df.pick { [INDEX_KEY, keys - [INDEX_KEY]] }
+      df = size > head + tail ? df[0, 0, 0...head, 0, -tail..-1] : df[0, 0, 0..-1]
+      df = df.assign do
+        vectors.each_with_object({}) do |v, assigner|
+          vec = v.replace(0, v.key == INDEX_KEY ? '' : v.key.to_s)
+                 .replace(1, v.key == INDEX_KEY ? '' : "<#{original[v.key].type}>")
+          assigner[v.key] = size > head + tail ? vec.replace(head + 2, ':') : vec
+        end
+      end
+
+      width_list = df.vectors.map { |v| v.to_a.map(&:length).max }
+      total_length = width_list[-1] # reserved for last column
+
+      formats = []
+      row_ellipsis = nil
+      df.vectors.each_with_index do |v, i|
+        w = width_list[i]
+        if total_length + w > width && i < df.n_keys - 1
+          row_ellipsis = i
+          formats << '%3s'
+          formats << format_for_column(df.vectors[-1], original, width_list[-1])
+          break
+        end
+        formats << format_for_column(v, original, w)
+        total_length += w
+      end
+      format_str = formats.join(' ')
+
+      str = StringIO.new
+      if row_ellipsis
+        df = df[df.keys[0..row_ellipsis], df.keys[-1]]
+        df = df.assign(df.keys[row_ellipsis] => ['...'] * df.size)
+      end
+
+      df.to_a.each do |row|
+        str.puts format(format_str, *row).rstrip
+      end
+
+      str.string
+    end
+
+    def format_for_column(vector, original, width)
+      if vector.key != INDEX_KEY && !original[vector.key].numeric?
+        "%-#{width}s"
+      else
+        "%#{width}s"
+      end
+    end
   end
 end

data/lib/red_amber/data_frame_selectable.rb

@@ -3,35 +3,94 @@
 module RedAmber
   # mix-in for the class DataFrame
   module DataFrameSelectable
-    # select columns: [symbol] or [string]
-    # select rows: [array of index], [range]
+    # select variables: [symbol] or [string]
+    # select observations: [array of index], [range]
     def [](*args)
+      args.flatten!
       raise DataFrameArgumentError, 'Empty dataframe' if empty?
-      raise DataFrameArgumentError, 'Empty argument' if args.empty?
-
-      if args.one?
-        case args[0]
-        when Vector
-          return select_obs_by_boolean(Arrow::BooleanArray.new(args[0].data))
-        when Arrow::BooleanArray
-          return select_obs_by_boolean(args[0])
-        when Array
-          return select_obs_by_boolean(Arrow::BooleanArray.new(args[0]))
-
-          # when Hash
-          # specify conditions to select by a Hash
-        end
+      return remove_all_values if args.empty? || args[0].nil?
+
+      vector = parse_to_vector(args)
+      if vector.boolean?
+        return filter_by_vector(vector.data) if vector.size == size
+
+        raise DataFrameArgumentError, "Size is not match in booleans: #{args}"
+      end
+      return take_by_array(vector) if vector.numeric?
+      return select_vars_by_keys(vector.to_a.map(&:to_sym)) if vector.string? || vector.type == :dictionary
+
+      raise DataFrameArgumentError, "Invalid argument: #{args}"
+    end
+
+    # slice and select some observations to create sub DataFrame
+    def slice(*args, &block)
+      slicer = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+
+        slicer = instance_eval(&block)
+      end
+      slicer = [slicer].flatten
+
+      raise DataFrameArgumentError, 'Empty dataframe' if empty?
+      return remove_all_values if slicer.empty? || slicer[0].nil?
+
+      vector = parse_to_vector(slicer)
+      if vector.boolean?
+        return filter_by_vector(vector.data) if vector.size == size
+
+        raise DataFrameArgumentError, "Size is not match in booleans: #{slicer}"
       end
+      return take_by_array(vector) if vector.numeric?
+
+      raise DataFrameArgumentError, "Invalid argument #{slicer}"
+    end
+
+    # remove selected observations to create sub DataFrame
+    def remove(*args, &block)
+      remover = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+
+        remover = instance_eval(&block)
+      end
+      remover = [remover].flatten
+
+      raise DataFrameArgumentError, 'Empty dataframe' if empty?
+      return self if remover.empty? || remover[0].nil?
+
+      vector = parse_to_vector(remover)
+      if vector.boolean?
+        return filter_by_vector(vector.primitive_invert.data) if vector.size == size
+
+        raise DataFrameArgumentError, "Size is not match in booleans: #{remover}"
+      end
+      if vector.numeric?
+        raise DataFrameArgumentError, "Index out of range: #{vector.min}" if vector.min <= -size - 1
+
+        normalized_indices = (vector < 0).if_else(vector + size, vector) # normalize index from tail
+        if normalized_indices.max >= size
+          raise DataFrameArgumentError, "Index out of range: #{normalized_indices.max}"
+        end
 
-      return select_obs_by_boolean(args) if booleans?(args)
+        normalized_indices = normalized_indices.floor.to_a.map(&:to_i) # round to integer array
+        return remove_all_values if normalized_indices == indices
+        return self if normalized_indices.empty?
 
-      # expand Range like [1..3, 4] to [1, 2, 3, 4]
-      expanded = expand_range(args)
-      return map_indices(*expanded) if integers?(expanded)
-      return select_vars_by_keys(expanded.map(&:to_sym)) if sym_or_str?(expanded)
+        index_array = indices - normalized_indices
 
-      raise DataFrameArgumentError, "Invalid argument #{args}"
+        datum = Arrow::Function.find(:take).execute([table, index_array])
+        return DataFrame.new(datum.value)
+      end
+
+      raise DataFrameArgumentError, "Invalid argument #{remover}"
+    end
+
+    def remove_nil
+      func = Arrow::Function.find(:drop_null)
+      DataFrame.new(func.execute([table]).value)
     end
+    alias_method :drop_nil, :remove_nil
 
     # Select a variable by a key in String or Symbol
     def v(key)
@@ -43,24 +102,57 @@ module RedAmber
       variables[key.to_sym]
     end
 
-    def head(n_rows = 5)
-      raise DataFrameArgumentError, "Index is out of range #{n_rows}" if n_rows.negative?
+    def head(n_obs = 5)
+      raise DataFrameArgumentError, "Index is out of range #{n_obs}" if n_obs.negative?
 
-      self[0...[n_rows, size].min]
+      self[0...[n_obs, size].min]
     end
 
-    def tail(n_rows = 5)
-      raise DataFrameArgumentError, "Index is out of range #{n_rows}" if n_rows.negative?
+    def tail(n_obs = 5)
+      raise DataFrameArgumentError, "Index is out of range #{n_obs}" if n_obs.negative?
 
-      self[-[n_rows, size].min..]
+      self[-[n_obs, size].min..]
     end
 
-    def first(n_rows = 1)
-      head(n_rows)
+    def first(n_obs = 1)
+      head(n_obs)
     end
 
-    def last(n_rows = 1)
-      tail(n_rows)
+    def last(n_obs = 1)
+      tail(n_obs)
+    end
+
+    # Undocumented
+    # TODO: support for option {boundscheck: true}
+    def take(*indices)
+      indices.flatten!
+      return remove_all_values if indices.empty?
+
+      indices = indices[0] if indices.one? && !indices[0].is_a?(Numeric)
+      indices = Vector.new(indices) unless indices.is_a?(Vector)
+
+      take_by_array(indices)
+    end
+
+    # Undocumented
+    # TODO: support for option {null_selection_behavior: :drop}
+    def filter(*booleans)
+      booleans.flatten!
+      return remove_all_values if booleans.empty?
+
+      b = booleans[0]
+      case b
+      when Vector
+        raise DataFrameArgumentError, 'Argument is not a boolean.' unless b.boolean?
+
+        filter_by_vector(b.data)
+      when Arrow::BooleanArray
+        filter_by_vector(b)
+      else
+        raise DataFrameArgumentError, 'Argument is not a boolean.' unless booleans?(booleans)
+
+        filter_by_vector(Arrow::BooleanArray.new(booleans))
+      end
     end
 
     private
@@ -75,5 +167,32 @@ module RedAmber
         DataFrame.new(@table[keys])
       end
     end
+
+    # Accepts indices by numeric Vector
+    def take_by_array(indices)
+      raise DataFrameArgumentError, "Indices must be a numeric Vector: #{indices}" unless indices.numeric?
+      raise DataFrameArgumentError, "Index out of range: #{indices.min}" if indices.min <= -size - 1
+
+      normalized_indices = (indices < 0).if_else(indices + size, indices) # normalize index from tail
+      raise DataFrameArgumentError, "Index out of range: #{normalized_indices.max}" if normalized_indices.max >= size
+
+      index_array = Arrow::UInt64ArrayBuilder.build(normalized_indices.data) # round to integer array
+
+      datum = Arrow::Function.find(:take).execute([table, index_array])
+      DataFrame.new(datum.value)
+    end
+
+    # Accepts booleans by Arrow::BooleanArray
+    def filter_by_vector(boolean_array)
+      raise DataFrameArgumentError, 'Booleans must be same size as self.' unless boolean_array.length == size
+
+      datum = Arrow::Function.find(:filter).execute([table, boolean_array])
+      DataFrame.new(datum.value)
+    end
+
+    # return a DataFrame with same keys as self without values
+    def remove_all_values
+      filter_by_vector(Arrow::BooleanArray.new([false] * size))
+    end
   end
 end

data/lib/red_amber/data_frame_variable_operation.rb

@@ -129,5 +129,9 @@ module RedAmber
         arrays << Arrow::ChunkedArray.new([a])
       end
     end
+
+    def keys_by_booleans(booleans)
+      keys.select.with_index { |_, i| booleans[i] }
+    end
   end
 end

data/lib/red_amber/group.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # group class
+  class Group
+    def initialize(dataframe, *group_keys)
+      @dataframe = dataframe
+      @table = @dataframe.table
+      @group_keys = group_keys.flatten
+
+      raise GroupArgumentError, 'group_keys is empty.' if @group_keys.empty?
+
+      d = @group_keys - @dataframe.keys
+      raise GroupArgumentError, "#{d} is not a key of\n #{@dataframe}." unless d.empty?
+
+      @group = @table.group(*@group_keys)
+    end
+
+    functions = %i[count sum product mean min max stddev variance]
+    functions.each do |function|
+      define_method(function) do |*summary_keys|
+        by(function, summary_keys)
+      end
+    end
+
+    def inspect
+      tallys = @dataframe.pick(@group_keys).vectors.map.with_object({}) do |v, h|
+        h[v.key] = v.tally
+      end
+      "#<#{self.class}:#{format('0x%016x', object_id)}\n#{tallys}>"
+    end
+
+    def summarize(&block)
+      agg = instance_eval(&block)
+      case agg
+      when DataFrame
+        agg
+      when Array
+        agg.reduce { |aggregated, df| aggregated.assign(df.to_h) }
+      else
+        raise GroupArgumentError, "Unknown argument: #{agg}"
+      end
+    end
+
+    private
+
+    def by(func, summary_keys)
+      summary_keys = Array(summary_keys).flatten
+      d = summary_keys - @dataframe.keys
+      raise GroupArgumentError, "#{d} is not a key of\n #{@dataframe}." unless summary_keys.empty? || d.empty?
+
+      df = RedAmber::DataFrame.new(@group.send(func, *summary_keys))
+      df = df[df.keys[-1], df.keys[0...-1]]
+      # if counts are the same (no nil included), aggregate count columns.
+      df = df[df.keys[0..1]].rename(df.keys[1], :count) if func == :count && df.to_h.values[1..].uniq.size == 1
+      df
+    end
+  end
+end

data/lib/red_amber/helper.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-in for the class DataFrame
+  module Helper
+    private
+
+    def pl(num)
+      num > 1 ? 's' : ''
+    end
+
+    def out_of_range?(indeces)
+      indeces.max >= size || indeces.min < -size
+    end
+
+    def integers?(enum)
+      enum.all?(Integer)
+    end
+
+    def sym_or_str?(enum)
+      enum.all? { |e| e.is_a?(Symbol) || e.is_a?(String) }
+    end
+
+    def booleans?(enum)
+      enum.all? { |e| e.is_a?(TrueClass) || e.is_a?(FalseClass) || e.is_a?(NilClass) }
+    end
+
+    def create_dataframe_from_vector(key, vector)
+      DataFrame.new(key => vector.data)
+    end
+
+    def parse_to_vector(args)
+      a = args.reduce([]) do |accum, elem|
+        accum.concat(normalize_element(elem))
+      end
+      Vector.new(a)
+    end
+
+    def normalize_element(elem)
+      case elem
+      when Numeric, String, Symbol, TrueClass, FalseClass, NilClass
+        [elem]
+      when Range
+        both_end = [elem.begin, elem.end]
+        both_end[1] -= 1 if elem.exclude_end? && elem.end.is_a?(Integer)
+
+        if both_end.any?(Integer) || both_end.all?(&:nil?)
+          if both_end.any? { |e| e&.>=(size) || e&.<(-size) }
+            raise DataFrameArgumentError, "Index out of range: #{elem} for 0..#{size - 1}"
+          end
+
+          (0...size).to_a[elem]
+        else
+          elem.to_a
+        end
+      else
+        Array(elem)
+      end
+    end
+  end
+end