⚡️ Speed up method ArrowParserWrapper._get_pyarrow_options by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for ArrowParserWrapper._get_pyarrow_options in pandas/io/parsers/arrow_parser_wrapper.py

⏱️ Runtime : 136 microseconds → 123 microseconds (best of 16 runs)

📝 Explanation and details

The optimized code achieves a 10% speedup by replacing inefficient dictionary comprehensions with direct assignments and eliminating redundant dictionary lookups.

Key optimizations:

1. Eliminated dict comprehension overhead for parse_options: Instead of creating a dictionary comprehension that iterates through all self.kwds.items() and filters by option names, the code now uses direct get() calls for the 4 specific options and conditional assignments. This avoids the overhead of creating intermediate tuples and filtering logic.

2. Reduced redundant lookups in mapping loop: Changed if pandas_name in self.kwds and self.kwds.get(pandas_name) is not None to option_value = self.kwds.get(pandas_name); if option_value is not None, eliminating the double dictionary lookup for each key.

3. Replaced dict comprehension for convert_options: Similar to parse_options, replaced the comprehension that scans all kwds with direct assignments for the 6 specific option names, avoiding iteration overhead.

4. Optimized strings_can_be_null logic: Added a null check before the membership test "" in null_values to avoid potential exceptions and make the logic more explicit.

The optimizations are particularly effective for the test cases with many options (20-43% speedup) because they eliminate the O(n) dictionary iterations in favor of O(1) direct lookups. Even with smaller option sets, the reduced function call overhead provides consistent 10-20% improvements. These gains are meaningful since this function is likely called during CSV parsing initialization, where every microsecond counts for data processing workflows.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 22 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	60.7%

🌀 Generated Regression Tests and Runtime

# imports
from pandas.io.parsers.arrow_parser_wrapper import ArrowParserWrapper


# Mocks and minimal dependencies to support the function
class ParserWarning(Warning):
    pass


def find_stack_level():
    return 1


class ParserBase:
    class BadLineHandleMethod:
        ERROR = "error"
        WARN = "warn"
        SKIP = "skip"

    def __init__(self, kwds):
        self.header = kwds.get("header", 0)
        self.date_format = kwds.pop("date_format", None)
        self.na_values = kwds.get("na_values", [""])
        self.kwds = kwds


# -------------------- UNIT TESTS --------------------

# Basic Test Cases


def test_basic_column_mapping():
    # Test that basic column renaming works
    kwds = {
        "usecols": ["a", "b"],
        "na_values": ["", "NA"],
        "escapechar": "\\",
        "skip_blank_lines": True,
        "decimal": ".",
        "quotechar": '"',
        "delimiter": ",",
        "header": 0,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 6.83μs -> 5.62μs (21.7% faster)
    # Should not have the original pandas keys anymore
    for k in [
        "usecols",
        "na_values",
        "escapechar",
        "skip_blank_lines",
        "decimal",
        "quotechar",
    ]:
        pass


def test_header_none_include_columns_prefix():
    # If header is None, include_columns values are prefixed with 'f'
    kwds = {
        "usecols": [0, 1, 2],
        "na_values": [""],
        "header": None,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()


def test_include_columns_empty_list():
    # usecols as empty list
    kwds = {
        "usecols": [],
        "na_values": [""],
        "header": 0,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 5.92μs -> 5.19μs (14.2% faster)


# Large Scale Test Cases


def test_large_number_of_columns():
    # Large number of columns in usecols
    cols = [f"col{i}" for i in range(500)]
    kwds = {
        "usecols": cols,
        "na_values": ["", "NA"],
        "header": 0,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 5.37μs -> 4.77μs (12.5% faster)


def test_large_include_columns_with_header_none():
    # Large include_columns and header None triggers prefixing
    cols = list(range(500))
    kwds = {
        "usecols": cols,
        "na_values": [""],
        "header": None,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()


def test_performance_many_options():
    # Many options set at once, should not error or be slow
    kwds = {
        "usecols": [f"col{i}" for i in range(100)],
        "na_values": [str(i) for i in range(100)],
        "escapechar": "\\",
        "skip_blank_lines": True,
        "decimal": ".",
        "quotechar": '"',
        "delimiter": ",",
        "header": None,
        "skiprows": 10,
        "encoding": "utf-8",
        "true_values": ["yes", "y"],
        "false_values": ["no", "n"],
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 13.1μs -> 10.9μs (20.4% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

# imports
from pandas.io.parsers.arrow_parser_wrapper import ArrowParserWrapper


# Minimal stubs for dependencies (to avoid importing pandas internals)
class ParserWarning(UserWarning):
    pass


def find_stack_level():
    return 1


class ParserBase:
    class BadLineHandleMethod:
        ERROR = "error"
        WARN = "warn"
        SKIP = "skip"

    def __init__(self, kwds):
        self.header = kwds.get("header")
        self.date_format = kwds.get("date_format", None)


# ===== UNIT TESTS =====

# -------- BASIC TEST CASES --------


def test_basic_column_mapping_and_options():
    """
    Test that basic pandas to pyarrow option mapping works,
    and that parse/convert/read options are set correctly.
    """
    kwds = {
        "usecols": [0, 1],
        "na_values": ["", "NA"],
        "escapechar": "\\",
        "skip_blank_lines": True,
        "decimal": ".",
        "quotechar": '"',
        "skiprows": 0,
        "header": 0,
        "date_format": "%Y-%m-%d",
        "true_values": ["yes"],
        "false_values": ["no"],
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 6.86μs -> 4.77μs (43.9% faster)


def test_include_columns_autogen_header():
    """
    Test that include_columns is prefixed with 'f' when header is None.
    """
    kwds = {
        "usecols": [0, 2, 4],
        "na_values": ["NA"],
        "skiprows": 1,
        "header": None,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 7.15μs -> 6.10μs (17.3% faster)


def test_large_number_of_columns():
    """
    Test with a large list of usecols and na_values.
    """
    usecols = list(range(1000))
    na_values = [str(i) for i in range(1000)]
    kwds = {
        "usecols": usecols,
        "na_values": na_values,
        "skiprows": 0,
        "header": 0,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 10.1μs -> 9.17μs (9.66% faster)


def test_large_autogen_column_names():
    """
    Test autogeneration of column names with large include_columns.
    """
    usecols = list(range(1000))
    kwds = {
        "usecols": usecols,
        "na_values": ["NA"],
        "skiprows": 0,
        "header": None,
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 70.7μs -> 68.7μs (2.82% faster)


def test_performance_large_options():
    """
    Sanity check that large option dicts do not cause errors or slowdowns.
    """
    # This test is not a timed benchmark, but should not error or hang.
    kwds = {
        "usecols": list(range(500)),
        "na_values": [str(i) for i in range(500)],
        "skiprows": 0,
        "header": 0,
        "true_values": [str(i) for i in range(100)],
        "false_values": [str(i) for i in range(100, 200)],
        "escapechar": "\\",
        "skip_blank_lines": True,
        "decimal": ".",
        "quotechar": '"',
    }
    wrapper = ArrowParserWrapper("dummy", **kwds)
    wrapper._get_pyarrow_options()  # 9.69μs -> 7.34μs (32.0% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-ArrowParserWrapper._get_pyarrow_options-miy54msj and push.