From 60287e1e8d3d54084da9e126670f2c87843861aa Mon Sep 17 00:00:00 2001 From: Thomas Li Date: Mon, 1 Jul 2024 17:56:34 +0000 Subject: [PATCH] address more comments --- python/cudf/cudf/_lib/pylibcudf/io/types.pyx | 8 ++++---- python/cudf/cudf/pylibcudf_tests/io/test_json.py | 9 ++++----- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/python/cudf/cudf/_lib/pylibcudf/io/types.pyx b/python/cudf/cudf/_lib/pylibcudf/io/types.pyx index cc609f1ca94..c71d4b27ec4 100644 --- a/python/cudf/cudf/_lib/pylibcudf/io/types.pyx +++ b/python/cudf/cudf/_lib/pylibcudf/io/types.pyx @@ -205,15 +205,15 @@ cdef class SinkInfo: if not sinks: raise ValueError("Need to pass at least one sink") + if isinstance(sinks[0], os.PathLike): + sinks = [os.path.expanduser(s) for s in sinks] + cdef object initial_sink_cls = type(sinks[0]) if not all(isinstance(s, initial_sink_cls) for s in sinks): raise ValueError("All sinks must be of the same type!") - if isinstance(sinks[0], os.PathLike): - sinks = [os.path.expanduser(s) for s in sinks] - - if isinstance(sinks[0], (io.StringIO, io.BytesIO, io.TextIOBase)): + if initial_sink_cls in {io.StringIO, io.BytesIO, io.TextIOBase}: data_sinks.reserve(len(sinks)) if isinstance(sinks[0], (io.StringIO, io.BytesIO)): for s in sinks: diff --git a/python/cudf/cudf/pylibcudf_tests/io/test_json.py b/python/cudf/cudf/pylibcudf_tests/io/test_json.py index 552baa5ee51..d6b8bfa6976 100644 --- a/python/cudf/cudf/pylibcudf_tests/io/test_json.py +++ b/python/cudf/cudf/pylibcudf_tests/io/test_json.py @@ -14,12 +14,11 @@ def test_write_json_basic(table_data, source_or_sink, lines, rows_per_chunk): plc_table_w_meta, pa_table = table_data sink = source_or_sink - kwargs = {} - if rows_per_chunk <= plc_table_w_meta.tbl.num_rows(): - kwargs["rows_per_chunk"] = rows_per_chunk - plc.io.json.write_json( - plc.io.SinkInfo([sink]), plc_table_w_meta, lines=lines, **kwargs + plc.io.SinkInfo([sink]), + plc_table_w_meta, + lines=lines, + rows_per_chunk=rows_per_chunk, ) exp = pa_table.to_pandas()