Skip to content

perf: avoid post-pass column reconstruction in NaN preprocessing #6

Merged
CyrilJl merged 6 commits into
mainfrom
perf/trial-05-preprocess-direct-cols
Mar 20, 2026
Merged

perf: avoid post-pass column reconstruction in NaN preprocessing #6
CyrilJl merged 6 commits into
mainfrom
perf/trial-05-preprocess-direct-cols

Conversation

@CyrilJl

@CyrilJl CyrilJl commented Mar 20, 2026

Copy link
Copy Markdown
Owner

Summary

This PR optimizes the preprocessing step used to build the NaN coordinate mapping in OptiMask.

Instead of reconstructing cols_with_nan at the end with np.flatnonzero(...) and an additional reordering pass, the branch records NaN column indices directly
during the main scan. That removes extra array work after preprocessing while preserving the existing output shape and semantics.

What Changed

  • Add a dedicated cols_with_nan buffer during preprocessing
  • Fill that buffer the first time a NaN is seen in a column
  • Slice the populated portion at the end instead of rebuilding the column list from cols_index_mapper

Why

The previous implementation did the column discovery twice:

  • once implicitly while scanning values
  • once again at the end to rebuild ordered column indices

This change keeps the same behavior but avoids that post-processing step, which should reduce preprocessing overhead on inputs with missing values.

Notes

  • Functional change is limited to optimask/_optimask.py

@CyrilJl CyrilJl self-assigned this Mar 20, 2026
@CyrilJl CyrilJl merged commit 0b3a628 into main Mar 20, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant