Add dt functions to koalas Series

Author: icexelloss

databricks/koalas/base.py

@@ -30,6 +30,7 @@
     to_arrow_type
 
 from databricks import koalas as ks  # For running doctests and reference resolution in PyCharm.
+from databricks.koalas.typedef import pandas_wraps
 
 
 def _column_op(f):
@@ -74,6 +75,25 @@ def wrapper(self, *args):
     return wrapper
 
 
+def _wrap_accessor_spark(accessor, fn, return_type=None):
+    """
+    Wrap an accessor property or method, e.g., Series.dt.date with a spark function.
+    """
+    if return_type:
+        return _column_op(
+            lambda col: fn(col).cast(return_type)
+        )(accessor._data)
+    else:
+        return _column_op(fn)(accessor._data)
+
+
+def _wrap_accessor_pandas(accessor, fn, return_type):
+    """
+    Wrap an accessor property or method, e.g, Series.dt.date with a pandas function.
+    """
+    return pandas_wraps(fn, return_col=return_type)(accessor._data)
+
+
 class IndexOpsMixin(object):
     """common ops mixin to support a unified interface / docs for Series / Index
 

databricks/koalas/datetime.py

@@ -0,0 +1,161 @@
+#
+# Copyright (C) 2019 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Date/Time related functions on Koalas Series
+"""
+import pyspark.sql.functions as F
+
+from pyspark.sql.types import DateType, TimestampType, LongType, StringType
+from databricks.koalas.base import (
+    _wrap_accessor_pandas,
+    _wrap_accessor_spark
+)
+
+import databricks.koalas as ks
+
+from databricks.koalas.utils import lazy_property
+
+
+class DatetimeMethods(object):
+    """Date/Time methods for Koalas Series"""
+    def __init__(self, series):
+        if not isinstance(series.spark_type, (DateType, TimestampType)):
+            raise ValueError(
+                "Cannot call DatetimeMethods on type {}"
+                .format(series.spark_type))
+        self._data = series
+
+    # Properties
+    @lazy_property
+    def date(self) -> ks.Series:
+        """
+        The date part of the datetime.
+        """
+        # TODO: Hit a weird exception
+        # syntax error in attribute name: `to_date(`start_date`)` with alias
+        return _wrap_accessor_spark(
+            self, lambda col: F.to_date(col).alias('date')
+        )
+
+    @lazy_property
+    def time(self) -> ks.Series:
+        raise NotImplementedError()
+
+    @lazy_property
+    def timetz(self) -> ks.Series:
+        raise NotImplementedError()
+
+    @lazy_property
+    def year(self) -> ks.Series:
+        """
+        The year of the datetime.
+        `"""
+        return _wrap_accessor_spark(self, F.year, LongType())
+
+    @lazy_property
+    def month(self) -> ks.Series:
+        """
+        The month of the timestamp as January = 1 December = 12.
+        """
+        return _wrap_accessor_spark(self, F.month, LongType())
+
+    @lazy_property
+    def day(self) -> ks.Series:
+        """
+        The days of the datetime.
+        """
+        return _wrap_accessor_spark(self, F.dayofmonth, LongType())
+
+    @lazy_property
+    def hour(self) -> ks.Series:
+        """
+        The hours of the datetime.
+        """
+        return _wrap_accessor_spark(self, F.hour, LongType())
+
+    @lazy_property
+    def minute(self) -> ks.Series:
+        """
+        The minutes of the datetime.
+        """
+        return _wrap_accessor_spark(self, F.minute, LongType())
+
+    @lazy_property
+    def second(self) -> ks.Series:
+        """
+        The seconds of the datetime.
+        """
+        return _wrap_accessor_spark(self, F.second, LongType())
+
+    @lazy_property
+    def millisecond(self) -> ks.Series:
+        """
+        The milliseconds of the datetime.
+        """
+        return _wrap_accessor_pandas(
+            self, lambda x: x.dt.millisecond, LongType())
+
+    @lazy_property
+    def microsecond(self) -> ks.Series:
+        """
+        The microseconds of the datetime.
+        """
+        return _wrap_accessor_pandas(
+            self, lambda x: x.dt.microsecond, LongType())
+
+    @lazy_property
+    def nanosecond(self) -> ks.Series:
+        raise NotImplementedError()
+
+    @lazy_property
+    def week(self) -> ks.Series:
+        """
+        The week ordinal of the year.
+        """
+        return _wrap_accessor_spark(self, F.weekofyear, LongType())
+
+    @lazy_property
+    def weekofyear(self) -> ks.Series:
+        """
+        The week ordinal of the year.
+        """
+        return _wrap_accessor_spark(self, F.weekofyear, LongType())
+
+    @lazy_property
+    def dayofweek(self) -> ks.Series:
+        """
+        The day of the week with Monday=0, Sunday=6.
+        """
+        return _wrap_accessor_pandas(self, lambda s: s.dt.dayofweek, LongType())
+
+    @lazy_property
+    def dayofyear(self) -> ks.Series:
+        """
+        The day ordinal of the year.
+        """
+        return _wrap_accessor_pandas(self, lambda s: s.dt.dayofyear, LongType())
+
+    # Methods
+    def strftime(self, date_format) -> ks.Series:
+        """
+        Convert to a String Series using specified date_format.
+        """
+        return _wrap_accessor_pandas(
+            self,
+            lambda x: x.dt.strftime(date_format),
+            StringType()
+        )

databricks/koalas/series.py

@@ -133,6 +133,11 @@ def _with_new_scol(self, scol: spark.Column) -> 'Series':
         """
         return Series(scol, anchor=self._kdf, index=self._index_map)
 
+    @property
+    def dt(self):
+        from databricks.koalas.datetime import DatetimeMethods
+        return DatetimeMethods(self)
+
     @property
     def dtypes(self):
         """Return the dtype object of the underlying data.

