Repository: databricks/koalas

Issue:
Leverage NamedTuple and TypedDict for type hinting
https://docs.python.org/3/library/typing.html#typing.NamedTuple
https://docs.python.org/3/library/typing.html#typing.TypedDict

We can use both to avoid schema inference in APIs such as `DataFrame.transform`. This looks added from Python 3.6. Koalas supports Python 3.5 yet. So we should wait for now.


Discussion:
cc @thunterdb 

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -292,6 +292,43 @@
 T = TypeVar("T")
 
 
+def _create_tuple_for_frame_type(params):
+    from databricks.koalas.typedef import NameTypeHolder
+
+    if isinstance(params, zip):
+        params = [slice(name, tpe) for name, tpe in params]
+
+    if isinstance(params, slice):
+        params = (params,)
+
+    if (
+        hasattr(params, "__len__")
+        and isinstance(params, Iterable)
+        and all(isinstance(param, slice) for param in params)
+    ):
+        for param in params:
+            if isinstance(param.start, str) and param.step is not None:
+                raise TypeError(
+                    "Type hints should be specified as "
+                    "DataFrame['name': type]; however, got %s" % param
+                )
+
+        name_classes = []
+        for param in params:
+            new_class = type("NameType", (NameTypeHolder,), {})
+            new_class.name = param.start
+            # When the given argument is a numpy's dtype instance.
+            new_class.tpe = param.stop.type if isinstance(param.stop, np.dtype) else param.stop
+            name_classes.append(new_class)
+
+        return Tuple[tuple(name_classes)]
+
+    if not isinstance(params, Iterable):
+        params = [params]
+    params = [param.type if isinstance(param, np.dtype) else param for param in params]
+    return Tuple[tuple(params)]
+
+
 if (3, 5) <= sys.version_info < (3, 7):
     from typing import GenericMeta  # type: ignore
 
@@ -302,7 +339,7 @@
 
     def new_getitem(self, params):
         if hasattr(self, "is_dataframe"):
-            return old_getitem(self, Tuple[params])
+            return old_getitem(self, _create_tuple_for_frame_type(params))
         else:
             return old_getitem(self, params)
 
@@ -2033,7 +2070,16 @@ def apply_batch(self, func, args=(), **kwds):
 
             If the return type is specified, the output column names become
             `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
-            DataFrame in ``func``. See examples below.
+            DataFrame in ``func``.
+
+            To specify the column names, you can assign them in a pandas friendly style as below:
+
+            >>> def plus_one(x) -> ks.DataFrame["a": float, "b": float]:
+            ...     return x + 1
+
+            >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
+            >>> def plus_one(x) -> ks.DataFrame[zip(pdf.dtypes, pdf.columns)]:
+            ...     return x + 1
 
 
         Parameters
@@ -2074,6 +2120,12 @@ def apply_batch(self, func, args=(), **kwds):
            c0  c1
         0   1   2
 
+        >>> def query_func(pdf) -> ks.DataFrame["A": int, "B": int]:
+        ...     return pdf.query('A == 1')
+        >>> df.apply_batch(query_func)
+           A  B
+        0  1  2
+
         You can also omit the type hints so Koalas infers the return schema as below:
 
         >>> df.apply_batch(lambda pdf: pdf.query('A == 1'))
@@ -2237,7 +2289,16 @@ def apply(self, func, axis=0, args=(), **kwds):
 
             If the return type is specified as `DataFrame`, the output column names become
             `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
-            DataFrame in ``func``. See examples below.
+            DataFrame in ``func``.
+
+            To specify the column names, you can assign them in a pandas friendly style as below:
+
+            >>> def plus_one(x) -> ks.DataFrame["a": float, "b": float]:
+            ...     return x + 1
+
+            >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
+            >>> def plus_one(x) -> ks.DataFrame[zip(pdf.dtypes, pdf.columns)]:
+            ...     return x + 1
 
             However, this way switches the index type to default index type in the output
             because the type hint cannot express the index type at this moment. Use
@@ -2331,14 +2392,14 @@ def apply(self, func, axis=0, args=(), **kwds):
         In order to specify the types when `axis` is '1', it should use DataFrame[...]
         annotation. In this case, the column names are automatically generated.
 
