Fix DataFrame.koalas.transform_batch to support additional dtypes.

[ueshin]

Fix DataFrame.koalas.transform_batch to support additional dtypes.

After this, additional dtypes can be specified in the return type annotation of the UDFs for DataFrame.koalas.transform_batch.

>>> kdf = ks.DataFrame(
...     {"a": ["a", "b", "c", "a", "b", "c"], "b": ["b", "a", "c", "c", "b", "a"]}
... )
>>> dtype = pd.CategoricalDtype(categories=["a", "b", "c", "d"])
>>> def to_category(pdf) -> ks.DataFrame["a":dtype, "b":dtype]:
...   return pdf.astype(dtype)
...
>>> applied = kdf.koalas.transform_batch(to_category)
>>> applied
   a  b
0  a  b
1  b  a
2  c  c
3  a  c
4  b  b
5  c  a
>>> applied.dtypes
a    category
b    category
dtype: object

[xinrong-meng]

Why is it called by_pass? Would you please help me understand?

[ueshin]

It uses some new Spark APIs to "by pass" a workaround:

You can see:

• koalas/databricks/koalas/accessors.py

  Lines 646 to 666 in 80a5893

  	if should_by_pass:
  	pudf = pandas_udf(
  	output_func, returnType=return_schema, functionType=PandasUDFType.SCALAR
  	)
  	temp_struct_column = verify_temp_column_name(
  	self_applied._internal.spark_frame, "__temp_struct__"
  	)
  	applied = pudf(F.struct(*columns)).alias(temp_struct_column)
  	sdf = self_applied._internal.spark_frame.select(applied)
  	sdf = sdf.selectExpr("%s.*" % temp_struct_column)
  	else:
  	applied = []
  	for field in return_schema.fields:
  	applied.append(
  	pandas_udf(
  	pandas_frame_func(output_func, field.name),
  	returnType=field.dataType,
  	functionType=PandasUDFType.SCALAR,
  	)(*columns).alias(field.name)
  	)
  	sdf = self_applied._internal.spark_frame.select(*applied)

• koalas/databricks/koalas/accessors.py

  Lines 715 to 735 in 80a5893

  	if should_by_pass:
  	pudf = pandas_udf(
  	output_func, returnType=return_schema, functionType=PandasUDFType.SCALAR
  	)
  	temp_struct_column = verify_temp_column_name(
  	self_applied._internal.spark_frame, "__temp_struct__"
  	)
  	applied = pudf(F.struct(*columns)).alias(temp_struct_column)
  	sdf = self_applied._internal.spark_frame.select(applied)
  	sdf = sdf.selectExpr("%s.*" % temp_struct_column)
  	else:
  	applied = []
  	for field in return_schema.fields:
  	applied.append(
  	pandas_udf(
  	pandas_frame_func(output_func, field.name),
  	returnType=field.dataType,
  	functionType=PandasUDFType.SCALAR,
  	)(*columns).alias(field.name)
  	)
  	sdf = self_applied._internal.spark_frame.select(*applied)

[xinrong-meng]

Looks great, thank you!

[ueshin]

Thanks! merging.