red_amber 0.1.3 → 0.1.6

data/lib/red_amber/data_frame.rb

@@ -5,19 +5,23 @@ module RedAmber
   #   @table   : holds Arrow::Table object
   class DataFrame
     # mix-in
+    include DataFrameDisplayable
+    include DataFrameIndexable
     include DataFrameSelectable
-    include DataFrameOutput
+    include DataFrameObservationOperation
+    include DataFrameVariableOperation
+    include Helper
 
     def initialize(*args)
-      # DataFrame.new, DataFrame.new([]), DataFrame.new({}), DataFrame.new(nil)
-      #   returns empty DataFrame
-      @table = Arrow::Table.new({}, [])
+      @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
-      return if args.empty? || args == [[]] || args == [{}] || [nil] == args
-
-      if args.size > 1
+      if args.empty? || args == [[]] || args == [{}] || [nil] == args
+        # DataFrame.new, DataFrame.new([]), DataFrame.new({}), DataFrame.new(nil)
+        #   returns empty DataFrame
+        @table = Arrow::Table.new({}, [])
+      elsif args.size > 1
         @table = Arrow::Table.new(*args)
       else
         arg = args[0]
@@ -39,67 +43,71 @@ module RedAmber
 
     attr_reader :table
 
+    def to_arrow
+      @table
+    end
+
     def save(output, options = {})
       @table.save(output, options)
     end
 
-    # Properties ===
-    def n_rows
+    def size
       @table.n_rows
     end
-    alias_method :nrow, :n_rows
-    alias_method :size, :n_rows
-    alias_method :length, :n_rows
+    alias_method :n_rows, :size
+    alias_method :n_obs, :size
 
-    def n_columns
+    def n_keys
       @table.n_columns
     end
-    alias_method :ncol, :n_columns
-    alias_method :width, :n_columns
+    alias_method :n_cols, :n_keys
+    alias_method :n_vars, :n_keys
 
     def shape
-      [n_rows, n_columns]
+      [size, n_keys]
     end
 
-    def column_names
-      @table.columns.map { |column| column.name.to_sym }
+    def variables
+      @variables || @variables = init_instance_vars(:variables)
     end
-    alias_method :keys, :column_names
-    alias_method :header, :column_names
+    alias_method :vars, :variables
+
+    def keys
+      @keys || @keys = init_instance_vars(:keys)
+    end
+    alias_method :column_names, :keys
+    alias_method :var_names, :keys
 
     def key?(key)
-      column_names.include?(key.to_sym)
+      @keys.include?(key.to_sym)
     end
     alias_method :has_key?, :key?
 
     def key_index(key)
-      column_names.find_index(key.to_sym)
+      @keys.find_index(key.to_sym)
     end
     alias_method :find_index, :key_index
     alias_method :index, :key_index
 
     def types
-      @table.columns.map do |column|
-        column.data_type.to_s.to_sym
-      end
+      @types || @types = @table.columns.map { |column| column.data.value_type.nick.to_sym }
     end
 
-    def data_types
-      @table.columns.map do |column|
-        column.data_type.class
-      end
+    def type_classes
+      @data_types || @data_types = @table.columns.map { |column| column.data_type.class }
     end
 
     def vectors
-      @table.columns.map do |column|
-        Vector.new(column.data)
-      end
+      @vectors || @vectors = init_instance_vars(:vectors)
+    end
+
+    def indices
+      (0...size).to_a
     end
+    alias_method :indexes, :indices
 
     def to_h
-      @table.columns.each_with_object({}) do |column, result|
-        result[column.name.to_sym] = column.entries
-      end
+      variables.transform_values(&:to_a)
     end
 
     def to_a
@@ -118,13 +126,59 @@ module RedAmber
     end
 
     def empty?
-      @table.columns.empty?
+      variables.empty?
     end
 
     def to_rover
       Rover::DataFrame.new(to_h)
     end
 
-    # def to_parquet() end
+    def to_iruby
+      require 'iruby'
+      return ['text/plain', '(empty DataFrame)'] if empty?
+
+      if ENV.fetch('RED_AMBER_OUTPUT_MODE', 'tdr') == 'table'
+        ['text/html', html_table]
+      elsif size <= 5
+        ['text/plain', tdr_str(tally: 0)]
+      else
+        ['text/plain', tdr_str]
+      end
+    end
+
+    private
+
+    # initialize @variable, @keys, @vectors and return one of them
+    def init_instance_vars(var)
+      ary = @table.columns.each_with_object([{}, [], []]) do |column, (variables, keys, vectors)|
+        v = Vector.new(column.data)
+        k = column.name.to_sym
+        v.key = k
+        variables[k] = v
+        keys << k
+        vectors << v
+      end
+      @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)
+      "#{size} x #{n_keys} vector#{pl(n_keys)} ; #{html}"
+    end
   end
 end

data/lib/red_amber/{data_frame_output.rb → data_frame_displayable.rb}

@@ -4,7 +4,7 @@ require 'stringio'
 
 module RedAmber
   # mix-ins for the class DataFrame
-  module DataFrameOutput
+  module DataFrameDisplayable
     def to_s
       @table.to_s
     end
@@ -13,19 +13,37 @@ module RedAmber
 
     # def summary() end
 
-    def inspect_raw
-      format "#<#{self.class}:0x%016x>\n#{self}", object_id
+    def inspect
+      if ENV.fetch('RED_AMBER_OUTPUT_MODE', 'tdr') == 'table'
+        "#<#{shape_str(with_id: true)}>\n#{self}"
+      else
+        "#<#{shape_str(with_id: true)}>\n#{dataframe_info(3)}"
+      end
     end
 
-    # - tally_level: max level to use tally mode
-    # - max_element: max element to show values in each row
-    # - TODO: Is it better to change name other than `inspect` ?
-    # - TODO: Fall back to inspect_raw when treating large dataset
-    # - TODO: Refactor code to smaller methods
-    def inspect(tally_level: 5, max_element: 5)
-      return '#<RedAmber::DataFrame (empty)>' if empty?
+    # - limit: max num of Vectors to show
+    # - tally: max level to use tally mode
+    # - elements: max element to show values in each vector
+    def tdr(limit = 10, tally: 5, elements: 5)
+      puts tdr_str(limit, tally: tally, elements: elements)
+    end
+
+    def tdr_str(limit = 10, tally: 5, elements: 5)
+      "#{shape_str}\n#{dataframe_info(limit, tally_level: tally, max_element: elements)}"
+    end
+
+    private # =====
+
+    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) : ''
+      "#{self.class} : #{shape_info}#{id}"
+    end
 
-      stringio = StringIO.new # output string buffer
+    def dataframe_info(limit, tally_level: 5, max_element: 5)
+      return '' if empty?
+
+      limit = n_keys if [:all, -1].include? limit
 
       tallys = vectors.map(&:tally)
       levels = tallys.map(&:size)
@@ -34,52 +52,41 @@ module RedAmber
       headers = { idx: '#', key: 'key', type: 'type', levels: 'level', data: 'data_preview' }
       header_format = make_header_format(levels, headers, quoted_keys)
 
-      # 1st row: show shape of the dataframe
-      vs = "Vector#{pl(ncol)}"
-      stringio.puts \
-        "#{self.class} : #{nrow} x #{ncol} #{vs}"
-
-      # 2nd row: show var counts by type
-      stringio.puts "#{vs} : #{var_type_count(type_groups).join(', ')}"
+      sio = StringIO.new # output string buffer
+      sio.puts "Vector#{pl(n_keys)} : #{var_type_count(type_groups).join(', ')}"
+      sio.printf header_format, *headers.values
 
-      # 3rd row: print header of rows
-      stringio.printf header_format, *headers.values
-
-      # 4th row ~: show details for each column (vector)
       vectors.each.with_index do |vector, i|
+        if i >= limit
+          sio << " ... #{n_keys - i} more Vector#{pl(n_keys - i)} ...\n"
+          break
+        end
         key = quoted_keys[i]
         type = types[i]
         type_group = type_groups[i]
         data_tally = tallys[i]
-
         a = case type_group
             when :numeric, :string, :boolean
-              if data_tally.size <= tally_level && data_tally.size != nrow
+              if data_tally.size <= tally_level && data_tally.size != size
                 [data_tally.to_s]
               else
-                [shorthand(vector, nrow, max_element)].concat na_string(vector)
+                [shorthand(vector, size, max_element)].concat na_string(vector)
               end
             else
-              shorthand(vector, nrow, max_element)
+              [shorthand(vector, size, max_element)]
             end
-        stringio.printf header_format, i + 1, key, type, data_tally.size, a.join(', ')
+        sio.printf header_format, i + 1, key, type, data_tally.size, a.join(', ')
       end
-      stringio.string
-    end
-
-    private # =====
-
-    def pl(num)
-      num > 1 ? 's' : ''
+      sio.string
     end
 
     def make_header_format(levels, headers, quoted_keys)
-      # find longest word to adjust column width
-      w_idx = ncol.to_s.size
+      # 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)
@@ -103,10 +110,11 @@ module RedAmber
       a
     end
 
-    def shorthand(vector, nrow, max_element)
-      a = vector.to_a.take(max_element)
+    def shorthand(vector, size, max_element)
+      max = vector.temporal? ? 2 : max_element
+      a = vector.to_a.take(max)
       a.map! { |e| e.nil? ? 'nil' : e.inspect }
-      a << '... ' if nrow > max_element
+      a << '... ' if size > max
       "[#{a.join(', ')}]"
     end
 

data/lib/red_amber/data_frame_indexable.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-ins for the class DataFrame
+  module DataFrameIndexable
+    # Common method
+    def map_indices(*indices)
+      return self if indices.empty?
+
+      indices = indices[0].data if indices[0].is_a?(Vector)
+
+      new_dataframe_by(indices)
+    end
+
+    # @param sort_keys [Arrow::SortKey]
+    #   :key, "key" or "+key" denotes ascending,
+    #   "-key" denotes descending order
+    # @return [RedAmber::Vector] Sorted indices in Vector
+    def sort_indices(*sort_keys)
+      indices = @table.sort_indices(sort_keys.flatten)
+      Vector.new(indices)
+    end
+
+    # @return [RedAmber::DataFrame] Sorted DataFrame
+    def sort(*sort_keys)
+      indices = @table.sort_indices(sort_keys.flatten)
+
+      new_dataframe_by(indices)
+    end
+
+    private
+
+    def new_dataframe_by(index_array)
+      t = Arrow::Function.find(:take).execute([@table, index_array]).value
+      RedAmber::DataFrame.new(t)
+    end
+  end
+end

data/lib/red_amber/data_frame_observation_operation.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-ins for the class DataFrame
+  module DataFrameObservationOperation
+    def group(aggregating_keys, func, target_keys)
+      t = table.group(*aggregating_keys)
+      RedAmber::DataFrame.new(t.send(func, *target_keys))
+    end
+  end
+end

data/lib/red_amber/data_frame_selectable.rb

@@ -3,89 +3,196 @@
 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?
+      return remove_all_values if args.empty? || args[0].nil?
 
-      # expand Range like [1..3, 4] to [1, 2, 3, 4]
-      expanded =
-        args.each_with_object([]) do |e, a|
-          e.is_a?(Range) ? a.concat(normalized_array(e)) : a.append(e)
+      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_rows(expanded) if integers?(expanded)
-      return select_columns(expanded.map(&:to_sym)) if sym_or_str?(expanded)
+        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?
+
+        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)
+      unless key.is_a?(Symbol) || key.is_a?(String)
+        raise DataFrameArgumentError, "Key is not a Symbol or String [#{key}]"
+      end
+      raise DataFrameArgumentError, "Key not exist [#{key}]" unless key?(key)
+
+      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
 
-    private # =====
+    # Undocumented
+    # TODO: support for option {boundscheck: true}
+    def take(*indices)
+      indices.flatten!
+      return remove_all_values if indices.empty?
 
-    def select_columns(keys)
-      if keys.one?
-        t = @table[*keys]
-        raise DataFrameArgumentError, "Key is not exists #{keys}" unless t
+      indices = indices[0] if indices.one? && !indices[0].is_a?(Numeric)
+      indices = Vector.new(indices) unless indices.is_a?(Vector)
 
-        Vector.new(t.data)
-      else
-        DataFrame.new(@table[keys])
-      end
+      take_by_array(indices)
     end
 
-    def select_rows(indeces)
-      out_of_range?(indeces) && raise(DataFrameArgumentError, "Invalid index: #{indeces} for 0..#{size - 1}")
+    # Undocumented
+    # TODO: support for option {null_selection_behavior: :drop}
+    def filter(*booleans)
+      booleans.flatten!
+      return remove_all_values if booleans.empty?
 
-      a = indeces.map { |i| @table.slice(i).to_a }
-      DataFrame.new(@table.schema, a)
+      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
 
-    def normalized_array(range)
-      both_end = [range.begin, range.end]
-      both_end[1] -= 1 if range.exclude_end? && range.end.is_a?(Integer)
+    private
 
-      if both_end.any?(Integer) || both_end.all?(&:nil?)
-        if both_end.any? { |e| e&.>=(size) || e&.<(-size) }
-          raise DataFrameArgumentError, "Index out of range: #{range} for 0..#{size - 1}"
-        end
+    def select_vars_by_keys(keys)
+      if keys.one?
+        key = keys[0].to_sym
+        raise DataFrameArgumentError, "Key does not exist #{keys}" unless key? key
 
-        (0...size).to_a[range]
+        variables[key]
       else
-        range.to_a
+        DataFrame.new(@table[keys])
       end
     end
 
-    def out_of_range?(indeces)
-      indeces.max >= size || indeces.min < -size
+    # 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
 
-    def integers?(enum)
-      enum.all?(Integer)
+    # 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
 
-    def sym_or_str?(enum)
-      enum.all? { |e| e.is_a?(Symbol) || e.is_a?(String) }
+    # 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