red_amber 0.1.2 → 0.1.5

data/lib/red_amber/data_frame_displayable.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+require 'stringio'
+
+module RedAmber
+  # mix-ins for the class DataFrame
+  module DataFrameDisplayable
+    def to_s
+      @table.to_s
+    end
+
+    # def describe() end
+
+    # def summary() end
+
+    def inspect
+      "#<#{shape_str(with_id: true)}>\n#{dataframe_info(3)}"
+    end
+
+    # - 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 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) : ''
+      "#{self.class} : #{shape_info}#{id}"
+    end
+
+    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)
+      type_groups = @table.columns.map { |column| type_group(column.data_type) }
+      quoted_keys = keys.map(&:inspect)
+      headers = { idx: '#', key: 'key', type: 'type', levels: 'level', data: 'data_preview' }
+      header_format = make_header_format(levels, headers, quoted_keys)
+
+      sio = StringIO.new # output string buffer
+      sio.puts "Vector#{pl(n_keys)} : #{var_type_count(type_groups).join(', ')}"
+      sio.printf header_format, *headers.values
+
+      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 != size
+                [data_tally.to_s]
+              else
+                [shorthand(vector, size, max_element)].concat na_string(vector)
+              end
+            else
+              [shorthand(vector, size, max_element)]
+            end
+        sio.printf header_format, i + 1, key, type, data_tally.size, a.join(', ')
+      end
+      sio.string
+    end
+
+    def make_header_format(levels, headers, quoted_keys)
+      # find longest word to adjust column 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"
+    end
+
+    def type_group(data_type)
+      case data_type
+      when Arrow::NumericDataType then :numeric
+      when Arrow::StringDataType then :string
+      when Arrow::BooleanDataType then :boolean
+      when Arrow::TemporalDataType then :temporal
+      else
+        :other
+      end
+    end
+
+    def var_type_count(type_groups)
+      tg = type_groups.tally
+      a = []
+      a << "#{tg[:numeric]} numeric" if tg[:numeric]
+      a << "#{tg[:string]} string#{pl(tg[:string])}" if tg[:string]
+      a << "#{tg[:boolean]} boolean" if tg[:boolean]
+      a << "#{tg[:temporal]} temporal" if tg[:temporal]
+      a
+    end
+
+    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 size > max
+      "[#{a.join(', ')}]"
+    end
+
+    def na_string(vector)
+      n_nan = vector.n_nans
+      n_nil = vector.n_nils
+      a = []
+      return a if (n_nan + n_nil).zero?
+
+      a << "#{n_nan} NaN#{pl(n_nan)}" unless n_nan.zero?
+      a << "#{n_nil} nil#{pl(n_nil)}" unless n_nil.zero?
+      a
+    end
+  end
+end

data/lib/red_amber/data_frame_helper.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-in for the class DataFrame
+  module DataFrameHelper
+    private
+
+    def expand_range(args)
+      ary = args.each_with_object([]) do |e, a|
+        e.is_a?(Range) ? a.concat(normalized_array(e)) : a.append(e)
+      end
+      ary.map do |e|
+        if e.is_a?(Integer) && e.negative?
+          e + size
+        else
+          e
+        end
+      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)
+
+      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
+
+        (0...size).to_a[range]
+      else
+        range.to_a
+      end
+    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 select_obs_by_boolean(array)
+      DataFrame.new(@table.filter(array))
+    end
+
+    def keys_by_booleans(booleans)
+      keys.select.with_index { |_, i| booleans[i] }
+    end
+  end
+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,83 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-ins for the class DataFrame
+  module DataFrameObservationOperation
+    # 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
+      return remove_all_values if slicer.empty? || slicer[0].nil?
+
+      # filter with same length
+      booleans = nil
+      if slicer[0].is_a?(Vector) || slicer[0].is_a?(Arrow::BooleanArray)
+        booleans = slicer[0].to_a
+      elsif slicer.size == size && booleans?(slicer)
+        booleans = slicer
+      end
+      return select_obs_by_boolean(booleans) if booleans
+
+      # filter with indexes
+      slicer = expand_range(slicer)
+      return map_indices(*slicer) if integers?(slicer)
+
+      raise DataFrameArgumentError, "Invalid argument #{args}"
+    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
+
+      return self if remover.empty?
+
+      # filter with same length
+      booleans = nil
+      if remover[0].is_a?(Vector) || remover[0].is_a?(Arrow::BooleanArray)
+        booleans = remover[0].to_a
+      elsif remover.size == size && booleans?(remover)
+        booleans = remover
+      end
+      if booleans
+        inverted = booleans.map(&:!)
+        return select_obs_by_boolean(inverted)
+      end
+
+      # filter with indexes
+      slicer = indexes.to_a - expand_range(remover)
+      return remove_all_values if slicer.empty?
+      return map_indices(*slicer) if integers?(slicer)
+
+      raise DataFrameArgumentError, "Invalid argument #{args}"
+    end
+
+    def remove_nil
+      func = Arrow::Function.find(:drop_null)
+      DataFrame.new(func.execute([table]).value)
+    end
+    alias_method :drop_nil, :remove_nil
+
+    def group(aggregating_keys, func, target_keys)
+      t = table.group(*aggregating_keys)
+      RedAmber::DataFrame.new(t.send(func, *target_keys))
+    end
+
+    private
+
+    # return a DataFrame with same keys as self without values
+    def remove_all_values
+      DataFrame.new(keys.each_with_object({}) { |key, h| h[key] = [] })
+    end
+  end
+end

