perf: cheapen Index construction and slicing

[ikrommyd]

A few small things in Index, which gets constructed everywhere: module-level numpy instead of Numpy.instance() per call, a dtype to class dict instead of constructing five dtypes to compare against, and __getitem__ skips normalize_slice and the _metadata dict when the data is known.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.89%. Comparing base (60390ac) to head (8d21bcf).

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/index.py	93.93% <100.00%> (+1.43%)	⬆️

... and 1 file with indirect coverage changes

[github-actions]

The documentation preview is ready to be viewed at https://awkward-array.org/doc/pr/4340/

[ikrommyd]

🤖 AI text below 🤖

CPU timings, main vs this branch, median of 7 repeats per case.

case	base (us)	branch (us)	change
array.layout.offsets[:] n=10	2.59	1.40	-45.8%
array.layout.offsets[:] n=100000	2.59	1.41	-45.8%
array.layout.offsets[:] n=1000	2.64	1.43	-45.7%
ak.num(array) n=10	26.52	22.75	-14.2%
ak.Array(layout) n=10	2.34	2.35	+0.4%

Median across all cases: -8.2% (27 cases total, showing the four best and the worst).

benchmark script

"""Benchmark for PR #4340 -- cheapen Index construction and slicing."""

from __future__ import annotations

import statistics
import timeit

import numpy as np

import awkward as ak


def report(label, stmt, glb, *, number, repeats=7):
    ts = timeit.repeat(stmt, number=number, repeat=repeats, globals=glb)
    us = sorted(t / number * 1e6 for t in ts)
    print(
        f"{label:<45} {statistics.median(us):10.3f} us  min {us[0]:10.3f}  sd {statistics.stdev(us):8.3f}  (n={number}x{repeats})"
    )


print(f"awkward {ak.__version__} from {ak.__file__}")

SIZES = [(10, 3000), (1_000, 1000), (100_000, 50)]

for n, number in SIZES:
    offsets = np.arange(0, 3 * n + 1, 3, dtype=np.int64)
    content = np.arange(3 * n, dtype=np.float64)
    layout = ak.contents.ListOffsetArray(
        ak.index.Index64(offsets), ak.contents.NumpyArray(content)
    )
    array = ak.Array(layout)
    glb = {
        "ak": ak,
        "np": np,
        "offsets": offsets,
        "content": content,
        "layout": layout,
        "array": array,
    }

    print(f"\n--- {n} sublists of 3 ---")
    report(
        f"build ListOffsetArray+Array   n={n}",
        "ak.Array(ak.contents.ListOffsetArray("
        "ak.index.Index64(offsets), ak.contents.NumpyArray(content)))",
        glb,
        number=number,
    )
    report(f"ak.Array(layout)              n={n}", "ak.Array(layout)", glb, number=number)
    report(f"array[:]                      n={n}", "array[:]", glb, number=number)
    report(f"array[1:-1]                   n={n}", "array[1:-1]", glb, number=number)
    report(f"array[::2]                    n={n}", "array[::2]", glb, number=number)
    report(f"array.layout.offsets[:]       n={n}", "array.layout.offsets[:]", glb, number=number)
    report(f"ak.num(array)                 n={n}", "ak.num(array)", glb, number=number)
    report(f"ak.flatten(array)             n={n}", "ak.flatten(array)", glb, number=number)
    report(f"ak.to_layout(array)           n={n}", "ak.to_layout(array)", glb, number=number)