AlexanderKalistratov · Apr 21, 2020 · Mar 27, 2020 · Apr 7, 2020 · Apr 8, 2020 · Apr 8, 2020
diff --git a/examples/dataframe/dataframe_loc.py b/examples/dataframe/dataframe_loc.py
@@ -0,0 +1,47 @@
+# *****************************************************************************
+# Copyright (c) 2020, Intel Corporation All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#     Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+
+"""
+   Expected result:
+    A    3.0
+    B    6.0
+    C    2.0
+    Name: 2, dtype: float64
+"""
+
+import pandas as pd
+from numba import njit
+
+
+@njit
+def dataframe_loc():
+    df = pd.DataFrame({'A': [1.0, 2.0, 3.0, 1.0], 'B': [4, 5, 6, 7], 'C': [4, 5, 2, 1]})
+
+    return df.loc[2]
+
+
+print(dataframe_loc())
diff --git a/sdc/datatypes/hpat_pandas_dataframe_functions.py b/sdc/datatypes/hpat_pandas_dataframe_functions.py
@@ -1867,13 +1867,63 @@ def _df_getitem_unicode_idx_impl(self, idx):
     ty_checker.raise_exc(idx, expected_types, 'idx')
 
 
+def df_getitem_single_label_loc_codegen(self, idx):
+    """
+    Example of generated implementation:
+        def _df_getitem_single_label_loc_impl(self, idx):
+            for i in numba.prange(len(self._dataframe.index)):
+                if self._dataframe._index[i] == idx:
+                data_0 = pandas.Series(self._dataframe._data[0], index=self._dataframe.index)
+                result_0 = data_0.at[idx]
+                data_1 = pandas.Series(self._dataframe._data[1], index=self._dataframe.index)
+                result_1 = data_1.at[idx]
+                return pandas.Series(data=[result_0[0], result_1[0]], index=['A', 'B'], name=str(idx))
+            raise IndexingError('Index is out of bounds for axis')
+    """
+    func_lines = ['def _df_getitem_single_label_loc_impl(self, idx):',
+                  '  for i in numba.prange(len(self._dataframe.index)):',
+                  '    if self._dataframe._index[i] == idx:']
+    if isinstance(self.index, types.NoneType):
+        func_lines = ['def _df_getitem_single_label_loc_impl(self, idx):',
+                      '  if -1 < idx < len(self._dataframe._data):']
+    results = []
+    result_index = []
+    for i, c in enumerate(self.columns):
+        result_c = f"result_{i}"
+        func_lines += [f"      data_{i} = pandas.Series(self._dataframe._data[{i}], index=self._dataframe.index)",
+                       f"      {result_c} = data_{i}.at[idx]"]
+        results.append(result_c)
+        result_index.append(c)
+    data = '[0], '.join(col for col in results) + '[0]'
+    func_lines += [f"      return pandas.Series(data=[{data}], index={result_index}, name=str(idx))",
+                   f"  raise IndexingError('Index is out of bounds for axis')"]
+
+    func_text = '\n'.join(func_lines)
+    global_vars = {'pandas': pandas, 'numpy': numpy,
+                   'numba': numba,
+                   'IndexingError': IndexingError}
+
+    return func_text, global_vars
+
+
+gen_df_getitem_loc_single_label_impl = gen_impl_generator(
+    df_getitem_single_label_loc_codegen, '_df_getitem_single_label_loc_impl')
+
+
 @sdc_overload(operator.getitem)
 def sdc_pandas_dataframe_accessor_getitem(self, idx):
     if not isinstance(self, DataFrameGetitemAccessorType):
         return None
 
     accessor = self.accessor.literal_value
 