data/lib/red_amber/data_frame_selectable.rb

@@ -9,18 +9,40 @@ module RedAmber
       raise DataFrameArgumentError, 'Empty dataframe' if empty?
       raise DataFrameArgumentError, 'Empty argument' if args.empty?
 
-      # 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)
+      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
+      end
 
-      return select_rows(expanded) if integers?(expanded)
-      return select_columns(expanded.map(&:to_sym)) if sym_or_str?(expanded)
+      return select_obs_by_boolean(args) if booleans?(args)
+
+      # 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)
 
       raise DataFrameArgumentError, "Invalid argument #{args}"
     end
 
+    # 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?
 
@@ -41,48 +63,17 @@ module RedAmber
       tail(n_rows)
     end
 
-    private # =====
+    private
 
-    def select_columns(keys)
+    def select_vars_by_keys(keys)
       if keys.one?
-        Vector.new(@table[*keys].data)
-      else
-        DataFrame.new(@table[keys])
-      end
-    end
-
-    def select_rows(indeces)
-      out_of_range?(indeces) && raise(DataFrameArgumentError, "Invalid index: #{indeces} for 0..#{size - 1}")
+        key = keys[0].to_sym
+        raise DataFrameArgumentError, "Key does not exist #{keys}" unless key? key
 
-      a = indeces.map { |i| @table.slice(i).to_a }
-      DataFrame.new(@table.schema, a)
-    end
-
-    def normalized_array(range)
-      both_end = [range.begin, range.end]
-      both_end[1] -= 1 if range.exclude_end? && range.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: #{range} for 0..#{size - 1}"
-        end
-
-        (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
-    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
   end
 end

data/lib/red_amber/data_frame_variable_operation.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-ins for the class DataFrame
+  module DataFrameVariableOperation
+    # pick up some variables to create sub DataFrame
+    def pick(*args, &block)
+      picker = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+
+        picker = instance_eval(&block)
+      end
+      picker = [picker].flatten
+      return DataFrame.new if picker.empty? || picker == [nil]
+
+      picker = keys_by_booleans(picker) if booleans?(picker)
+
+      # DataFrame#[] creates a Vector with single key is specified.
+      # DataFrame#pick creates a DataFrame with single key.
+      return DataFrame.new(@table[picker]) if sym_or_str?(picker)
+
+      raise DataFrameArgumentError, "Invalid argument #{args}"
+    end
+
+    # drop some variables to create remainer sub DataFrame
+    def drop(*args, &block)
+      dropper = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+
+        dropper = instance_eval(&block)
+      end
+      dropper = [dropper].flatten
+      dropper = keys_by_booleans(dropper) if booleans?(dropper)
+
+      picker = keys - dropper
+      return DataFrame.new if picker.empty?
+
+      # DataFrame#[] creates a Vector with single key is specified.
+      # DataFrame#drop creates a DataFrame with single key.
+      return DataFrame.new(@table[picker]) if sym_or_str?(picker)
+
+      raise DataFrameArgumentError, "Invalid argument #{args}"
+    end
+
+    # rename variables to create new DataFrame
+    def rename(*args, &block)
+      renamer = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and a block' unless args.empty?
+
+        renamer = instance_eval(&block)
+      end
+      renamer = [renamer].flatten
+      return self if renamer.empty?
+
+      return rename_by_hash([renamer].to_h) if renamer.size == 2 && sym_or_str?(renamer) # rename(from, to)
+      return rename_by_hash(renamer[0]) if renamer.one? && renamer[0].is_a?(Hash) # rename({from => to})
+
+      raise DataFrameArgumentError, "Invalid argument #{args}"
+    end
+
+    # assign variables to create new DataFrame
+    def assign(*args, &block)
+      assigner = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and a block' unless args.empty?
+
+        assigner = instance_eval(&block)
+      end
+      assigner = [assigner].flatten
+      return self if assigner.empty? || assigner == [nil]
+
+      raise DataFrameArgumentError, "Invalid argument #{args}" unless assigner.one? && assigner[0].is_a?(Hash)
+
+      updater = {}
+      appender = {}
+      assigner[0].each do |key, value|
+        if keys.include? key
+          updater[key] = value
+        else
+          appender[key] = value
+        end
+      end
+      fields, arrays = update_fields_and_arrays(updater)
+      append_to_fields_and_arrays(appender, fields, arrays) unless appender.empty?
+
+      DataFrame.new(Arrow::Table.new(Arrow::Schema.new(fields), arrays))
+    end
+
+    private
+
+    def rename_by_hash(key_pairs)
+      fields = keys.map do |key|
+        new_key = key_pairs[key]
+        if new_key
+          Arrow::Field.new(new_key.to_sym, @table[key].data_type)
+        else
+          @table.schema[key]
+        end
+      end
+      schema = Arrow::Schema.new(fields)
+      DataFrame.new(Arrow::Table.new(schema, @table.columns))
+    end
+
+    def update_fields_and_arrays(updater)
+      fields = @table.columns.map(&:field)
+      arrays = @table.columns.map(&:data) # chunked_arrays
+      keys.each_with_index do |key, i|
+        data = updater[key]
+        next unless data
+
+        raise DataFrameArgumentError, "Data size mismatch (#{data.size} != #{size})" if data.size != size
+
+        a = Arrow::Array.new(data.is_a?(Vector) ? data.to_a : data)
+        fields[i] = Arrow::Field.new(key, a.value_data_type)
+        arrays[i] = Arrow::ChunkedArray.new([a])
+      end
+      [fields, arrays]
+    end
+
+    def append_to_fields_and_arrays(appender, fields, arrays)
+      appender.each do |key, data|
+        raise DataFrameArgumentError, "Data size mismatch (#{data.size} != #{size})" if data.size != size
+
+        a = Arrow::Array.new(data.is_a?(Vector) ? data.to_a : data)
+        fields << Arrow::Field.new(key.to_sym, a.value_data_type)
+        arrays << Arrow::ChunkedArray.new([a])
+      end
+    end
+  end
+end

data/lib/red_amber/vector.rb

@@ -5,10 +5,12 @@ module RedAmber
   #   @data : holds Arrow::ChunkedArray
   class Vector
     # mix-in
+    include VectorCompensable
     include VectorFunctions
 
     # chunked_array may come from column.data
     def initialize(array)
+      @key = nil # default is 'headless'
       case array
       when Vector
         @data = array.data
@@ -17,18 +19,31 @@ module RedAmber
       when Array
         @data = Arrow::Array.new(array)
       else
-        raise ArgumentError, 'Unknown array in argument'
+        raise VectorArgumentError, 'Unknown array in argument'
       end
     end
 
     attr_reader :data
+    attr_accessor :key
 
     def to_s
       @data.to_a.inspect
     end
 
-    def inspect
-      format "#<#{self.class}(:#{type}, size=#{size}):0x%016x>\n#{self}", object_id
+    def inspect(limit: 80)
+      sio = StringIO.new << '['
+      to_a.each_with_object(sio).with_index do |(e, s), i|
+        next_str = "#{s.size > 1 ? ', ' : ''}#{e.inspect}"
+        if (s.size + next_str.size) < limit
+          s << next_str
+        else
+          s << ', ... ' if i < size
+          break
+        end
+      end
+      sio << ']'
+
+      format "#<#{self.class}(:#{type}, size=#{size}):0x%016x>\n%s\n", object_id, sio.string
     end
 
     def values
@@ -49,8 +64,24 @@ module RedAmber
       @data.value_type.nick.to_sym
     end
 
-    def data_type
-      @data.value_type
+    def boolean?
+      type == :boolean
+    end
+
+    def numeric?
+      type_class < Arrow::NumericDataType
+    end
+
+    def string?
+      type == :string
+    end
+
+    def temporal?
+      type_class < Arrow::TemporalDataType
+    end
+
+    def type_class
+      @data.value_data_type.class
     end
 
     # def each() end
@@ -66,11 +97,32 @@ module RedAmber
     # def each_chunk() end
 
     def tally
-      values.tally
+      hash = values.tally
+      if (type_class < Arrow::FloatingPointDataType) && is_nan.any
+        a = 0
+        hash.each do |key, value|
+          if key.is_a?(Float) && key.nan?
+            hash.delete(key)
+            a += value
+          end
+        end
+        hash[Float::NAN] = a
+      end
+      hash
+    end
+
+    def value_counts
+      values, counts = Arrow::Function.find(:value_counts).execute([data]).value.fields
+      values.zip(counts).to_h
     end
 
     def n_nulls
       @data.n_nulls
     end
+    alias_method :n_nils, :n_nulls
+
+    def n_nans
+      numeric? ? is_nan.to_a.count(true) : 0
+    end
   end
 end