def test_selectByIndex(self):
dataLocal = [((1,), arange(12))]
index = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
data = Series(self.sc.parallelize(dataLocal), index=index)
result = data.selectByIndex(1)
assert_true(array_equal(result.values().first(), array([4, 5, 6, 7])))
assert_true(array_equal(result.index, array([1, 1, 1, 1])))
result = data.selectByIndex(1, squeeze=True)
assert_true(array_equal(result.index, array([0, 1, 2, 3])))
index = [
[0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
[0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1],
[0, 1, 0, 1, 2, 3, 0, 1, 0, 1, 2, 3]
]
data.index = array(index).T
result, mask = data.selectByIndex(0, level=2, returnMask=True)
assert_true(array_equal(result.values().first(), array([0, 2, 6, 8])))
assert_true(array_equal(result.index, array([[0, 0, 0], [0, 1, 0], [1, 0, 0], [1, 1, 0]])))
assert_true(array_equal(mask, array([1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0])))
result = data.selectByIndex(0, level=2, squeeze=True)
assert_true(array_equal(result.values().first(), array([0, 2, 6, 8])))
assert_true(array_equal(result.index, array([[0, 0], [0, 1], [1, 0], [1, 1]])))
result = data.selectByIndex([1, 0], level=[0, 1])
assert_true(array_equal(result.values().first(), array([6, 7])))
assert_true(array_equal(result.index, array([[1, 0, 0], [1, 0, 1]])))
result = data.selectByIndex(val=[0, [2,3]], level=[0, 2])
assert_true(array_equal(result.values().first(), array([4, 5])))
assert_true(array_equal(result.index, array([[0, 1, 2], [0, 1, 3]])))
result = data.selectByIndex(1, level=1, filter=True)
assert_true(array_equal(result.values().first(), array([0, 1, 6, 7])))
assert_true(array_equal(result.index, array([[0, 0, 0], [0, 0, 1], [1, 0, 0], [1, 0, 1]])))