databricks/koalas/tests/test_series_datetime.py

@@ -18,19 +18,37 @@
 
 import numpy as np
 import pandas as pd
+import pandas.testing as mt
 
 from databricks import koalas
-from databricks.koalas.testing.utils import ReusedSQLTestCase, TestUtils
+from databricks.koalas.testing.utils import ReusedSQLTestCase, SQLTestUtils
 
 
-class SeriesDatetimeTest(ReusedSQLTestCase, TestUtils):
+class SeriesDateTimeTest(ReusedSQLTestCase, SQLTestUtils):
 
     @property
     def pdf1(self):
         date1 = pd.Series(pd.date_range('2012-1-1 12:00:00', periods=3, freq='M'))
         date2 = pd.Series(pd.date_range('2013-3-11 21:45:00', periods=3, freq='W'))
         return pd.DataFrame(dict(start_date=date1, end_date=date2))
 
+    @property
+    def pd_start_date(self):
+        return self.pdf1['start_date']
+
+    @property
+    def ks_start_date(self):
+        return koalas.from_pandas(self.pd_start_date)
+
+    def check_func(self, func):
+        # import pdb; pdb.set_trace()
+
+        mt.assert_series_equal(
+            func(self.ks_start_date).to_pandas(),
+            func(self.pd_start_date),
+            check_names=False
+        )
+
     @unittest.skip(
         "It fails in certain OSs presumably due to different "
         "timezone behaviours inherited from C library.")
@@ -49,3 +67,51 @@ def test_div(self):
             self.assert_eq(
                 (kdf['end_date'] - kdf['start_date']) / duration,
                 (pdf['end_date'] - pdf['start_date']) / duration)
+
+    def test_date(self):
+        self.check_func(lambda x: x.dt.date)
+
+    def test_time(self):
+        with self.assertRaises(NotImplementedError):
+            self.check_func(lambda x: x.dt.time)
+
+    def test_timetz(self):
+        with self.assertRaises(NotImplementedError):
+            self.check_func(lambda x: x.dt.timetz)
+
+    def test_year(self):
+        self.check_func(lambda x: x.dt.year)
+
+    def test_month(self):
+        self.check_func(lambda x: x.dt.month)
+
+    def test_day(self):
+        self.check_func(lambda x: x.dt.day)
+
+    def test_hour(self):
+        self.check_func(lambda x: x.dt.hour)
+
+    def test_minute(self):
+        self.check_func(lambda x: x.dt.minute)
+
+    def test_second(self):
+        self.check_func(lambda x: x.dt.second)
+
+    def test_microsecond(self):
+        self.check_func(lambda x: x.dt.microsecond)
+
+    def test_nanosecond(self):
+        with self.assertRaises(NotImplementedError):
+            self.check_func(lambda x: x.dt.nanosecond)
+
+    def test_week(self):
+        self.check_func(lambda x: x.dt.week)
+
+    def test_weekofyear(self):
+        self.check_func(lambda x: x.dt.weekofyear)
+
+    def test_dayofweek(self):
+        self.check_func(lambda x: x.dt.dayofweek)
+
+    def test_strftime(self):
+        self.check_func(lambda x: x.dt.strftime('%Y-%m-%d'))

databricks/koalas/typedef.py

@@ -18,6 +18,7 @@
 Utilities to deal with types. This is mostly focused on python3.
 """
 import typing
+import datetime
 from inspect import getfullargspec
 from functools import wraps
 
@@ -108,6 +109,7 @@ def _to_stype(tpe) -> X:
     types.FloatType(): [float, 'float', np.float],
     types.DoubleType(): [np.float64, 'float64', 'double'],
     types.TimestampType(): [np.datetime64],
+    types.DateType(): [datetime.date],
     types.BooleanType(): [bool, 'boolean', 'bool', np.bool],
 }
 
@@ -361,9 +363,10 @@ def wrapper(*args, **kwargs):
             spark_return_type = sig_return.inner
             return _make_fun(f, spark_return_type, *args, **kwargs)
         return wrapper
-    if return_col is not None or return_scalar is not None:
+    if callable(function):
+        return function_wrapper(function)
+    else:
         return function_wrapper
-    return function_wrapper(function)
 
 
 def _infer_return_type(f, return_col_hint=None, return_scalar_hint=None) -> X:

databricks/koalas/utils.py

@@ -17,6 +17,7 @@
 Commonly used utils in Koalas.
 """
 
+import functools
 from typing import Callable, Dict, Union
 
 from pyspark import sql as spark
@@ -90,6 +91,7 @@ def lazy_property(fn):
     attr_name = '_lazy_' + fn.__name__
 
     @property
+    @functools.wraps(fn)
     def _lazy_property(self):
         if not hasattr(self, attr_name):
             setattr(self, attr_name, fn(self))

docs/source/reference/series.rst

@@ -114,3 +114,27 @@ Serialization / IO / Conversion
    Series.to_json
    Series.to_csv
    Series.to_excel
+
+Datetime Methods
+----------------
+Methods accessible through `Series.dt`
+
+.. currentmodule:: databricks.koalas.datetime
+.. autosummary::
+   :toctree: api/
+
+   DatetimeMethods.date
+   DatetimeMethods.year
+   DatetimeMethods.month
+   DatetimeMethods.week
+   DatetimeMethods.weekofyear
+   DatetimeMethods.day
+   DatetimeMethods.dayofweek
+   DatetimeMethods.dayofyear
+   DatetimeMethods.hour
+   DatetimeMethods.minute
+   DatetimeMethods.second
+   DatetimeMethods.millisecond
+   DatetimeMethods.microsecond
+
+   DatetimeMethods.strftime