Merge 617d972 into 8e3a220

hpat/datatypes/hpat_pandas_series_functions.py

@@ -534,6 +534,44 @@ def hpat_pandas_series_copy_impl(self, deep=True):
             return hpat_pandas_series_copy_impl
 
 
+@overload_method(SeriesType, 'head')
+def hpat_pandas_series_head(self, n=5):
+    """
+    Pandas Series method :meth:`pandas.Series.head` implementation.
+
+    .. only:: developer
+       Test: python -m -k hpat.runtests hpat.tests.test_series.TestSeries.test_series_head*
+
+    Parameters
+    -----------
+    n: :obj:`int`, default 5
+        input argument, default 5
+    Returns
+    -------
+    :obj:`pandas.Series`
+        returns: The first n rows of the caller object.
+    """
+
+    _func_name = 'Method head().'
+
+    if not isinstance(self, SeriesType):
+        raise TypingError('{} The object must be a pandas.series. Given: {}'.format(_func_name, self))
+
+    if not isinstance(n, (types.Integer, types.Omitted)) and n != 5:
+        raise TypingError('{} The parameter must be an integer type. Given type n: {}'.format(_func_name, n))
+
+    if isinstance(self.index, types.NoneType):
+        def hpat_pandas_series_head_impl(self, n=5):
+            return pandas.Series(self._data[:n])
+
+        return hpat_pandas_series_head_impl
+    else:
+        def hpat_pandas_series_head_index_impl(self, n=5):
+            return pandas.Series(self._data[:n], self._index[:n], self._name)
+
+        return hpat_pandas_series_head_index_impl
+
+
 @overload_method(SeriesType, 'groupby')
 def hpat_pandas_series_groupby(
         self,

hpat/hiframes/pd_series_ext.py

@@ -996,6 +996,7 @@ def generic_expand_cumulative_series(self, args, kws):
                            'resolve_take', 'resolve_max', 'resolve_min', 'resolve_nunique',
                            'resolve_prod', 'resolve_count']
 
+
 # use ArrayAttribute for attributes not defined in SeriesAttribute
 for attr, func in numba.typing.arraydecl.ArrayAttribute.__dict__.items():
     if (attr.startswith('resolve_')

hpat/tests/test_series.py

@@ -1987,6 +1987,86 @@ def test_impl(S):
         pd.testing.assert_series_equal(hpat_func(S[start:end]), test_impl(S))
         self.assertTrue(count_array_OneDs() > 0)
 
+    def test_series_head_noidx_float(self):
+        def test_impl(S, n):
+            return S.head(n)
+        hpat_func = hpat.jit(test_impl)
+        for input_data in test_global_input_data_float64:
+            S = pd.Series(input_data)
+            for n in [-1, 0, 2, 3]:
+                result_ref = test_impl(S, n)
+                result_jit = hpat_func(S, n)
+                pd.testing.assert_series_equal(result_jit, result_ref)
+
+    @unittest.skip("Need fix test_global_input_data_integer64")
+    def test_series_head_noidx_int(self):
+        def test_impl(S, n):
+            return S.head(n)
+        hpat_func = hpat.jit(test_impl)
+        for input_data in test_global_input_data_integer64:
+            S = pd.Series(input_data)
+            for n in [-1, 0, 2, 3]:
+                result_ref = test_impl(S, n)
+                result_jit = hpat_func(S, n)
+                pd.testing.assert_series_equal(result_jit, result_ref)
+
+    @unittest.skip("Need fix test_global_input_data_integer64")
+    def test_series_head_noidx_num(self):
+        def test_impl(S, n):
+            return S.head(n)
+        hpat_func = hpat.jit(test_impl)
+        for input_data in test_global_input_data_numeric:
+            S = pd.Series(input_data)
+            for n in [-1, 0, 2, 3]:
+                result_ref = test_impl(S, n)
+                result_jit = hpat_func(S, n)
+                pd.testing.assert_series_equal(result_jit, result_ref)
+
+    @unittest.skip("Old implementation not work with n negative and data str")
+    def test_series_head_noidx_str(self):
+        def test_impl(S, n):
+            return S.head(n)
+        hpat_func = hpat.jit(test_impl)
+        input_data = test_global_input_data_unicode_kind4
+        S = pd.Series(input_data)
+        for n in [-1, 0, 2, 3]:
+            result_ref = test_impl(S, n)
+            result_jit = hpat_func(S, n)
+            pd.testing.assert_series_equal(result_jit, result_ref)
+
+    @unittest.skip("Broke another three tests")
+    def test_series_head_idx(self):
+        def test_impl(S):
+            return S.head()
+
+        def test_impl_param(S, n):
+            return S.head(n)
+
+        hpat_func = hpat.jit(test_impl)
+
+        data_test = [[6, 6, 2, 1, 3, 3, 2, 1, 2],
+                     [1.1, 0.3, 2.1, 1, 3, 0.3, 2.1, 1.1, 2.2],
+                     [6, 6.1, 2.2, 1, 3, 0, 2.2, 1, 2],
+                     ['as', 'b', 'abb', 'sss', 'ytr65', '', 'qw', 'a', 'b'],
+                     [6, 6, 2, 1, 3, np.inf, np.nan, np.nan, np.nan],
+                     [3., 5.3, np.nan, np.nan, np.inf, np.inf, 4.4, 3.7, 8.9]
+                     ]
+
+        for input_data in data_test:
+            for index_data in data_test:
+                S = pd.Series(input_data, index_data)
+
+                result_ref = test_impl(S)
+                result = hpat_func(S)
+                pd.testing.assert_series_equal(result, result_ref)
+
+                hpat_func_param1 = hpat.jit(test_impl_param)
+
+                for param1 in [1, 3, 7]:
+                    result_param1_ref = test_impl_param(S, param1)
+                    result_param1 = hpat_func_param1(S, param1)
+                    pd.testing.assert_series_equal(result_param1, result_param1_ref)
+
     def test_series_median1(self):
         '''Verifies median implementation for float and integer series of random data'''
         def test_impl(S):