-        >>> def identify(x) -> ks.DataFrame[np.int64, np.int64]:
+        >>> def identify(x) -> ks.DataFrame['A': np.int64, 'B': np.int64]:
         ...     return x
         ...
         >>> df.apply(identify, axis=1)
-           c0  c1
-        0   4   9
-        1   4   9
-        2   4   9
+           A  B
+        0  4  9
+        1  4  9
+        2  4  9
 
         You can also specify extra arguments.
 
@@ -2648,7 +2709,16 @@ def transform_batch(self, func, *args, **kwargs):
 
             If the return type is specified, the output column names become
             `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
-            DataFrame in ``func``. See examples below.
+            DataFrame in ``func``.
+
+            To specify the column names, you can assign them in a pandas friendly style as below:
+
+            >>> def plus_one(x) -> ks.DataFrame['a': float, 'b': float]:
+            ...     return x + 1
+
+            >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
+            >>> def plus_one(x) -> ks.DataFrame[zip(pdf.dtypes, pdf.columns)]:
+            ...     return x + 1
 
 
         Parameters
@@ -2686,6 +2756,14 @@ def transform_batch(self, func, *args, **kwargs):
         1   4   5
         2   6   7
 
+        >>> def plus_one_func(pdf) -> ks.DataFrame['A': int, 'B': int]:
+        ...     return pdf + 1
+        >>> df.transform_batch(plus_one_func)
+           A  B
+        0  2  3
+        1  4  5
+        2  6  7
+
         >>> def plus_one_func(pdf) -> ks.Series[int]:
         ...     return pdf.B + 1
         >>> df.transform_batch(plus_one_func)
@@ -10282,7 +10360,7 @@ def __class_getitem__(cls, params):
             # This is a workaround to support variadic generic in DataFrame in Python 3.7.
             # See https://github.com/python/typing/issues/193
             # we always wraps the given type hints by a tuple to mimic the variadic generic.
-            return Tuple[params]
+            return _create_tuple_for_frame_type(params)
 
     elif (3, 5) <= sys.version_info < (3, 7):
         # This is a workaround to support variadic generic in DataFrame in Python 3.5+
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -848,17 +848,26 @@ def apply(self, func, *args, **kwargs):
         use them before reaching for `apply`.
 
         .. note:: this API executes the function once to infer the type which is
-             potentially expensive, for instance, when the dataset is created after
-             aggregations or sorting.
+            potentially expensive, for instance, when the dataset is created after
+            aggregations or sorting.
 
-             To avoid this, specify return type in ``func``, for instance, as below:
+            To avoid this, specify return type in ``func``, for instance, as below:
+
+            >>> def pandas_div(x) -> ks.DataFrame[float, float]:
+            ...     return x[['B', 'C']] / x[['B', 'C']]
+
+            If the return type is specified, the output column names become
+            `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
+            DataFrame in ``func``.
+
+            To specify the column names, you can assign them in a pandas friendly style as below:
 
-             >>> def pandas_div(x) -> ks.DataFrame[float, float]:
-             ...     return x[['B', 'C']] / x[['B', 'C']]
+            >>> def pandas_div(x) -> ks.DataFrame["a": float, "b": float]:
+            ...     return x[['B', 'C']] / x[['B', 'C']]
 
-             If the return type is specified, the output column names become
-             `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
-             DataFrame in ``func``. See examples below.
+            >>> pdf = pd.DataFrame({'B': [1.], 'C': [3.]})
+            >>> def plus_one(x) -> ks.DataFrame[zip(pdf.columns, pdf.dtypes)]:
+            ...     return x[['B', 'C']] / x[['B', 'C']]
 
         .. note:: the dataframe within ``func`` is actually a pandas dataframe. Therefore,
             any pandas APIs within this function is allowed.
diff --git a/databricks/koalas/typedef/typehints.py b/databricks/koalas/typedef/typehints.py
--- a/databricks/koalas/typedef/typehints.py
+++ b/databricks/koalas/typedef/typehints.py
@@ -19,7 +19,7 @@
 """
 import typing
 import datetime
-from inspect import getfullargspec
+from inspect import getfullargspec, isclass
 
 import numpy as np
 import pandas as pd
@@ -48,12 +48,16 @@ def __repr__(self):
 
 
 class DataFrameType(object):
-    def __init__(self, tpe):
-        # Seems we cannot specify field names. I currently gave some default names
-        # `c0, c1, ... cn`.
-        self.tpe = types.StructType(
-            [types.StructField("c%s" % i, tpe[i]) for i in range(len(tpe))]
-        )  # type: types.StructType
+    def __init__(self, tpe, names=None):
+        if names is None:
+            # Default names `c0, c1, ... cn`.
+            self.tpe = types.StructType(
+                [types.StructField("c%s" % i, tpe[i]) for i in range(len(tpe))]
+            )  # type: types.StructType
+        else:
+            self.tpe = types.StructType(
+                [types.StructField(n, t) for n, t in zip(names, tpe)]
+            )  # type: types.StructType
 
     def __repr__(self):
         return "DataFrameType[{}]".format(self.tpe)
@@ -77,6 +81,11 @@ def __repr__(self):
         return "UnknownType[{}]".format(self.tpe)
 
 
+class NameTypeHolder(object):
+    name = None
+    tpe = None
+
+
 def as_spark_type(tpe) -> types.DataType:
     """
     Given a python type, returns the equivalent spark type.
@@ -109,8 +118,11 @@ def as_spark_type(tpe) -> types.DataType:
         return types.DateType()
     elif tpe in (bool, "boolean", "bool", np.bool):
         return types.BooleanType()
-    elif tpe in ():
+    elif tpe in (np.ndarray,):
+        # TODO: support other child types
         return types.ArrayType(types.StringType())
+    else:
+        raise TypeError("Type %s was not understood." % tpe)
 
 
 def spark_type_to_pandas_dtype(spark_type):
@@ -184,11 +196,33 @@ def infer_return_type(f) -> typing.Union[SeriesType, DataFrameType, ScalarType,
     ...    pass
     >>> infer_return_type(func).tpe
     StructType(List(StructField(c0,FloatType,true)))
+
+    >>> def func() -> ks.DataFrame['a': np.float, 'b': int]:
+    ...     pass
+    >>> infer_return_type(func).tpe
+    StructType(List(StructField(a,FloatType,true),StructField(b,IntegerType,true)))
+
+    >>> def func() -> "ks.DataFrame['a': np.float, 'b': int]":
+    ...     pass
+    >>> infer_return_type(func).tpe
+    StructType(List(StructField(a,FloatType,true),StructField(b,IntegerType,true)))
+
+    >>> pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
+    >>> def func() -> ks.DataFrame[pdf.dtypes]:
+    ...     pass
+    >>> infer_return_type(func).tpe
+    StructType(List(StructField(c0,LongType,true),StructField(c1,LongType,true)))
+
+    >>> pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
+    >>> def func() -> ks.DataFrame[zip(pdf.columns, pdf.dtypes)]:
+    ...     pass
+    >>> infer_return_type(func).tpe
+    StructType(List(StructField(a,LongType,true),StructField(b,LongType,true)))
     """
     # We should re-import to make sure the class 'SeriesType' is not treated as a class
     # within this module locally. See Series.__class_getitem__ which imports this class
     # canonically.
-    from databricks.koalas.typedef import SeriesType
+    from databricks.koalas.typedef import SeriesType, NameTypeHolder
 
     spec = getfullargspec(f)
     tpe = spec.annotations.get("return", None)
@@ -219,6 +253,12 @@ def infer_return_type(f) -> typing.Union[SeriesType, DataFrameType, ScalarType,
             parameters = getattr(tuple_type, "__tuple_params__")
         else:
             parameters = getattr(tuple_type, "__args__")
+        if len(parameters) > 0 and all(
+            isclass(p) and issubclass(p, NameTypeHolder) for p in parameters
+        ):
+            names = [p.name for p in parameters if issubclass(p, NameTypeHolder)]
+            types = [p.tpe for p in parameters if issubclass(p, NameTypeHolder)]
+            return DataFrameType([as_spark_type(t) for t in types], names)
         return DataFrameType([as_spark_type(t) for t in parameters])
     inner = as_spark_type(tpe)
     if inner is None:


Test patch:
diff --git a/databricks/koalas/tests/test_series.py b/databricks/koalas/tests/test_series.py
--- a/databricks/koalas/tests/test_series.py
+++ b/databricks/koalas/tests/test_series.py
@@ -1050,7 +1050,7 @@ def test_astype(self):
 
         self.assert_eq(kser.astype(bool), pser.astype(bool))
 
-        with self.assertRaisesRegex(ValueError, "Type int63 not understood"):
+        with self.assertRaisesRegex(TypeError, "not understood"):
             kser.astype("int63")
 
     def test_aggregate(self):
diff --git a/databricks/koalas/tests/test_typedef.py b/databricks/koalas/tests/test_typedef.py
--- a/databricks/koalas/tests/test_typedef.py
+++ b/databricks/koalas/tests/test_typedef.py
@@ -70,9 +70,124 @@ def func() -> pd.DataFrame[np.float]:
         expected = StructType([StructField("c0", FloatType())])
         self.assertEqual(infer_return_type(func).tpe, expected)
 
+        pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
+
+        def func() -> pd.DataFrame[pdf.dtypes]:  # type: ignore
+            pass
+
+        expected = StructType([StructField("c0", LongType()), StructField("c1", LongType())])
+        self.assertEqual(infer_return_type(func).tpe, expected)
+
     def test_if_pandas_implements_class_getitem(self):
         # the current type hint implementation of pandas DataFrame assumes pandas doesn't
         # implement '__class_getitem__'. This test case is to make sure pandas
         # doesn't implement them.
         assert not ks._frame_has_class_getitem
         assert not ks._series_has_class_getitem
+
+    @unittest.skipIf(
+        sys.version_info < (3, 7),
+        "Type inference from pandas instances is supported with Python 3.7+",
+    )
+    def test_infer_schema_with_names_pandas_instances(self):
+        def func() -> 'pd.DataFrame["a" : np.float, "b":str]':
+            pass
+
+        expected = StructType([StructField("a", FloatType()), StructField("b", StringType())])
+        self.assertEqual(infer_return_type(func).tpe, expected)
+
+        def func() -> "pd.DataFrame['a': np.float, 'b': int]":
+            pass
+
+        expected = StructType([StructField("a", FloatType()), StructField("b", IntegerType())])
+        self.assertEqual(infer_return_type(func).tpe, expected)
+
+        pdf = pd.DataFrame({"a": [1, 2, 3], "b": [3, 4, 5]})
+
+        def func() -> pd.DataFrame[zip(pdf.columns, pdf.dtypes)]:
+            pass
+
+        expected = StructType([StructField("a", LongType()), StructField("b", LongType())])
+        self.assertEqual(infer_return_type(func).tpe, expected)
+
+    @unittest.skipIf(
+        sys.version_info < (3, 7),
+        "Type inference from pandas instances is supported with Python 3.7+",
+    )
+    def test_infer_schema_with_names_pandas_instances_negative(self):
+        def try_infer_return_type():
+            def f() -> 'pd.DataFrame["a" : np.float : 1, "b":str:2]':
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type)
+
+        class A:
+            pass
+
+        def try_infer_return_type():
+            def f() -> pd.DataFrame[A]:
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "not understood", try_infer_return_type)
+
+        def try_infer_return_type():
+            def f() -> 'pd.DataFrame["a" : np.float : 1, "b":str:2]':
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type)
+
+        # object type
+        pdf = pd.DataFrame({"a": ["a", 2, None]})
+
+        def try_infer_return_type():
+            def f() -> pd.DataFrame[pdf.dtypes]:  # type: ignore
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "object.*not understood", try_infer_return_type)
+
+    def test_infer_schema_with_names_negative(self):
+        def try_infer_return_type():
+            def f() -> 'ks.DataFrame["a" : np.float : 1, "b":str:2]':
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type)
+
+        class A:
+            pass
+
+        def try_infer_return_type():
+            def f() -> ks.DataFrame[A]:
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "not understood", try_infer_return_type)
+
+        def try_infer_return_type():
+            def f() -> 'ks.DataFrame["a" : np.float : 1, "b":str:2]':
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "Type hints should be specified", try_infer_return_type)
+
+        # object type
+        pdf = pd.DataFrame({"a": ["a", 2, None]})
+
+        def try_infer_return_type():
+            def f() -> ks.DataFrame[pdf.dtypes]:  # type: ignore
+                pass
+
+            infer_return_type(f)
+
+        self.assertRaisesRegex(TypeError, "object.*not understood", try_infer_return_type)