+    if accessor == 'loc':
+        if isinstance(idx, types.Integer):
+            return gen_df_getitem_loc_single_label_impl(self.dataframe, idx)
+
+        raise TypingError('Operator getitem(). The index must be a single label, a  list or array of labels,\
+                          a slice object with labels, a boolean array or a callable. Given: {}'.format(idx))
-        raise TypingError('Operator getitem(). The index must be a single label, a  list or array of labels,\
-                          a slice object with labels, a boolean array or a callable. Given: {}'.format(idx))
+        ty_checker = TypeChecker('Operator loc().')
+        ty_checker.raise_exc(idx, 'int', 'idx')
-        raise TypingError('Operator getitem(). The index must be a single label, a  list or array of labels,\
-                          a slice object with labels, a boolean array or a callable. Given: {}'.format(idx))
+        ty_checker = TypeChecker('Operator loc().')
+        ty_checker.raise_exc(idx, 'int', 'idx')
+
     if accessor == 'iat':
         if isinstance(idx, types.Tuple) and isinstance(idx[1], types.Literal):
             col = idx[1].literal_value
@@ -1944,6 +1994,52 @@ def sdc_pandas_dataframe_iat_impl(self):
     return sdc_pandas_dataframe_iat_impl
 
 
+@sdc_overload_attribute(DataFrameType, 'loc')
+def sdc_pandas_dataframe_loc(self):
+    """
+    Intel Scalable Dataframe Compiler User Guide
+    ********************************************
+
+    Pandas API: pandas.DataFrame.loc
+
+    Examples
+    --------
+    .. literalinclude:: ../../../examples/dataframe/dataframe_loc.py
+       :language: python
+       :lines: 34-
-       :lines: 34-
+       :lines: 36-
-       :lines: 34-
+       :lines: 36-
+       :caption: Access a group of rows and columns by label(s) or a boolean array.
+       :name: ex_dataframe_loc
+
+    .. command-output:: python ./dataframe/dataframe_loc.py
+       :cwd: ../../../examples
+
+    .. seealso::
+        :ref:`DataFrame.at <pandas.DataFrame.at>`
+            Access a single value for a row/column label pair.
+        :ref:`DataFrame.iloc <pandas.DataFrame.iloc>`
+            Access group of rows and columns by integer position(s).
+        :ref:`DataFrame.xs <pandas.DataFrame.xs>`
+            Returns a cross-section (row(s) or column(s)) from the Series/DataFrame.
+        :ref:`Series.loc <pandas.Series.loc>`
+            Access group of values using labels.
+
+    Intel Scalable Dataframe Compiler Developer Guide
+    *************************************************
+    Pandas DataFrame method :meth:`pandas.DataFrame.loc` implementation.
+
+    .. only:: developer
+        Test: python -m sdc.runtests -k sdc.tests.test_dataframe.TestDataFrame.test_df_loc*
+    """
+
+    ty_checker = TypeChecker('Attribute loc().')
+    ty_checker.check(self, DataFrameType)
+
+    def sdc_pandas_dataframe_loc_impl(self):
+        return sdc.datatypes.hpat_pandas_dataframe_getitem_types.dataframe_getitem_accessor_init(self, 'loc')
+
+    return sdc_pandas_dataframe_loc_impl
+
+
 @sdc_overload_method(DataFrameType, 'pct_change')
 def pct_change_overload(df, periods=1, fill_method='pad', limit=None, freq=None):
     """

diff --git a/sdc/tests/test_dataframe.py b/sdc/tests/test_dataframe.py
@@ -1100,6 +1100,17 @@ def test_impl(df):
         msg = 'Index is out of bounds for axis'
         self.assertIn(msg, str(raises.exception))
 
+    def test_df_loc(self):
+        def test_impl(df):
+            return df.loc[4]
+
+        sdc_func = sdc.jit(test_impl)
+        idx = [3, 4, 1, 2, 0]
+        df = pd.DataFrame({"A": [3.2, 4.4, 7.0, 3.3, 1.0],
+                           "B": [3, 4, 1, 0, 222],
+                           "C": [3.1, 8.4, 7.1, 3.2, 1]}, index=idx)
+        pd.testing.assert_series_equal(sdc_func(df), test_impl(df), check_names=False)
+
     def test_df_head(self):
         def get_func(n):
             def impl(a):