⚡️ Speed up function _maybe_prepare_times by 14%

[codeflash-ai]

📄 14% (0.14x) speedup for _maybe_prepare_times in xarray/backends/netcdf3.py

⏱️ Runtime : 1.32 milliseconds → 1.15 milliseconds (best of 25 runs)

📝 Explanation and details

The optimized code achieves a 14% speedup through three key optimizations that reduce computational overhead in data processing pipelines:

1. Set-based Lookup Optimization in _is_time_like()

• Replaced the list time_strings with a module-level set _TIME_STRINGS_SET
• Changed any(tstr == units for tstr in time_strings) to units in _TIME_STRINGS_SET
• This converts O(n) linear search to O(1) hash lookup, eliminating the 137.6μs bottleneck visible in line profiling (reduced to 21.9μs)
• The set is created once at module load instead of recreating the list on every function call

2. Precomputed Constants in _maybe_prepare_times()

• Moved np.iinfo(np.int64).min computation to module-level constant _INT64_MIN
• This eliminates repeated expensive numpy introspection calls, reducing the comparison line from 616.7μs to 314.5μs
• Cached var.attrs reference to avoid repeated attribute access

3. Optimized Conditional Logic

• Restructured the fill value retrieval to only call attrs.get("_FillValue", np.nan) when actually needed (inside the mask check)
• This reduces unnecessary dictionary lookups when no replacement is required

Performance Impact Analysis:
The optimizations are particularly effective for:

• Non-time-like units (36-43% faster): Quick rejection via set lookup instead of expensive list iteration
• Time-like units without replacement (23-31% faster): Benefits from precomputed constants and reduced attribute access
• Large arrays (15-28% faster): Constant-time improvements scale well with data size

Hot Path Context:
Given that _maybe_prepare_times() is called from encode_nc3_variable() in NetCDF encoding workflows, these micro-optimizations compound significantly when processing large datasets or many variables. The function processes integer arrays to handle sentinel values, making it critical in data serialization pipelines where every millisecond matters.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np

# imports
import pytest  # used for unit tests
from xarray.backends.netcdf3 import _maybe_prepare_times


# Helper class to mimic xarray.Variable-like object
class DummyVar:
    def __init__(self, data, attrs=None):
        self.data = data
        self.attrs = attrs or {}


# -------------------------
# Unit tests for _maybe_prepare_times
# -------------------------

# 1. Basic Test Cases


def test_basic_int64_time_like_with_min_replacement():
    # Data contains np.iinfo(np.int64).min, units is time-like, _FillValue is set
    arr = np.array([1, 2, np.iinfo(np.int64).min, 4], dtype=np.int64)
    fill = -9999
    var = DummyVar(arr, {"units": "days", "_FillValue": fill})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 25.7μs -> 23.0μs (11.8% faster)


