Use of koalas instead of pandas for numpy where function

I am new to koalas. I have been told to implement koalas instead of pandas in my work. Earlier when we have dataframe we convert that to pandas and use that for np.where with condition check inside. Example in pandas we used to do like

np.where(condition,action1,action2)

When I try to use koalas for the same we get error below

PandasNotImplementedError: The method pd.Series.__iter__() is not implemented. If you want to collect your data as an NumPy array, use 'to_numpy()' instead.

I even tried ks.series and ks.dataframe but the error didn't go.

Is there any method/function in koalas to accept 3 params (condition,action1,action2) like we use np.where in pandas. It will be much helpful if anyone explains through example too.

Solution

A possible solution similar to np.where is to create a function yourself (like below) that uses koalas' where:

import databricks.koalas as ks


# sample dataframe
df = ks.DataFrame({
  'id': [1, 2, 3, 4, 5],
  'cost': [5000, 4000, 3000, 4500, 2000],
  'class': ['A', 'A', 'B', 'C', 'A']
})


# your custom function
def numpy_where(s, cond, action1, action2):
  return s.where(cond, action2).where(~cond, action1)


# create sample new column
df['new_col'] = numpy_where(df['class'], df['class'] == 'A', 'yes', 'no')
print(df)
#    id  cost class new_col
# 0   1  5000     A     yes
# 1   2  4000     A     yes
# 2   3  3000     B      no
# 3   4  4500     C      no
# 4   5  2000     A     yes

Basically:

s is the ks.Series on which you compute the where
cond is the condition to be satisfied
action1 is the value to be inserted when cond is true
action2 is the value to be inserted when cond is false