DataFrame#set_index can take column name array, which results in multi-index

lib/daru/dataframe.rb

@@ -752,7 +752,7 @@ def rolling_fillna(direction=:forward)
  # 3 4 d
  #
  def uniq(*vtrs)
- vecs = vtrs.empty? ? vectors.map(&:to_s) : Array(vtrs)
+ vecs = vtrs.empty? ? vectors.to_a : Array(vtrs)
  grouped = group_by(vecs)
  indexes = grouped.groups.values.map { |v| v[0] }.sort
  row[*indexes]
@@ -1527,14 +1527,52 @@ def union other_df
  df
  end
 
+ module SetSingleIndexStrategy
+ def self.uniq_size(df, col)
+ df[col].uniq.size
+ end
+
+ def self.new_index(df, col)
+ Daru::Index.new(df[col].to_a)
+ end
+
+ def self.delete_vector(df, col)
+ df.delete_vector(col)
+ end
+ end
+
+ module SetMultiIndexStrategy
+ def self.uniq_size(df, cols)
+ df[*cols].uniq.size
+ end
+
+ def self.new_index(df, cols)
+ Daru::MultiIndex.from_arrays(df[*cols].map_vectors(&:to_a)).tap do |mi|
+ mi.name = cols
+ mi
+ end
+ end
+
+ def self.delete_vector(df, cols)
+ df.delete_vectors(*cols)
+ end
+ end
+
  # Set a particular column as the new DF
- def set_index new_index, opts={}
- raise ArgumentError, 'All elements in new index must be unique.' if
- @size != self[new_index].uniq.size
+ def set_index new_index_col, opts={}
+ if new_index_col.respond_to?(:to_a)
+ strategy = SetMultiIndexStrategy
+ new_index_col = new_index_col.to_a
+ else
+ strategy = SetSingleIndexStrategy
+ end
 
- self.index = Daru::Index.new(self[new_index].to_a)
- delete_vector(new_index) unless opts[:keep]
+ uniq_size = strategy.uniq_size(self, new_index_col)
+ raise ArgumentError, 'All elements in new index must be unique.' if
+ @size != uniq_size
 
+ self.index = strategy.new_index(self, new_index_col)
+ strategy.delete_vector(self, new_index_col) unless opts[:keep]
  self
  end
 

spec/dataframe_spec.rb

@@ -3712,6 +3712,26 @@ def create_test(*args, &_proc)
  jholu.set_index(:a)
  }.to raise_error(ArgumentError)
  end
+
+ it "sets multiindex if array is given" do
+ df = Daru::DataFrame.new({
+ a: %w[a a b b],
+ b: [1, 2, 1, 2],
+ c: %w[a b c d]
+ })
+ df.set_index(%i[a b])
+ expected =
+ Daru::DataFrame.new(
+ { c: %w[a b c d] },
+ index: Daru::MultiIndex.from_tuples(
+ [['a', 1], ['a', 2], ['b', 1], ['b', 2]]
+ )
+ ).tap do |df|
+ df.index.name = %i[a b]
+ df
+ end
+ expect(df).to eq(expected)
+ end
  end
 
  context "#concat" do