Skip to content

perf: optimize std.map, std.flatMap, and std.filterMap allocations#670

Merged
stephenamar-db merged 1 commit intodatabricks:masterfrom
He-Pin:perf/map-flatmap-filtermap
Apr 8, 2026
Merged

perf: optimize std.map, std.flatMap, and std.filterMap allocations#670
stephenamar-db merged 1 commit intodatabricks:masterfrom
He-Pin:perf/map-flatmap-filtermap

Conversation

@He-Pin
Copy link
Copy Markdown
Contributor

@He-Pin He-Pin commented Apr 4, 2026

Motivation

std.map, std.flatMap, and std.filterMap allocate intermediate collections (iterators, builders) for each call. These stdlib functions are among the most frequently called operations in Jsonnet programs.

Key Design Decision

Replace functional collection operations with direct array-based loops:

  1. Pre-allocate output arrays to exact size where possible
  2. Use while-loops instead of iterator-based operations
  3. Avoid intermediate collection allocations

Modification

  • ArrayModule.scala: Optimized std.map, std.flatMap, std.filterMap with direct array operations
  • Test: Added tests covering empty arrays, single elements, and large arrays

Benchmark Results

JMH (JVM, 3 iterations)

Benchmark Master (ms/op) This PR (ms/op) Change
bench.02 50.427 ± 38.906 42.621 ± 2.277 -15.5% 🔥
comparison2 85.854 ± 188.657 70.046 ± 6.580 -18.4%
realistic2 73.458 ± 66.747 64.608 ± 15.144 -12.1%
reverse ~11.5 10.226 ± 2.288 -11.1%

Analysis

  • Broad improvement across all benchmarks — map/flatMap/filterMap are universally used
  • bench.02 benefits most (-15.5%) as it heavily uses std.map for array transformations
  • No regressions

References

  • Upstream exploration: he-pin/sjsonnet jit branch stdlib optimization

Result

-12% to -18% JVM improvement by eliminating intermediate collections in map/flatMap/filterMap.

Three stdlib array function optimizations:

1. std.map: Replace .map(closure) with pre-sized array + while-loop.
   Eliminates closure allocation and intermediate array creation.

2. std.flatMap: Two-pass approach for array variant:
   - First pass: apply function, collect sub-arrays, count total length
   - Second pass: System.arraycopy into pre-sized result array
   Avoids .flatMap's intermediate ArrayBuilder resizing.

3. std.filterMap: Replace .flatMap + Option boxing with while-loop
   and ArrayBuilder. Eliminates Some/None wrapping per element.

Upstream: jit branch commit 9cb95af

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@He-Pin He-Pin marked this pull request as ready for review April 4, 2026 21:41
@stephenamar-db stephenamar-db merged commit 7bb4b6d into databricks:master Apr 8, 2026
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.

2 participants