Skip to content
This repository was archived by the owner on Feb 2, 2024. It is now read-only.

Commit 1ebf55c

Browse files
Allow global values to be used in read_csv params (names, usecols) (#998)
Motivation: fixes a bug in RewriteReadCsv rewrite that fails compilation when exception is raised during the match phase because variable defining value of read_csv parameter is not ir.Expr node but ir.Global (and thus has no op attribute).
1 parent 3b22273 commit 1ebf55c

File tree

2 files changed

+21
-1
lines changed

2 files changed

+21
-1
lines changed

sdc/rewrites/read_csv_consts.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ def match(self, func_ir, block, typemap, calltypes):
6363
if key in self._read_csv_const_args:
6464
arg_def = guard(get_definition, func_ir, var)
6565
ops = ['build_list', 'build_set', 'build_map']
66-
if arg_def.op in ops:
66+
if isinstance(arg_def, ir.Expr) and arg_def.op in ops:
6767
args.append(arg_def)
6868

6969
return len(args) > 0

sdc/tests/test_io.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,26 @@ def test_impl():
398398
sdc_func = self.jit(test_impl)
399399
pd.testing.assert_frame_equal(sdc_func(), test_impl())
400400

401+
def test_csv_infer_file_param_names_no_rewrite(self):
402+
"""Test verifies pandas read_csv impl supports names parameters as tuple of columns
403+
both captured as global or local variable """
404+
405+
def test_impl_1():
406+
const_names = ('A', 'B', 'C', 'D')
407+
return pd.read_csv("csv_data_date1.csv",
408+
names=const_names)
409+
410+
global_csv_names = ('A', 'B', 'C', 'D')
411+
412+
def test_impl_2():
413+
return pd.read_csv("csv_data_date1.csv",
414+
names=global_csv_names)
415+
416+
for test_impl in [test_impl_1, test_impl_2]:
417+
with self.subTest(tested_func_name=test_impl.__name__):
418+
sdc_func = self.jit(test_impl)
419+
pd.testing.assert_frame_equal(sdc_func(), test_impl())
420+
401421
def test_csv_infer_file_param_converters_1(self):
402422
"""Test verifies pandas read_csv impl supports conversion of all columns using converters parameter"""
403423

0 commit comments

Comments
 (0)