Repository navigation
Optimize Dataset.subset() lookups - #686
harshitgupta31415 wants to merge 5 commits into
Conversation
|
Please comment with a benchmark for this. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #686 +/- ##
==========================================
+ Coverage 93.48% 93.92% +0.44%
==========================================
Files 29 30 +1
Lines 3315 3574 +259
==========================================
+ Hits 3099 3357 +258
- Misses 216 217 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Thanks, @hugovk! I benchmarked the base commit ( Square datasetsEach result is the median of five runs.
Rectangular datasetI also measured two rectangular datasets using the median of ten runs:
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 codeimport 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") |
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as resolved.
|
Can you check the behaviour change when |
|
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. |
|
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 |
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.
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. |
|
Follow-up update (78ecabd): Added a follow-up fix for the edge cases. Small datasets and custom inputs keep the original column-lookup path, while normal larger datasets still use cached indexes. The header scan now stops once the needed columns are found, which fixes the wide-table slowdown. All 87 focused subset checks pass, and I’ve updated the benchmarks, including slower cases. |
What changed?
Dataset.subset() was looking up the same rows and columns over and over. This keeps ordinary integer row selections in a set and caches column positions for ordinary, consistent datasets.
The set is only used for membership checks. Rows still follow the source dataset's order, and requested column order, duplicate columns and the first occurrence of duplicate source headers stay the same.
The latest follow-up avoids scanning every header just to check its type. Small datasets keep the original column-lookup loop; larger datasets stop the index scan once every needed first occurrence is found. This fixes the very wide dataset / first-column slowdown without dropping the full-table optimization.
I also kept the original behavior in the compatibility cases I tested:
For an n × n ordinary dataset selecting all rows and columns, the expected work stays O(n²), down from upstream's O(n³). The initial column-selection membership checks are still present. Fallback paths retain the original column-lookup work and can be slower than the earlier PR; these checks are intentional compatibility overhead, not a promise that every input gets faster.
Extra storage is O(k + c) for k selected integer rows and c requested columns, in addition to the output. This does not double the whole dataset, but is not always negligible. An earlier isolated measurement on 64-bit CPython 3.14.5 found that a set of five million existing integer indexes occupied 128 MiB, with a 192 MiB construction peak. That excludes the existing index list, integer objects and dataset, and is interpreter-dependent.
The focused suite has 87 passing tests (83 compatibility tests and four existing subset tests), plus two passing subtests. The full local suite has 268 passing tests and two passing subtests, with the same timezone-dependent timestamp failure as upstream. Ruff and whitespace checks pass. Optional pandas paths and other Python versions were not exercised locally.
Fixes #685.
Fresh benchmarks — all 15 subset workloads, including slower cases
CPython 3.14.5 on Windows. Upstream base: a36c965. This targeted follow-up has 40 workloads in total: 15 subset and 25 filter. Each version gets five timing batches and five separate peak-allocation samples; the table shows medians.
The actual upstream, previous PR and updated method source snapshots run on identical inputs using the otherwise unchanged Dataset/Row implementation. Fixtures, validation and tracing are excluded from timing. One measurement runs at a time, allowed on a rotating pair of physical cores; the methods themselves are not parallelized.
Fixtures are synthetic, not collected personal data. Peak includes the live output but excludes preexisting inputs. It measures additional Python allocations during one call, not total dataset storage or RAM.
These are local observations, not speed guarantees. Thermal state, hybrid cores and OS activity add substantial timing noise, especially on tiny calls; near-1× ratios should not be treated as definite wins or losses.
Against the previous PR (8e98231), the 2 × 10,000 first-column selection went from 349.27 µs to 7.47 µs (upstream: 7.25 µs). The 1,000 × 50 first-column selection went from 59.91 µs to 144.36 µs because of row validation; upstream was 259.43 µs. The no-row selection also uses more temporary mapping storage than the previous PR. So this follow-up removes the wide-table regression, but is not faster than the previous PR in every case.
Follow-up update — 78ecabd
Added a follow-up fix for the edge cases. Small datasets and custom inputs keep the original column-lookup path, while normal larger datasets still use cached indexes. The header scan now stops once the needed columns are found, which fixes the wide-table slowdown. All 87 focused subset checks pass, and I’ve updated the benchmarks, including slower cases.