Repository: IntelPython/sdc

Issue:
Support for empty strings in read_csv()
Closes #438 


Discussion:

Resolving patch:
diff --git a/sdc/datatypes/hpat_pandas_functions.py b/sdc/datatypes/hpat_pandas_functions.py
new file mode 100644
--- /dev/null
+++ b/sdc/datatypes/hpat_pandas_functions.py
@@ -0,0 +1,89 @@
+# *****************************************************************************
+# Copyright (c) 2019, 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.
+# *****************************************************************************
+
+'''
+| :module:`pandas` functions and operators implementations in Intel SDC
+'''
+
+import pandas
+import pyarrow.csv
+
+import numba
+import sdc
+
+from numba.extending import overload
+from sdc.hiframes.api import fix_df_array
+
+import pandas as pd
+from numba import types
+
+from sdc.io.csv_ext import (
+    _gen_csv_reader_py_pyarrow_py_func,
+    _gen_csv_reader_py_pyarrow_func_text_dataframe,
+)
+from sdc.hiframes.hiframes_untyped import HiFramesPassImpl
+
+
+def infer_usecols(col_names):
+    # usecols_var = self._get_arg('read_csv', rhs.args, kws, 6, 'usecols', '')
+    usecols = list(range(len(col_names)))
+    # if usecols_var != '':
+    #     err_msg = "pd.read_csv() usecols should be constant list of ints"
+    #     usecols = self._get_str_or_list(usecols_var, err_msg=err_msg, typ=int)
+    return usecols
+
+
+@overload(pandas.read_csv)
+def sdc_pandas_read_csv(fname, skiprows=0):
+    signature =        "fname, skiprows=0"  # nopep8
+
+    assert isinstance(fname, numba.types.Literal)
+    assert isinstance(skiprows, numba.types.Literal) or skiprows == 0
+
+    sep = ','
+
+    fname_const = fname.literal_value
+    if isinstance(skiprows, numba.types.Literal):
+        skiprows = skiprows.literal_value
+    col_names = 0
+    skiprows, col_names, dtype_map = \
+        HiFramesPassImpl.infer_column_names_and_types_from_constant_filename(
+            fname_const, skiprows, col_names)
+
+    usecols = infer_usecols(col_names)
+
+    date_cols = []
+    columns, out_types = HiFramesPassImpl._get_csv_col_info_core(dtype_map, date_cols, col_names)
+
+    # generate function text with signature and returning DataFrame
+    func_text, func_name = _gen_csv_reader_py_pyarrow_func_text_dataframe(
+        columns, out_types, usecols, sep, skiprows, signature)
+    # print(func_text)
+
+    # compile with Python
+    csv_reader_py = _gen_csv_reader_py_pyarrow_py_func(func_text, func_name)
+
+    return csv_reader_py
diff --git a/sdc/hiframes/hiframes_untyped.py b/sdc/hiframes/hiframes_untyped.py
--- a/sdc/hiframes/hiframes_untyped.py
+++ b/sdc/hiframes/hiframes_untyped.py
@@ -755,6 +755,23 @@ def get_dtypes(df):
             dtypes.append(array_type)
         return dtypes
 
+    @staticmethod
+    def infer_column_names_and_types_from_constant_filename(fname_const, skiprows, col_names):
+        rows_to_read = 100  # TODO: tune this
+        df = pd.read_csv(fname_const, nrows=rows_to_read, skiprows=skiprows)
+        # TODO: string_array, categorical, etc.
+        dtypes = HiFramesPassImpl.get_dtypes(df)
+        cols = df.columns.to_list()
+        # overwrite column names like Pandas if explicitly provided
+        if col_names != 0:
+            cols[-len(col_names):] = col_names
+        else:
+            # a row is used for names if not provided
+            skiprows += 1
+        col_names = cols
+        dtype_map = {c: d for c, d in zip(col_names, dtypes)}
+        return skiprows, col_names, dtype_map
+
     def _handle_pd_read_csv(self, assign, lhs, rhs, label):
         """transform pd.read_csv(names=[A], dtype={'A': np.int32}) call
         """
@@ -790,19 +807,9 @@ def _handle_pd_read_csv(self, assign, lhs, rhs, label):
             if fname_const is None:
                 raise ValueError("pd.read_csv() requires explicit type"
                                  "annotation using 'dtype' if filename is not constant")
