在Koalas DataFrame中按索引值筛选出现PandasNotImplementedError如何解决

解决方案 #

有两种可行的实现方式，效果和Pandas侧原代码完全一致：

方式1：索引转序列后筛选 #

import databricks.koalas as ks

df = ks.DataFrame({'foo':['a','b','c','d','e'], 'bar':['1', '2', '3','4','5']})
ci = [4,32,12,1]

# 将索引对象转为Koalas Series后调用isin方法
result = df[df.index.get_level_values(0).to_series().isin(ci)]

执行输出：

foo bar
1   b   2
4   e   5

方式2：重置索引筛选后还原 #

该方式对单级、多级索引都兼容：

result = df.reset_index()\
           .query("index in @ci")\
           .set_index("index")

报错原因 #

Koalas为了避免大数据量场景下全量索引被拉取到Driver节点导致内存溢出，没有实现pd.Index.__iter__()方法。直接调用Index.isin()时内部会触发索引迭代逻辑，因此抛出未实现异常。将索引转为普通Koalas Series后调用isin，底层会转换为Spark分布式算子执行，无需迭代索引即可完成筛选。

内容的提问来源于stack exchange，提问作者Yuvaraj Subramanian