Koalas: DataFrame.transform operates per-element when transform function is type-annotated

[chunyang]

DataFrame.transform is behaving differently than its pandas equivalent when its return type is annotated. Consider the example below.

>>> import databricks.koalas as ks
>>> import numpy as np
>>> import pandas as pd
>>>
>>> pdf = pd.DataFrame([[1, 2, 3, 4], [2, 3, 4, 5]], columns=("a", "b", "c", "d"))
>>> kdf = ks.DataFrame([[1, 2, 3, 4], [2, 3, 4, 5]], columns=("a", "b", "c", "d"))
>>>
>>> def normalize(v):
...     return v / sum(v)
...
>>> pdf.transform(normalize)
          a    b         c         d
0  0.333333  0.4  0.428571  0.444444
1  0.666667  0.6  0.571429  0.555556
>>> kdf.transform(normalize)
          a    b         c         d                                            
0  0.333333  0.4  0.428571  0.444444
1  0.666667  0.6  0.571429  0.555556
>>>
>>> def typed_normalize(v) -> ks.Series[np.float64]:
...     return v / sum(v)
...
>>> pdf.transform(typed_normalize)
          a    b         c         d
0  0.333333  0.4  0.428571  0.444444
1  0.666667  0.6  0.571429  0.555556
>>> kdf.transform(typed_normalize)
     a    b    c    d                                                           
0  1.0  1.0  1.0  1.0
1  1.0  1.0  1.0  1.0

Koalas version: 0.24.0
Numpy version: 1.18.0
Pandas version: 0.25.3

[HyukjinKwon]

I believe expecting global aggregation within the function doesn't work in Koalas:

>>> def typed_normalize(v) -> ks.Series[np.float64]:
...     return v / sum(v)

because v will be multiple Series's as the segments of a whole Series v. I think the example above worked because it executed the shortcut path where it just uses pandas API (when the schema inference is required).

I think this limitation has to be clarified in the documentation at this moment.

[chunyang]

Ah I see, yes clarifying in the documentation would be great. Thanks! I can confirm the non-shortcut path without type annotations does not work.

>>> kdf = ks.DataFrame(np.ones((2000, 1))
>>> 1 == sum(kdf.transform(normalize).to_numpy())
array([False])

[chunyang]

Would it be possible to pass a koalas Series instead of a pandas Series to the transform function? Conceptually if we have a reference to the whole Series, then we can perform global aggregation on it?

[HyukjinKwon]

Valid point. Problem is there are many missing APIs in Koalas at this moment.

My intention was that we should expose pandas APIs as are so people can work around via transform or apply, for some APIs Koalas currently does not support.

[HyukjinKwon]

If we think about this case only, I think it can make sense to pass Koalas series but if we think about other APIs like groupby.apply(), it's difficult to pass Koalas series as is.

Probably we should revisit about this later. For now, I am sure you can work around via just looping Series in the Frame.

[chunyang]

Thanks for the explanation Hyukjin, I understand the reasoning better now. Looking forward to a revisit in the future, I think it could be a very user friendly feature!