Added access_row_tuples_by_indexs method (#463)

Prakriti-nith · zverok · commit e64a1d5eaf4b · 2018-07-07T16:45:35.000+03:00
diff --git a/lib/daru/dataframe.rb b/lib/daru/dataframe.rb
@@ -2274,6 +2274,50 @@ def split_by_category cat_name
       end
     end
 
+    # @param indexes [Array] index(s) at which row tuples are retrieved
+    # @return [Array] returns array of row tuples at given index(s)
+    # @example Using Daru::Index
+    #   df = Daru::DataFrame.new({
+    #     a: [1, 2, 3],
+    #     b: ['a', 'a', 'b']
+    #   })
+    #
+    #   df.access_row_tuples_by_indexs(1,2)
+    #   # => [[2, "a"], [3, "b"]]
+    #
+    #   df.index = Daru::Index.new([:one,:two,:three])
+    #   df.access_row_tuples_by_indexs(:one,:three)
+    #   # => [[1, "a"], [3, "b"]]
+    #
+    # @example Using Daru::MultiIndex
+    #   mi_idx = Daru::MultiIndex.from_tuples [
+    #     [:a,:one,:bar],
+    #     [:a,:one,:baz],
+    #     [:b,:two,:bar],
+    #     [:a,:two,:baz],
+    #   ]
+    #   df_mi = Daru::DataFrame.new({
+    #     a: 1..4,
+    #     b: 'a'..'d'
+    #   }, index: mi_idx )
+    #
+    #   df_mi.access_row_tuples_by_indexs(:b, :two, :bar)
+    #   # => [[3, "c"]]
+    #   df_mi.access_row_tuples_by_indexs(:a)
+    #   # => [[1, "a"], [2, "b"], [4, "d"]]
+    def access_row_tuples_by_indexs *indexes
+      return get_sub_dataframe(indexes, by_position: false).map_rows(&:to_a) if
+      @index.is_a?(Daru::MultiIndex)
+      positions = @index.pos(*indexes)
+      if positions.is_a? Numeric
+        row = populate_row_for(positions)
+        row.first.is_a?(Array) ? row : [row]
+      else
+        new_rows = @data.map { |vec| vec[*indexes] }
+        indexes.map { |index| new_rows.map { |r| r[index] } }
+      end
+    end
+
     # Function to use for aggregating the data.
     #
     # @param options [Hash] options for column, you want in resultant dataframe
diff --git a/spec/dataframe_spec.rb b/spec/dataframe_spec.rb
@@ -4037,6 +4037,56 @@ def create_test(*args, &_proc)
     end
   end
 
+  context '#access_row_tuples_by_indexs' do
+    let(:df) {
+      Daru::DataFrame.new({col: [:a, :b, :c, :d, :e], num: [52,12,07,17,01]}) }
+    let(:df_idx) {
+      Daru::DataFrame.new({a: [52, 12, 07], b: [1, 2, 3]}, index: [:one, :two, :three])
+    }
+    let (:mi_idx) do
+      Daru::MultiIndex.from_tuples [
+        [:a,:one,:bar],
+        [:a,:one,:baz],
+        [:b,:two,:bar],
+        [:a,:two,:baz],
+      ]
+    end
+    let (:df_mi) do
+      Daru::DataFrame.new({
+        a: 1..4,
+        b: 'a'..'d'
+      }, index: mi_idx )
+    end
+    context 'when no index is given' do
+      it 'returns empty Array' do
+        expect(df.access_row_tuples_by_indexs()).to eq([])
+      end
+    end
+    context 'when index(s) are given' do
+      it 'returns Array of row tuples' do
+        expect(df.access_row_tuples_by_indexs(1)).to eq([[:b, 12]])
+        expect(df.access_row_tuples_by_indexs(0,3)).to eq([[:a, 52], [:d, 17]])
+      end
+    end
+    context 'when custom index(s) are given' do
+      it 'returns Array of row tuples' do
+        expect(df_idx.access_row_tuples_by_indexs(:one,:three)).to eq(
+          [[52, 1], [7, 3]]
+        )
+      end
+    end
+    context 'when multi index is given' do
+      it 'returns Array of row tuples' do
+        expect(df_mi.access_row_tuples_by_indexs(:a)).to eq(
+          [[1, "a"], [2, "b"], [4, "d"]]
+        )
+        expect(df_mi.access_row_tuples_by_indexs(:a, :one, :baz)).to eq(
+          [[2, "b"]]
+        )
+      end
+    end
+  end
+
   context '#aggregate' do
     let(:cat_idx) { Daru::CategoricalIndex.new [:a, :b, :a, :a, :c] }
     let(:df) { Daru::DataFrame.new(num: [52,12,07,17,01], cat_index: cat_idx) }
