Fix for merge_offset_ranges becoming quadratic

docs/source/changelog.rst

@@ -12,6 +12,7 @@ Enhancements
 
 Fixes
 
+- Make ``merge_offset_ranges`` `O(n log n)` and never emit overlapping ranges, replacing the quadratic nested-range filter added in #1982 (#2091)
 - Fix incorrect glob docstring for '[!]' (#2084)
 - Propagate storage_options to all backends resolved by GenericFileSystem (#2083)
 - Handle end=None in FirstChunkCache._fetch like the other caches (#2082)

fsspec/tests/test_utils.py

@@ -1,5 +1,7 @@
 import io
+import random
 import sys
+import time
 from pathlib import Path, PurePath
 from unittest.mock import Mock
 
@@ -450,6 +452,124 @@ def test_merge_offset_ranges(max_gap, max_block):
     assert expect_ends == result_ends
 
 
+def test_merge_offset_ranges_drops_nested():
+    paths = ["f", "f", "f", "f", "g"]
+    starts = [0, 10, 0, 100, 0]
+    ends = [50, 20, 80, 150, 10]
+
+    result_paths, result_starts, result_ends = merge_offset_ranges(
+        paths, starts, ends, max_gap=0, max_block=None
+    )
+
+    assert result_paths == ["f", "f", "g"]
+    assert result_starts == [0, 100, 0]
+    assert result_ends == [80, 150, 10]
+
+
+@pytest.mark.parametrize("sort", [True, False])
+@pytest.mark.parametrize(
+    "starts,ends,expected",
+    [
+        # Exact duplicates: keep one
+        ([0, 0], [50, 50], [(0, 50)]),
+        # Nested range
+        ([0, 10], [80, 20], [(0, 80)]),
+        # None end covers to EOF
+        ([0, 0], [None, 50], [(0, None)]),
+        ([0, 0], [50, None], [(0, None)]),
+        # None end past max_block: keep separate
+        ([0, 50], [10, None], [(0, 10), (50, None)]),
+    ],
+)
+def test_merge_offset_ranges_edges(starts, ends, expected, sort):
+    result = merge_offset_ranges(
+        ["f"] * len(starts),
+        list(starts),
+        list(ends),
+        max_gap=100,
+        max_block=1000,
+        sort=sort,
+    )
+
+    assert list(zip(*result)) == [("f", *rng) for rng in expected]
+
+
+def test_merge_offset_ranges_sorts_unsorted_nested_ranges():
+    result = merge_offset_ranges(
+        ["f", "f"],
+        [10, 0],
+        [20, 80],
+        max_gap=100,
+        max_block=1000,
+        sort=True,
+    )
+
+    assert list(zip(*result)) == [("f", 0, 80)]
+
+
+@pytest.mark.parametrize("max_block", [None, 4, 128])
+def test_merge_offset_ranges_never_overlap(max_block):
+    # Overlaps must merge even past max_block
+    paths = ["f"] * 3
+    starts = [0, 8, 12]
+    ends = [10, 40, 20]
+
+    result_paths, result_starts, result_ends = merge_offset_ranges(
+        paths, starts, ends, max_gap=0, max_block=max_block
+    )
+
+    assert result_paths == ["f"]
+    assert result_starts == [0]
+    assert result_ends == [40]
+
+
+def test_merge_offset_ranges_covers_every_input_range():
+    # Output ranges must not overlap; every input must be covered
+    rand = random.Random(42)
+    paths, starts, ends = [], [], []
+    for _ in range(200):
+        path = f"f{rand.randint(0, 2)}"
+        start = rand.randrange(0, 4000)
+        paths.append(path)
+        starts.append(start)
+        ends.append(start + rand.randrange(1, 500))
+
+    result = merge_offset_ranges(paths, starts, ends, max_gap=8, max_block=512)
+
+    blocks = sorted(zip(*result))
+    for (path1, _, end1), (path2, start2, _) in zip(blocks, blocks[1:]):
+        assert path1 != path2 or start2 >= end1
+
+    for path, start, end in zip(paths, starts, ends):
+        assert any(
+            path == block_path and block_start <= start and end <= block_end
+            for block_path, block_start, block_end in blocks
+        )
+
+
+def test_merge_offset_ranges_many_sequential_is_fast():
+    # Regression: O(n²) nested-range filter added in #1982
+    def run(n):
+        paths = ["file"] * n
+        starts = list(range(0, n * 240, 240))
+        ends = [s + 240 for s in starts]
+        t0 = time.perf_counter()
+        result = merge_offset_ranges(
+            paths, starts, ends, max_block=8_388_608, sort=True
+        )
+        return time.perf_counter() - t0, result, ends[-1]
+
+    large_elapsed, result, expected_end = run(20_000)
+    result_paths, result_starts, result_ends = result
+
+    # Generous ceiling: the linear implementation needs ~10ms here, while the
+    # quadratic one needs tens of seconds.
+    assert large_elapsed < 1.0
+    assert result_paths == ["file"]
+    assert result_starts == [0]
+    assert result_ends == [expected_end]
+
+
 def test_size():
     f = io.BytesIO(b"hello")
     assert fsspec.utils.file_size(f) == 5

fsspec/utils.py

@@ -537,18 +537,21 @@ def nullcontext(obj: T) -> Iterator[T]:
 
 def merge_offset_ranges(
     paths: list[str],
-    starts: list[int] | int,
-    ends: list[int] | int,
+    starts: list[int | None] | int | None,
+    ends: list[int | None] | int | None,
     max_gap: int = 0,
     max_block: int | None = None,
     sort: bool = True,
-) -> tuple[list[str], list[int], list[int]]:
+) -> tuple[list[str], list[int], list[int | None]]:
     """Merge adjacent byte-offset ranges when the inter-range
     gap is <= `max_gap`, and when the merged byte range does not
-    exceed `max_block` (if specified). By default, this function
-    will re-order the input paths and byte ranges to ensure sorted
-    order. If the user can guarantee that the inputs are already
-    sorted, passing `sort=False` will skip the re-ordering.
+    exceed `max_block` (if specified). Overlapping ranges are always
+    merged, even past `max_block`, so returned ranges never overlap
+    when the input is sorted. An `end` of `None` means to the end of
+    the file. By default, this function will re-order the input paths
+    and byte ranges to ensure sorted order. If the user can guarantee
+    that the inputs are already sorted, passing `sort=False` will skip
+    the re-ordering.
     """
     # Check input
     if not isinstance(paths, list):
@@ -560,59 +563,57 @@ def merge_offset_ranges(
     if len(starts) != len(paths) or len(ends) != len(paths):
         raise ValueError
 
+    starts_i: list[int] = [s or 0 for s in starts]
+    ends_i: list[int | None] = ends
+
     # Early Return
-    if len(starts) <= 1:
-        return paths, starts, ends
+    if len(starts_i) <= 1:
+        return paths, starts_i, ends_i
 
-    starts = [s or 0 for s in starts]
     # Sort by paths and then ranges if `sort=True`
     if sort:
-        paths, starts, ends = (
-            list(v)
-            for v in zip(
-                *sorted(
-                    zip(paths, starts, ends),
-                )
-            )
+        ranges = sorted(
+            zip(paths, starts_i, ends_i),
+            # None end sorts last (covers furthest into the file)
+            key=lambda pse: (pse[0], pse[1], math.inf if pse[2] is None else pse[2]),
         )
-    remove = []
-    for i, (path, start, end) in enumerate(zip(paths, starts, ends)):
-        if any(
-            e is not None and p == path and start >= s and end <= e and i != i2
-            for i2, (p, s, e) in enumerate(zip(paths, starts, ends))
+        paths = [r[0] for r in ranges]
+        starts_i = [r[1] for r in ranges]
+        ends_i = [r[2] for r in ranges]
+
+    # Loop through the coupled `paths`, `starts`, and
+    # `ends`, and merge adjacent blocks when appropriate
+    new_paths = paths[:1]
+    new_starts = starts_i[:1]
+    new_ends = ends_i[:1]
+    for path, start, end in zip(paths[1:], starts_i[1:], ends_i[1:]):
+        prev_end = new_ends[-1]
+        if path != new_paths[-1]:
+            # Cannot merge with previous block
+            new_paths.append(path)
+            new_starts.append(start)
+            new_ends.append(end)
+        elif prev_end is None:
+            # Previous block already covers the rest of the file
+            continue
+        elif start < prev_end:
+            # Overlap / nested: merge into the union even past max_block
+            new_starts[-1] = min(new_starts[-1], start)
+            if end is None or end > prev_end:
+                new_ends[-1] = end
+        elif (start - prev_end) > max_gap or (
+            max_block is not None
+            and (end is None or (end - new_starts[-1]) > max_block)
         ):
-            remove.append(i)
-    paths = [p for i, p in enumerate(paths) if i not in remove]
-    starts = [s for i, s in enumerate(starts) if i not in remove]
-    ends = [e for i, e in enumerate(ends) if i not in remove]
-
-    if paths:
-        # Loop through the coupled `paths`, `starts`, and
-        # `ends`, and merge adjacent blocks when appropriate
-        new_paths = paths[:1]
-        new_starts = starts[:1]
-        new_ends = ends[:1]
-        for i in range(1, len(paths)):
-            if paths[i] == paths[i - 1] and new_ends[-1] is None:
-                continue
-            elif (
-                paths[i] != paths[i - 1]
-                or ((starts[i] - new_ends[-1]) > max_gap)
-                or (max_block is not None and (ends[i] - new_starts[-1]) > max_block)
-            ):
-                # Cannot merge with previous block.
-                # Add new `paths`, `starts`, and `ends` elements
-                new_paths.append(paths[i])
-                new_starts.append(starts[i])
-                new_ends.append(ends[i])
-            else:
-                # Merge with the previous block by updating the
-                # last element of `ends`
-                new_ends[-1] = ends[i]
-        return new_paths, new_starts, new_ends
+            # Gap too large, or merging would exceed `max_block`
+            new_paths.append(path)
+            new_starts.append(start)
+            new_ends.append(end)
+        else:
+            # Merge with the previous block
+            new_ends[-1] = end
 
-    # `paths` is empty. Just return input lists
-    return paths, starts, ends
+    return new_paths, new_starts, new_ends
 
 
 def file_size(filelike: IO[bytes]) -> int: