Skip to content

Commit 6492dd2

Browse files
authored
Add is_type_compatible for Index & MultiIndex (#1765)
This PR proposes `is_type_compatible()` for `Index` & `MultiIndex` ```python >>> kidx = ks.Index([1, 2, 3]) >>> kidx.is_type_compatible('integer') True >>> kidx = ks.Index([1.0, 2.0, 3.0]) >>> kidx.is_type_compatible('integer') False >>> kidx.is_type_compatible('floating') True ```
1 parent defc229 commit 6492dd2

File tree

3 files changed

+50
-2
lines changed

3 files changed

+50
-2
lines changed

databricks/koalas/indexes.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -927,6 +927,24 @@ def is_object(self):
927927
"""
928928
return is_object_dtype(self.dtype)
929929

930+
def is_type_compatible(self, kind):
931+
"""
932+
Whether the index type is compatible with the provided type.
933+
934+
Examples
935+
--------
936+
>>> kidx = ks.Index([1, 2, 3])
937+
>>> kidx.is_type_compatible('integer')
938+
True
939+
940+
>>> kidx = ks.Index([1.0, 2.0, 3.0])
941+
>>> kidx.is_type_compatible('integer')
942+
False
943+
>>> kidx.is_type_compatible('floating')
944+
True
945+
"""
946+
return kind == self.inferred_type
947+
930948
def dropna(self):
931949
"""
932950
Return Index or MultiIndex without NA/NaN values

databricks/koalas/missing/indexes.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ class MissingPandasLikeIndex(object):
5353
intersection = _unsupported_function("intersection")
5454
is_ = _unsupported_function("is_")
5555
is_lexsorted_for_tuple = _unsupported_function("is_lexsorted_for_tuple")
56-
is_type_compatible = _unsupported_function("is_type_compatible")
5756
join = _unsupported_function("join")
5857
map = _unsupported_function("map")
5958
putmask = _unsupported_function("putmask")
@@ -124,7 +123,6 @@ class MissingPandasLikeMultiIndex(object):
124123
is_ = _unsupported_function("is_")
125124
is_lexsorted = _unsupported_function("is_lexsorted")
126125
is_lexsorted_for_tuple = _unsupported_function("is_lexsorted_for_tuple")
127-
is_type_compatible = _unsupported_function("is_type_compatible")
128126
join = _unsupported_function("join")
129127
map = _unsupported_function("map")
130128
putmask = _unsupported_function("putmask")

databricks/koalas/tests/test_indexes.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1446,6 +1446,38 @@ def test_inferred_type(self):
14461446
kmidx = ks.from_pandas(pmidx)
14471447
self.assert_eq(pmidx.inferred_type, kmidx.inferred_type)
14481448

1449+
def test_is_type_compatible(self):
1450+
data_types = ["integer", "floating", "string", "boolean"]
1451+
# Integer
1452+
pidx = pd.Index([1, 2, 3])
1453+
kidx = ks.from_pandas(pidx)
1454+
for data_type in data_types:
1455+
self.assert_eq(pidx.is_type_compatible(data_type), kidx.is_type_compatible(data_type))
1456+
1457+
# Floating
1458+
pidx = pd.Index([1.0, 2.0, 3.0])
1459+
kidx = ks.from_pandas(pidx)
1460+
for data_type in data_types:
1461+
self.assert_eq(pidx.is_type_compatible(data_type), kidx.is_type_compatible(data_type))
1462+
1463+
# String
1464+
pidx = pd.Index(["a", "b", "c"])
1465+
kidx = ks.from_pandas(pidx)
1466+
for data_type in data_types:
1467+
self.assert_eq(pidx.is_type_compatible(data_type), kidx.is_type_compatible(data_type))
1468+
1469+
# Boolean
1470+
pidx = pd.Index([True, False, True, False])
1471+
kidx = ks.from_pandas(pidx)
1472+
for data_type in data_types:
1473+
self.assert_eq(pidx.is_type_compatible(data_type), kidx.is_type_compatible(data_type))
1474+
1475+
# MultiIndex
1476+
pmidx = pd.MultiIndex.from_tuples([("a", "x")])
1477+
kmidx = ks.from_pandas(pmidx)
1478+
for data_type in data_types:
1479+
self.assert_eq(pmidx.is_type_compatible(data_type), kmidx.is_type_compatible(data_type))
1480+
14491481
def test_asi8(self):
14501482
# Integer
14511483
pidx = pd.Index([1, 2, 3])

0 commit comments

Comments
 (0)