def test_basic_int64_time_like_with_nan_replacement():
    # Data contains np.iinfo(np.int64).min, units is time-like, _FillValue is not set
    arr = np.array([np.iinfo(np.int64).min, 2], dtype=np.int64)
    var = DummyVar(arr, {"units": "hours"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 26.9μs -> 23.4μs (14.9% faster)


def test_basic_int64_non_time_like():
    # Data contains np.iinfo(np.int64).min, units is not time-like
    arr = np.array([np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, {"units": "meters"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.68μs -> 2.57μs (43.1% faster)


def test_basic_float_data():
    # Data is float, no replacement should occur even if units is time-like
    arr = np.array([1.0, np.nan, 3.0])
    var = DummyVar(arr, {"units": "days"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 968ns -> 890ns (8.76% faster)


def test_basic_uint_data():
    # Data is unsigned int, units is time-like, np.iinfo(np.int64).min cannot appear in uint
    arr = np.array([1, 2, 3], dtype=np.uint64)
    var = DummyVar(arr, {"units": "seconds"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 25.3μs -> 19.2μs (31.5% faster)


def test_basic_no_units():
    # Data is int, units missing
    arr = np.array([1, 2, np.iinfo(np.int64).min], dtype=np.int64)
    var = DummyVar(arr, {})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 1.02μs -> 1.07μs (4.93% slower)


def test_basic_units_none():
    # Data is int, units explicitly None
    arr = np.array([1, 2, np.iinfo(np.int64).min], dtype=np.int64)
    var = DummyVar(arr, {"units": None})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 1.02μs -> 1.08μs (5.01% slower)


# 2. Edge Test Cases


def test_edge_empty_array():
    # Empty array, nothing to replace
    arr = np.array([], dtype=np.int64)
    var = DummyVar(arr, {"units": "days"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 22.8μs -> 18.1μs (26.2% faster)


def test_edge_all_min_values():
    # All values are np.iinfo(np.int64).min
    arr = np.full(5, np.iinfo(np.int64).min, dtype=np.int64)
    fill = 12345
    var = DummyVar(arr, {"units": "hours", "_FillValue": fill})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 23.6μs -> 20.2μs (17.1% faster)


def test_edge_mixed_dtype_int32():
    # Data is int32, units is time-like, np.iinfo(np.int64).min cannot appear in int32
    arr = np.array([1, 2, -2147483648], dtype=np.int32)
    var = DummyVar(arr, {"units": "days"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 23.9μs -> 18.6μs (28.7% faster)


def test_edge_units_with_since_valid():
    # Data contains np.iinfo(np.int64).min, units is "days since 2000-01-01"
    arr = np.array([np.iinfo(np.int64).min, 123], dtype=np.int64)
    var = DummyVar(arr, {"units": "days since 2000-01-01", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 33.7μs -> 31.6μs (6.68% faster)


def test_edge_units_with_since_invalid():
    # Data contains np.iinfo(np.int64).min, units is "nonsense since 2000-01-01"
    arr = np.array([np.iinfo(np.int64).min, 123], dtype=np.int64)
    var = DummyVar(arr, {"units": "nonsense since 2000-01-01"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 33.8μs -> 32.9μs (2.78% faster)


def test_edge_units_as_bytes():
    # Units provided as bytes
    arr = np.array([np.iinfo(np.int64).min, 5], dtype=np.int64)
    var = DummyVar(arr, {"units": b"days", "_FillValue": 0})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 4.07μs -> 2.93μs (38.8% faster)


def test_edge_fillvalue_is_nan():
    # _FillValue is np.nan
    arr = np.array([np.iinfo(np.int64).min, 1], dtype=np.int64)
    var = DummyVar(arr, {"units": "days", "_FillValue": np.nan})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 27.2μs -> 24.1μs (12.8% faster)


def test_edge_fillvalue_is_none():
    # _FillValue is None, should fallback to np.nan
    arr = np.array([np.iinfo(np.int64).min, 1], dtype=np.int64)
    var = DummyVar(arr, {"units": "days", "_FillValue": None})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 27.3μs -> 24.2μs (12.7% faster)


def test_edge_units_case_sensitivity():
    # Units is "Days" (capital D), should not match
    arr = np.array([np.iinfo(np.int64).min, 1], dtype=np.int64)
    var = DummyVar(arr, {"units": "Days", "_FillValue": 999})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.56μs -> 2.58μs (37.7% faster)


def test_edge_units_with_extra_text():
    # Units is "days accumulated", should not match
    arr = np.array([np.iinfo(np.int64).min, 1], dtype=np.int64)
    var = DummyVar(arr, {"units": "days accumulated", "_FillValue": 999})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.44μs -> 2.53μs (36.1% faster)


def test_edge_units_with_since_and_extra_whitespace():
    # Units is "days   since   2000-01-01", should match
    arr = np.array([np.iinfo(np.int64).min, 2], dtype=np.int64)
    var = DummyVar(arr, {"units": "days   since   2000-01-01", "_FillValue": 0})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 33.7μs -> 32.1μs (5.06% faster)


def test_edge_units_with_since_and_invalid_date():
    # Units is "days since nonsense-date", should still match as regex doesn't validate date
    arr = np.array([np.iinfo(np.int64).min, 2], dtype=np.int64)
    var = DummyVar(arr, {"units": "days since nonsense-date", "_FillValue": 0})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 12.7μs -> 12.8μs (0.431% slower)


# 3. Large Scale Test Cases


def test_large_scale_no_min_values():
    # Large array, no np.iinfo(np.int64).min present
    arr = np.arange(1000, dtype=np.int64)
    var = DummyVar(arr, {"units": "days"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 25.1μs -> 19.8μs (26.4% faster)


def test_large_scale_some_min_values():
    # Large array, some np.iinfo(np.int64).min present
    arr = np.arange(1000, dtype=np.int64)
    arr[::100] = np.iinfo(np.int64).min  # every 100th value is min
    fill = -9999
    var = DummyVar(arr, {"units": "days", "_FillValue": fill})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 25.7μs -> 23.0μs (11.8% faster)
    for i in range(0, 1000, 100):
        pass


def test_large_scale_all_min_values():
    # Large array, all values are np.iinfo(np.int64).min
    arr = np.full(1000, np.iinfo(np.int64).min, dtype=np.int64)
    fill = 123456
    var = DummyVar(arr, {"units": "days", "_FillValue": fill})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 24.3μs -> 21.4μs (13.3% faster)


def test_large_scale_non_time_like():
    # Large array, units is not time-like
    arr = np.full(1000, np.iinfo(np.int64).min, dtype=np.int64)
    var = DummyVar(arr, {"units": "meters", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.69μs -> 2.61μs (41.5% faster)


def test_large_scale_float_data():
    # Large float array, units is time-like
    arr = np.random.rand(1000)
    var = DummyVar(arr, {"units": "days"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 953ns -> 1.01μs (5.46% slower)


def test_large_scale_uint_data():
    # Large unsigned int array, units is time-like
    arr = np.arange(1000, dtype=np.uint64)
    var = DummyVar(arr, {"units": "days"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 26.3μs -> 20.4μs (28.7% faster)


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

import numpy as np

# imports
import pytest  # used for our unit tests
from xarray.backends.netcdf3 import _maybe_prepare_times


# Helper class to mimic xarray.Variable interface
class DummyVar:
    def __init__(self, data, attrs=None):
        self.data = data
        self.attrs = attrs or {}


# unit tests

# ------------------- Basic Test Cases -------------------


def test_basic_int_time_like_with_fillvalue():
    # Test replacement of np.iinfo(np.int64).min with _FillValue for time-like units
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": -9999})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 25.9μs -> 22.9μs (12.7% faster)


def test_basic_int_time_like_without_fillvalue():
    # Should replace np.iinfo(np.int64).min with np.nan if _FillValue is not given
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "hours"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 26.9μs -> 24.0μs (11.9% faster)


def test_basic_int_not_time_like():
    # Should not replace anything if units are not time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "meters"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.56μs -> 2.60μs (36.9% faster)


def test_basic_float_time_like():
    # Should not replace anything for float dtype
    arr = np.array([1.0, float(np.iinfo(np.int64).min), 3.0], dtype=np.float64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": -9999})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 811ns -> 860ns (5.70% slower)


def test_basic_no_units():
    # Should not replace anything if units are missing
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 1.08μs -> 1.10μs (1.72% slower)


def test_basic_units_none():
    # Should not replace anything if units is None
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": None})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 973ns -> 1.02μs (5.07% slower)


def test_basic_units_string_not_time_like():
    # Should not replace anything if units string is not time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "banana"})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.74μs -> 2.66μs (40.2% faster)


def test_basic_units_time_like_with_since():
    # Should detect "days since 2000-01-01" as time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days since 2000-01-01", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 33.9μs -> 32.8μs (3.47% faster)


def test_basic_units_invalid_since():
    # Should not treat invalid "since" units as time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days since", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 7.89μs -> 8.44μs (6.50% slower)


# ------------------- Edge Test Cases -------------------


def test_edge_all_min_values():
    # All values are np.iinfo(np.int64).min
    arr = np.full(5, np.iinfo(np.int64).min, dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "minutes", "_FillValue": 999})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 23.8μs -> 20.9μs (13.7% faster)


def test_edge_no_min_values():
    # No values are np.iinfo(np.int64).min
    arr = np.array([10, 20, 30], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "seconds", "_FillValue": 0})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 22.7μs -> 17.8μs (27.2% faster)


def test_edge_empty_array():
    # Empty array should be handled gracefully
    arr = np.array([], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 22.1μs -> 17.1μs (29.4% faster)


def test_edge_non_integer_dtype():
    # Non-integer dtype (e.g., string) should not be processed
    arr = np.array(["a", "b", "c"])
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 1.03μs -> 985ns (4.47% faster)


def test_edge_units_with_extra_whitespace():
    # Units with extra whitespace should still be detected as time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "  days  ", "_FillValue": -5})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.69μs -> 2.65μs (38.9% faster)


def test_edge_units_with_since_and_extra_whitespace():
    # Units with "since" and extra whitespace
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "hours since   1999-12-31", "_FillValue": 123})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 34.0μs -> 32.2μs (5.78% faster)


def test_edge_fillvalue_is_none():
    # _FillValue is explicitly set to None, should use np.nan
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": None})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 26.9μs -> 24.6μs (9.28% faster)


def test_edge_fillvalue_is_nan():
    # _FillValue is np.nan, should use np.nan
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": np.nan})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 26.4μs -> 23.4μs (12.7% faster)


def test_edge_units_is_integer():
    # Units as integer should not be time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": 123, "_FillValue": 0})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.66μs -> 2.72μs (34.5% faster)


def test_edge_units_is_empty_string():
    # Units as empty string should not be time-like
    arr = np.array([1, np.iinfo(np.int64).min, 3], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "", "_FillValue": 0})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.44μs -> 2.55μs (34.9% faster)


# ------------------- Large Scale Test Cases -------------------


def test_large_scale_all_min_values():
    # Large array, all values are np.iinfo(np.int64).min
    arr = np.full(1000, np.iinfo(np.int64).min, dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "seconds", "_FillValue": 111})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 25.0μs -> 21.7μs (15.1% faster)


def test_large_scale_no_min_values():
    # Large array, no values are np.iinfo(np.int64).min
    arr = np.arange(1000, dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "minutes", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 24.1μs -> 19.5μs (23.7% faster)


def test_large_scale_some_min_values():
    # Large array, some values are np.iinfo(np.int64).min
    arr = np.arange(1000, dtype=np.int64)
    arr[100] = np.iinfo(np.int64).min
    arr[999] = np.iinfo(np.int64).min
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": 42})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 26.4μs -> 24.2μs (9.26% faster)


def test_large_scale_non_time_like():
    # Large array, units not time-like
    arr = np.full(1000, np.iinfo(np.int64).min, dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "kilograms", "_FillValue": 999})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 3.81μs -> 2.79μs (36.6% faster)


def test_large_scale_float_dtype():
    # Large array, float dtype
    arr = np.full(1000, float(np.iinfo(np.int64).min), dtype=np.float64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 1.27μs -> 865ns (46.8% faster)


def test_large_scale_empty_array():
    # Large scale edge: empty array
    arr = np.array([], dtype=np.int64)
    var = DummyVar(arr, attrs={"units": "days", "_FillValue": -1})
    codeflash_output = _maybe_prepare_times(var)
    result = codeflash_output  # 23.2μs -> 18.2μs (27.6% faster)


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

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_xarrayteststest_concat_py_xarrayteststest_computation_py_xarrayteststest_formatting_py_xarray__replay_test_0.py::test_xarray_backends_netcdf3__maybe_prepare_times	504μs	443μs	13.7%✅

To edit these changes git checkout codeflash/optimize-_maybe_prepare_times-mi9r5k4c and push.