Optimize Dataset.subset() lookups

[harshitgupta31415]

What changed?

Dataset.subset() was repeatedly looking up the same rows and columns while creating a subset. This made it noticeably slower for larger datasets.

This change:

• remembers the selected rows for quicker checks
• finds each selected column's position only once
• processes only the rows that are needed

For a 400×400 dataset, the execution time improved from about 1.03 seconds to 0.03 seconds.

A test was added to make sure the existing selection behavior remains unchanged. All 184 tests pass.

Fixes #685

[hugovk]

Please comment with a benchmark for this.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.56%. Comparing base (a36c965) to head (bf2a307).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #686      +/-   ##
==========================================
+ Coverage   93.48%   93.56%   +0.08%     
==========================================
  Files          29       29              
  Lines        3315     3327      +12     
==========================================
+ Hits         3099     3113      +14     
+ Misses        216      214       -2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[harshitgupta31415]

Thanks, @hugovk! I benchmarked the base commit (a36c965) against this PR (98b9bd1) using Python 3.14.5 on Windows. Dataset construction was completed before timing, so the results measure only Dataset.subset().

Square datasets

Each result is the median of five runs.

Dataset size	Base	This PR	Speedup
100×100	0.017916 s	0.002422 s	7.4×
200×200	0.116337 s	0.010250 s	11.3×
400×400	0.973558 s	0.024506 s	39.7×

Rectangular dataset

I also measured two rectangular datasets using the median of ten runs:

Dataset size	Base	This PR	Speedup
4,000×50	0.235769 s	0.026538 s	8.9×
400×4,092	106.261264 s	0.344110 s	308.8×

For an n×n selection, the change reduces runtime from approximately O(n³) to O(n²). It uses a row set and a precomputed column-index list, but auxiliary space remains O(n); the O(n²) returned dataset still dominates total memory usage.

Benchmark code

import platform
from statistics import median
from timeit import repeat

from tablib import Dataset


print(f"Python {platform.python_version()}")

cases = (
    (100, 100, 5),
    (200, 200, 5),
    (400, 400, 5),
    (4_000, 50, 10),
    (400, 4_092, 10),
)

for row_count, column_count, repetitions in cases:
    headers = [f"column_{index}" for index in range(column_count)]
    row = tuple(range(column_count))
    dataset = Dataset(*(row for _ in range(row_count)), headers=headers)

    elapsed = median(
        repeat(
            lambda: dataset.subset(
                rows=range(row_count),
                cols=headers,
            ),
            repeat=repetitions,
            number=1,
        )
    )

    print(f"{row_count} rows x {column_count} columns: {elapsed:.6f}s")

[hugovk]

Can you check the behaviour change when cols=[] or every requested column is missing?

[harshitgupta31415]

Good catch, @hugovk this was changing the behavior. I’ve fixed it so both cols=[] and an all-missing column list behave the same as old again. Added tests for both cases too.

[harshitgupta31415]

Hey @hugovk, just following up when you have a moment. I fixed the empty and all-missing column cases, then pushed 4357722 for another edge case: main accepts an unhashable row selector that compares equal to a row number. The PR now uses a set for plain integer row numbers and keeps the original list check for other values. For an n×n selection, the runtime is still O(n²).

The result still walks source rows in order, so repeated row numbers don't duplicate or reorder output. Requested columns keep their order and duplicates, missing names are ignored, and repeated source headers still resolve to their first match. On a nonempty dataset, no matching columns still raises the same TypeError as main. I added regression tests for the edge cases, and the current checks are green. The benchmark is above.

I left out the old else: raise KeyError because the earlier column filtering makes that branch unreachable for normal inputs. Checking again for every cell would bring back the repeated lookup. Could you review when you get a chance? Thanks!

[harshitgupta31415]

Benchmark after the row-selection fix (4357722)

I reran the same benchmark on the updated PR with Python 3.14.5 on Windows. Dataset construction was outside the timed part. The base values below are from my earlier benchmark of a36c965; that commit has not changed. The updated PR values were measured now. Square cases use the median of five runs, and rectangular cases use the median of ten.

Dataset size	Base (earlier run)	Updated PR	Speedup vs base
100×100	0.017916 s	0.001694 s	10.6×
200×200	0.116337 s	0.005802 s	20.1×
400×400	0.973558 s	0.022345 s	43.6×
4,000×50	0.235769 s	0.039248 s	6.0×
400×4,092	106.261264 s	0.825082 s	128.8×

The widest case varied across runs: a separate ten-run batch on the updated PR had a 0.487386 s median. So its exact speedup is approximate, but it remains much faster than base. The fast set lookup is still used for ordinary integer row selections.