Repository: IntelPython/sdc

Issue:
Support strings in read_csv()



Discussion:

Resolving patch:
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
@@ -743,6 +743,18 @@ def _handle_pd_DataFrame(self, assign, lhs, rhs, label):
         # # remove DataFrame call
         # return nodes
 
+    @staticmethod
+    def get_dtypes(df):
+        dtypes = []
+        for d in df.dtypes.values:
+            try:
+                numba_type = numba.typeof(d).dtype
+                array_type = types.Array(numba_type, 1, 'C')
+            except:
+                array_type = string_array_type
+            dtypes.append(array_type)
+        return dtypes
+
     def _handle_pd_read_csv(self, assign, lhs, rhs, label):
         """transform pd.read_csv(names=[A], dtype={'A': np.int32}) call
         """
@@ -781,8 +793,7 @@ def _handle_pd_read_csv(self, assign, lhs, rhs, label):
             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 = [types.Array(numba.typeof(d).dtype, 1, 'C')
-                      for d in df.dtypes.values]
+            dtypes = HiFramesPassImpl.get_dtypes(df)
             cols = df.columns.to_list()
             # overwrite column names like Pandas if explicitly provided
             if col_names != 0:
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
@@ -499,6 +499,7 @@ def pandas_read_csv(
 
     convert_options = pyarrow.csv.ConvertOptions(
         column_types=dtype,
+        strings_can_be_null=True,
         include_columns=include_columns,
     )
 


Test patch:
diff --git a/sdc/tests/gen_test_data.py b/sdc/tests/gen_test_data.py
--- a/sdc/tests/gen_test_data.py
+++ b/sdc/tests/gen_test_data.py
@@ -99,10 +99,10 @@ def generate_other_data():
     df.to_parquet("pivot2.pq")
 
     # CSV reader test
-    data = ("0,2.3,4.6,47736\n"
-            "1,2.3,4.6,47736\n"
-            "2,2.3,4.6,47736\n"
-            "4,2.3,4.6,47736\n")
+    data = ("0,2.3,4.6,A\n"
+            "1,2.3,4.6,B\n"
+            "2,2.3,4.6,\n"
+            "4,2.3,4.6,D\n")
 
     with open("csv_data1.csv", "w") as f:
         f.write(data)
diff --git a/sdc/tests/test_hpat_jit.py b/sdc/tests/test_hpat_jit.py
--- a/sdc/tests/test_hpat_jit.py
+++ b/sdc/tests/test_hpat_jit.py
@@ -354,7 +354,7 @@ def test_csv1_issue(self):
         def test_impl():
             return pd.read_csv("csv_data1.csv",
                                names=['A', 'B', 'C', 'D'],
-                               dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},
+                               dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},
                                )
         hpat_func = self.jit(test_impl)
         pd.testing.assert_frame_equal(hpat_func(), test_impl())
@@ -367,7 +367,7 @@ def test_impl():
                      '[right]: int32')
     def test_csv_keys1_issue(self):
         def test_impl():
-            dtype = {'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int}
+            dtype = {'A': np.int, 'B': np.float, 'C': np.float, 'D': str}
             return pd.read_csv("csv_data1.csv",
                                names=dtype.keys(),
                                dtype=dtype,
@@ -383,7 +383,7 @@ def test_impl():
                      '[right]: int32')
     def test_csv_const_dtype1_issue(self):
         def test_impl():
-            dtype = {'A': 'int', 'B': 'float64', 'C': 'float', 'D': 'int64'}
+            dtype = {'A': 'int', 'B': 'float64', 'C': 'float', 'D': 'str'}
             return pd.read_csv("csv_data1.csv",
                                names=dtype.keys(),
                                dtype=dtype,
@@ -401,7 +401,7 @@ def test_csv_skip1_issue(self):
         def test_impl():
             return pd.read_csv("csv_data1.csv",
                                names=['A', 'B', 'C', 'D'],
-                               dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},
+                               dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},
                                skiprows=2,
                                )
         hpat_func = self.jit(test_impl)
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
@@ -276,13 +276,13 @@ def pd_csv1(self):
             def test_impl():
                 return pd.read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.intp, 'B': np.float, 'C': np.float, 'D': np.intp},
+                                   dtype={'A': np.intp, 'B': np.float, 'C': np.float, 'D': str},
                                    )
         else:
             def test_impl():
                 return pd.read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},
+                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},
                                    )
         return test_impl
 
@@ -291,13 +291,13 @@ def pa_csv1(self):
             def test_impl():
                 return pd_read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.intp, 'B': np.float, 'C': np.float, 'D': np.intp},
+                                   dtype={'A': np.intp, 'B': np.float, 'C': np.float, 'D': str},
                                    )
         else:
             def test_impl():
                 return pd_read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},
+                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},
                                    )
         return test_impl
 
@@ -310,14 +310,14 @@ def test_csv1(self):
     def pd_csv_keys1(self):
         if platform.system() == 'Windows' and not IS_32BITS:
             def test_impl():
-                dtype = {'A': np.intp, 'B': np.float, 'C': np.float, 'D': np.intp}
+                dtype = {'A': np.intp, 'B': np.float, 'C': np.float, 'D': str}
                 return pd.read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
                                    )
         else:
             def test_impl():
-                dtype = {'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int}
+                dtype = {'A': np.int, 'B': np.float, 'C': np.float, 'D': str}
                 return pd.read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
@@ -327,14 +327,14 @@ def test_impl():
     def pa_csv_keys1(self):
         if platform.system() == 'Windows' and not IS_32BITS:
             def test_impl():
-                dtype = {'A': np.intp, 'B': np.float, 'C': np.float, 'D': np.intp}
+                dtype = {'A': np.intp, 'B': np.float, 'C': np.float, 'D': str}
                 return pd_read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
                                    )
         else:
             def test_impl():
