Skip to content

Commit

Permalink
Multi Index should return a Daru::Index when a single level is left
Browse files Browse the repository at this point in the history
  • Loading branch information
jonspalmer committed Aug 22, 2020
1 parent 3ac3d15 commit 3a0c571
Show file tree
Hide file tree
Showing 5 changed files with 42 additions and 19 deletions.
7 changes: 6 additions & 1 deletion lib/daru/index/multi_index.rb
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,12 @@ def dup
end

def drop_left_level by=1
MultiIndex.from_arrays to_a.transpose[by..-1]
arrays = to_a.transpose[by..-1]
if arrays.length == 1
Index.new(arrays[0])
else
MultiIndex.from_arrays to_a.transpose[by..-1]
end
end

def | other
Expand Down
6 changes: 2 additions & 4 deletions spec/category_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -689,15 +689,13 @@
end

it "returns sub vector when passed first and second layer of tuple" do
mi = Daru::MultiIndex.from_tuples([
[:foo],
[:bar]])
mi = Daru::Index.new([:foo,:bar])
expect(@vector[:c,:two]).to eq(Daru::Vector.new([10,11], index: mi,
name: :sub_sub_vector, type: :category))
end

it "returns sub vector not a single element when passed the partial tuple" do
mi = Daru::MultiIndex.from_tuples([[:foo]])
mi = Daru::Index.new([:foo])
expect(@vector[:d, :one]).to eq(Daru::Vector.new([12], index: mi,
name: :sub_sub_vector, type: :category))
end
Expand Down
14 changes: 9 additions & 5 deletions spec/dataframe_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -507,13 +507,20 @@
sub_order = Daru::MultiIndex.from_tuples([
[:one, :bar],
[:two, :baz]
])
])
expect(@df_mi[:a]).to eq(Daru::DataFrame.new([
@vector_arry1,
@vector_arry2
], index: @multi_index, order: sub_order))
end

it "returns DataFrame with simple index when specified first and second layer of MultiIndex" do
sub_order = Daru::Index.new([:baz])
expect(@df_mi[:a][:two]).to eq(Daru::DataFrame.new([
@vector_arry2
], index: @multi_index, order: sub_order))
end

it "returns a Vector if the last level of MultiIndex is tracked" do
expect(@df_mi[:a, :one, :bar]).to eq(
Daru::Vector.new(@vector_arry1, index: @multi_index))
Expand Down Expand Up @@ -1525,10 +1532,7 @@
end

it "returns DataFrame when specifying first and second layer of MultiIndex" do
sub_index = Daru::MultiIndex.from_tuples([
[:bar],
[:baz]
])
sub_index = Daru::Index.new([:bar,:baz])
expect(@df_mi.row[:c,:one]).to eq(Daru::DataFrame.new([
[11,12],
[1,2],
Expand Down
24 changes: 21 additions & 3 deletions spec/index/multi_index_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -453,6 +453,16 @@
])
)
end

it "returns a Daru::Index if a single level remains" do
expect(
Daru::MultiIndex.from_tuples([
[:c,:one,:bar],
[:c,:one,:baz]
]).drop_left_level(2)).to eq(
Daru::Index.new([:bar, :baz])
)
end
end

context 'other forms of tuple list representation' do
Expand Down Expand Up @@ -540,12 +550,20 @@
])
end

context "multiple indexes" do
subject { idx.subset :b, :one }
context "first level" do
subject { idx.subset :b }

it { is_expected.to be_a described_class }
its(:size) { is_expected.to eq 4 }
its(:to_a) { is_expected.to eq [[:one, :bar], [:two, :bar], [:two, :baz], [:one, :foo]] }
end

context "first and second level" do
subject { idx.subset :b, :one }

it { is_expected.to be_a Daru::Index }
its(:size) { is_expected.to eq 2 }
its(:to_a) { is_expected.to eq [[:bar], [:foo]] }
its(:to_a) { is_expected.to eq [:bar, :foo] }
end

context "multiple positional indexes" do
Expand Down
10 changes: 4 additions & 6 deletions spec/vector_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -239,16 +239,14 @@
dtype: dtype, name: :sub_vector))
end

it "returns sub vector when passed first and second layer of tuple" do
mi = Daru::MultiIndex.from_tuples([
[:foo],
[:bar]])
expect(@vector[:c,:two]).to eq(Daru::Vector.new([10,11], index: mi,
it "returns sub vector with simple index when passed first and second layer of tuple" do
i = Daru::Index.new([:foo, :bar])
expect(@vector[:c,:two]).to eq(Daru::Vector.new([10,11], index: i,
dtype: dtype, name: :sub_sub_vector))
end

it "returns sub vector not a single element when passed the partial tuple" do
mi = Daru::MultiIndex.from_tuples([[:foo]])
mi = Daru::Index.new([:foo])
expect(@vector[:d, :one]).to eq(Daru::Vector.new([12], index: mi,
dtype: dtype, name: :sub_sub_vector))
end
Expand Down

0 comments on commit 3a0c571

Please sign in to comment.