Support for Array type hints in APIs that take Python native functions (e.g., DataFrame.apply)

import databricks.koalas as ks

def tokenizeDF(col1) -> ks.Series[np.array(... ???)]:
    pass

[Callum027]

Hi there, I was wondering if there has been any progress with this? I’d like to be able to annotate my functions that generate lists, so it doesn’t have to infer the type and incur a performance penalty.

from typing import List
import databricks.koalas as ks

# Equivalent Spark UDF return type: ArrayType(StringType())
def to_list(x: str) -> List[str]:
    return [x]

ks.Series(["x", "y"]).apply(to_list)

Result:

Traceback (most recent call last)
<ipython-input-17-13d49d0c81a1> in <module>
      6     return [x]
      7 
----> 8 ks.Series(["x", "y"]).apply(to_list)

/opt/conda/lib/python3.7/site-packages/databricks/koalas/series.py in apply(self, func, args, **kwds)
   2664                 raise ValueError(
   2665                     "Expected the return type of this function to be of scalar type, "
-> 2666                     "but found type {}".format(sig_return)
   2667                 )
   2668             return_schema = sig_return.tpe

ValueError: Expected the return type of this function to be of scalar type, but found type UnknownType[typing.List[str]]

[itholic]

Will work on this from now on. Maybe we can do this like List[...] one.