-                dtype = {'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int}
+                dtype = {'A': np.int, 'B': np.float, 'C': np.float, 'D': str}
                 return pd_read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
@@ -350,14 +350,14 @@ def test_csv_keys1(self):
     def pd_csv_const_dtype1(self):
         if platform.system() == 'Windows' and not IS_32BITS:
             def test_impl():
-                dtype = {'A': 'int64', 'B': 'float64', 'C': 'float', 'D': 'int64'}
+                dtype = {'A': 'int64', 'B': 'float64', 'C': 'float', 'D': 'str'}
                 return pd.read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
                                    )
         else:
             def test_impl():
-                dtype = {'A': 'int', 'B': 'float64', 'C': 'float', 'D': 'int64'}
+                dtype = {'A': 'int', 'B': 'float64', 'C': 'float', 'D': 'str'}
                 return pd.read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
@@ -367,14 +367,14 @@ def test_impl():
     def pa_csv_const_dtype1(self):
         if platform.system() == 'Windows' and not IS_32BITS:
             def test_impl():
-                dtype = {'A': 'int64', 'B': 'float64', 'C': 'float', 'D': 'int64'}
+                dtype = {'A': 'int64', 'B': 'float64', 'C': 'float', 'D': 'str'}
                 return pd_read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
                                    )
         else:
             def test_impl():
-                dtype = {'A': 'int', 'B': 'float64', 'C': 'float', 'D': 'int64'}
+                dtype = {'A': 'int', 'B': 'float64', 'C': 'float', 'D': 'str'}
                 return pd_read_csv("csv_data1.csv",
                                    names=dtype.keys(),
                                    dtype=dtype,
@@ -406,13 +406,13 @@ def test_csv_infer1(self):
     def pd_csv_infer_parallel1(self):
         def test_impl():
             df = pd.read_csv("csv_data_infer1.csv")
-            return df.A.sum(), df.B.sum(), df.C.sum(), df.D.sum()
+            return df.A.sum(), df.B.sum(), df.C.sum()
         return test_impl
 
     def pa_csv_infer_parallel1(self):
         def test_impl():
             df = pd_read_csv("csv_data_infer1.csv")
-            return df.A.sum(), df.B.sum(), df.C.sum(), df.D.sum()
+            return df.A.sum(), df.B.sum(), df.C.sum()
         return test_impl
 
     @skip_numba_jit
@@ -426,14 +426,14 @@ def pd_csv_skip1(self):
             def test_impl():
                 return pd.read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.int64, 'B': np.float, 'C': np.float, 'D': np.int64},
+                                   dtype={'A': np.int64, 'B': np.float, 'C': np.float, 'D': str},
                                    skiprows=2,
                                    )
         else:
             def test_impl():
                 return pd.read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},
+                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},
                                    skiprows=2,
                                    )
         return test_impl
@@ -443,14 +443,14 @@ def pa_csv_skip1(self):
             def test_impl():
                 return pd_read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.int64, 'B': np.float, 'C': np.float, 'D': np.int64},
+                                   dtype={'A': np.int64, 'B': np.float, 'C': np.float, 'D': str},
                                    skiprows=2,
                                    )
         else:
             def test_impl():
                 return pd_read_csv("csv_data1.csv",
                                    names=['A', 'B', 'C', 'D'],
-                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},
+                                   dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},
                                    skiprows=2,
                                    )
         return test_impl
@@ -481,14 +481,14 @@ def pd_csv_infer_skip_parallel1(self):
         def test_impl():
             df = pd.read_csv("csv_data_infer1.csv", skiprows=2,
                              names=['A', 'B', 'C', 'D'])
-            return df.A.sum(), df.B.sum(), df.C.sum(), df.D.sum()
+            return df.A.sum(), df.B.sum(), df.C.sum()
         return test_impl
 
     def pa_csv_infer_skip_parallel1(self):
         def test_impl():
             df = pd_read_csv("csv_data_infer1.csv", skiprows=2,
                              names=['A', 'B', 'C', 'D'])
-            return df.A.sum(), df.B.sum(), df.C.sum(), df.D.sum()
+            return df.A.sum(), df.B.sum(), df.C.sum()
         return test_impl
 
     @skip_numba_jit
@@ -501,7 +501,7 @@ def pd_csv_rm_dead1(self):
         def test_impl():
             df = pd.read_csv("csv_data1.csv",
                              names=['A', 'B', 'C', 'D'],
-                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},)
+                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},)
             return df.B.values
         return test_impl
 
@@ -509,7 +509,7 @@ def pa_csv_rm_dead1(self):
         def test_impl():
             df = pd_read_csv("csv_data1.csv",
                              names=['A', 'B', 'C', 'D'],
-                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int},)
+                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str},)
             return df.B.values
         return test_impl
 
@@ -591,16 +591,16 @@ def pd_csv_parallel1(self):
         def test_impl():
             df = pd.read_csv("csv_data1.csv",
                              names=['A', 'B', 'C', 'D'],
-                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int})
-            return (df.A.sum(), df.B.sum(), df.C.sum(), df.D.sum())
+                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str})
+            return (df.A.sum(), df.B.sum(), df.C.sum())
         return test_impl
 
     def pa_csv_parallel1(self):
         def test_impl():
             df = pd_read_csv("csv_data1.csv",
                              names=['A', 'B', 'C', 'D'],
-                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': np.int})
-            return (df.A.sum(), df.B.sum(), df.C.sum(), df.D.sum())
+                             dtype={'A': np.int, 'B': np.float, 'C': np.float, 'D': str})
+            return (df.A.sum(), df.B.sum(), df.C.sum())
         return test_impl
 
     @skip_numba_jit