-            rows_to_read = 100  # TODO: tune this
-            df = pd.read_csv(fname_const, nrows=rows_to_read, skiprows=skiprows)
-            # TODO: string_array, categorical, etc.
-            dtypes = HiFramesPassImpl.get_dtypes(df)
-            cols = df.columns.to_list()
-            # overwrite column names like Pandas if explicitly provided
-            if col_names != 0:
-                cols[-len(col_names):] = col_names
-            else:
-                # a row is used for names if not provided
-                skiprows += 1
-            col_names = cols
-            dtype_map = {c: d for c, d in zip(col_names, dtypes)}
+            skiprows, col_names, dtype_map = \
+                self.infer_column_names_and_types_from_constant_filename(
+                    fname_const, skiprows, col_names)
         else:
             dtype_map = guard(get_definition, self.state.func_ir, dtype_var)
             if (not isinstance(dtype_map, ir.Expr)
@@ -862,15 +869,13 @@ def _handle_pd_read_csv(self, assign, lhs, rhs, label):
         nodes[-1].target = lhs
         return nodes
 
-    def _get_csv_col_info(self, dtype_map, date_cols, col_names, lhs):
+    @staticmethod
+    def _get_csv_col_info_core(dtype_map, date_cols, col_names):
         if isinstance(dtype_map, types.Type):
             typ = dtype_map
-            data_arrs = [ir.Var(lhs.scope, mk_unique_var(cname), lhs.loc)
-                         for cname in col_names]
-            return col_names, data_arrs, [typ] * len(col_names)
+            return col_names, [typ] * len(col_names)
 
         columns = []
-        data_arrs = []
         out_types = []
         for i, (col_name, typ) in enumerate(dtype_map.items()):
             columns.append(col_name)
@@ -878,11 +883,23 @@ def _get_csv_col_info(self, dtype_map, date_cols, col_names, lhs):
             if i in date_cols:
                 typ = types.Array(types.NPDatetime('ns'), 1, 'C')
             out_types.append(typ)
+        return columns, out_types
+
+    @staticmethod
+    def _get_csv_col_info(dtype_map, date_cols, col_names, lhs):
+        col_names, out_types = HiFramesPassImpl._get_csv_col_info_core(dtype_map, date_cols, col_names)
+
+        if isinstance(dtype_map, types.Type):
+            data_arrs = [ir.Var(lhs.scope, mk_unique_var(cname), lhs.loc)
+                         for cname in col_names]
+            return col_names, data_arrs, out_types
+
+        data_arrs = []
+        for i, (col_name, typ) in enumerate(dtype_map.items()):
             # output array variable
             data_arrs.append(
                 ir.Var(lhs.scope, mk_unique_var(col_name), lhs.loc))
-
-        return columns, data_arrs, out_types
+        return col_names, data_arrs, out_types
 
     def _get_const_dtype(self, dtype_var):
         dtype_def = guard(get_definition, self.state.func_ir, dtype_var)
@@ -1925,3 +1942,6 @@ def simple_block_copy_propagate(block):
             for k in lhs_kill:
                 var_dict.pop(k, None)
     return
+
+
+from sdc.datatypes.hpat_pandas_functions import *
diff --git a/sdc/io/csv_ext.py b/sdc/io/csv_ext.py
--- a/sdc/io/csv_ext.py
+++ b/sdc/io/csv_ext.py
@@ -521,13 +521,21 @@ def pandas_read_csv(
 
 
 def _gen_csv_reader_py_pyarrow(col_names, col_typs, usecols, sep, typingctx, targetctx, parallel, skiprows):
+    func_text, func_name = _gen_csv_reader_py_pyarrow_func_text(col_names, col_typs, usecols, sep, skiprows)
+    csv_reader_py = _gen_csv_reader_py_pyarrow_py_func(func_text, func_name)
+    return _gen_csv_reader_py_pyarrow_jit_func(csv_reader_py)
+
+
+def _gen_csv_reader_py_pyarrow_func_text_core(col_names, col_typs, usecols, sep, skiprows, signature=None):
     # TODO: support non-numpy types like strings
     date_inds = ", ".join(str(i) for i, t in enumerate(col_typs) if t.dtype == types.NPDatetime('ns'))
     typ_strs = ", ".join(["{}='{}'".format(to_varname(cname), _get_dtype_str(t))
                           for cname, t in zip(col_names, col_typs)])
     pd_dtype_strs = ", ".join(["'{}':{}".format(cname, _get_pd_dtype_str(t)) for cname, t in zip(col_names, col_typs)])
 
-    func_text =  "def csv_reader_py(fname):\n"
+    if signature is None:
+        signature = "fname"
+    func_text = "def csv_reader_py({}):\n".format(signature)
     func_text += "  with objmode({}):\n".format(typ_strs)
     func_text += "    df = pandas_read_csv(fname, names={},\n".format(col_names)
     func_text += "       parse_dates=[{}],\n".format(date_inds)
@@ -537,16 +545,42 @@ def _gen_csv_reader_py_pyarrow(col_names, col_typs, usecols, sep, typingctx, tar
     for cname in col_names:
         func_text += "    {} = df['{}'].values\n".format(to_varname(cname), cname)
         # func_text += "    print({})\n".format(cname)
+    return func_text, 'csv_reader_py'
+
+
+def _gen_csv_reader_py_pyarrow_func_text(col_names, col_typs, usecols, sep, skiprows):
+    func_text, func_name = _gen_csv_reader_py_pyarrow_func_text_core(col_names, col_typs, usecols, sep, skiprows)
+
     func_text += "  return ({},)\n".format(", ".join(to_varname(c) for c in col_names))
 
+    return func_text, func_name
+
+
+def _gen_csv_reader_py_pyarrow_func_text_dataframe(col_names, col_typs, usecols, sep, skiprows, signature):
+    func_text, func_name = _gen_csv_reader_py_pyarrow_func_text_core(
+        col_names, col_typs, usecols, sep, skiprows, signature)
+
+    func_text += "  return sdc.hiframes.pd_dataframe_ext.init_dataframe({}, None, {})\n".format(
+        ", ".join(to_varname(c) for c in col_names),
+        ", ".join("'{}'".format(c) for c in col_names)
+    )
+
+    return func_text, func_name
+
+
+def _gen_csv_reader_py_pyarrow_py_func(func_text, func_name):
     # print(func_text)
     glbls = globals()  # TODO: fix globals after Numba's #3355 is resolved
     # {'objmode': objmode, 'csv_file_chunk_reader': csv_file_chunk_reader,
     # 'pd': pd, 'np': np}
     loc_vars = {}
     exec(func_text, glbls, loc_vars)
-    csv_reader_py = loc_vars['csv_reader_py']
+    csv_reader_py = loc_vars[func_name]
+
+    return csv_reader_py
+
 
+def _gen_csv_reader_py_pyarrow_jit_func(csv_reader_py):
     # TODO: no_cpython_wrapper=True crashes for some reason
     jit_func = numba.njit(csv_reader_py)
     compiled_funcs.append(jit_func)


Test patch:
diff --git a/buildscripts/test.py b/buildscripts/test.py
--- a/buildscripts/test.py
+++ b/buildscripts/test.py
@@ -209,7 +209,7 @@
         failed_examples = []
         expected_failures = []
         unexpected_success = []
-        expected_failures_list = ['basic_workflow.py', 'basic_workflow_parallel.py']
+        expected_failures_list = []
 
         os.chdir(sdc_examples)
         sdc_packages = get_sdc_build_packages(build_folder)
diff --git a/sdc/tests/test_io.py b/sdc/tests/test_io.py
--- a/sdc/tests/test_io.py
+++ b/sdc/tests/test_io.py
@@ -397,7 +397,6 @@ def test_impl():
             return pd_read_csv("csv_data_infer1.csv")
         return test_impl
 
-    @skip_numba_jit
     def test_csv_infer1(self):
         test_impl = self.pd_csv_infer1()
         hpat_func = self.jit(test_impl)
@@ -471,7 +470,6 @@ def test_impl():
             return pd_read_csv("csv_data_infer1.csv", skiprows=2)
         return test_impl
 
-    @skip_numba_jit
     def test_csv_infer_skip1(self):
         test_impl = self.pd_csv_infer_skip1()
         hpat_func = self.jit(test_impl)