Skip to content

Commit 29c18da

Browse files
charlesdong1991ueshin
authored andcommitted
ENH: Add from_arrays for MultiIndex (#1001)
I was trying to create a MultiIndex from an array just now, and found out it hadn't been implemented, so have this PR to create MI from arrays. ```python >>> arrays = [['a', 'a', 'b', 'b'], ['red', 'blue', 'red', 'blue']] >>> ks.MultiIndex.from_arrays(arrays) MultiIndex([('a', 'red'), ('a', 'blue'), ('b', 'red'), ('b', 'blue')], names=['number', 'color']) ```
1 parent 23dc764 commit 29c18da

File tree

2 files changed

+41
-0
lines changed

2 files changed

+41
-0
lines changed

databricks/koalas/indexes.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -569,6 +569,40 @@ def from_tuples(tuples, sortorder=None, names=None):
569569
return DataFrame(index=pd.MultiIndex.from_tuples(
570570
tuples=tuples, sortorder=sortorder, names=names)).index
571571

572+
@staticmethod
573+
def from_arrays(arrays, sortorder=None, names=None):
574+
"""
575+
Convert arrays to MultiIndex.
576+
577+
Parameters
578+
----------
579+
arrays: list / sequence of array-likes
580+
Each array-like gives one level’s value for each data point. len(arrays)
581+
is the number of levels.
582+
sortorder: int or None
583+
Level of sortedness (must be lexicographically sorted by that level).
584+
names: list / sequence of str, optional
585+
Names for the levels in the index.
586+
587+
Returns
588+
-------
589+
index: MultiIndex
590+
591+
Examples
592+
--------
593+
594+
>>> arrays = [[1, 1, 2, 2], ['red', 'blue', 'red', 'blue']]
595+
>>> ks.MultiIndex.from_arrays(arrays, names=('number', 'color')) # doctest: +SKIP
596+
MultiIndex([(1, 'red'),
597+
(1, 'blue'),
598+
(2, 'red'),
599+
(2, 'blue')],
600+
names=['number', 'color'])
601+
"""
602+
return DataFrame(index=pd.MultiIndex.from_arrays(
603+
arrays=arrays, sortorder=sortorder, names=names
604+
)).index
605+
572606
@property
573607
def name(self) -> str:
574608
raise PandasNotImplementedError(class_name='pd.MultiIndex', property_name='name')

databricks/koalas/tests/test_indexes.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,13 @@ def test_multiindex_nlevel(self):
294294

295295
self.assertEqual(kdf.index.nlevels, 2)
296296

297+
def test_multiindex_from_arrays(self):
298+
arrays = [['a', 'a', 'b', 'b'], ['red', 'blue', 'red', 'blue']]
299+
pidx = pd.MultiIndex.from_arrays(arrays)
300+
kidx = ks.MultiIndex.from_arrays(arrays)
301+
302+
self.assert_eq(pidx, kidx)
303+
297304
def test_multiindex_levels(self):
298305
tuples = [[list('abc'), list('def')], [list('aac'), list('fed')]]
299306

0 commit comments

Comments
 (0)