Proposal: data-shape-aware fast paths and AArch64 vectorization guidance

[GENG-CHOVYYYY]

Summary

I use Daft on AArch64 servers for multimodal data-engineering workloads. Profiling has exposed a few narrow, data-shape-dependent hot paths in sort, join, and take. I would like to contribute guarded fast paths for those cases and get guidance on evaluating portable vectorization and AArch64 affinity tuning without introducing CPU-family-specific behavior.

Candidate work

Boolean and dense integer argsort

• For null-free boolean arrays, build the permutation from the false/true counts.
• For null-free integer arrays with a small, dense key range, use a counting-index path.
• Retain the existing implementation for nulls, sparse or wide ranges, and unsupported sort semantics.

Dense integer join probe

When integer build keys occupy a compact, bounded range, a direct-index structure can avoid a general hash lookup on the probe side. This seems related to #7076, so I would prefer maintainer advice on whether to continue the discussion there before opening another issue.

Identity and range take

Detect identity indices and contiguous ranges, then use clone or slice semantics instead of a materializing gather. Arbitrary indices would continue through the generic gather path.

Vectorization and affinity on AArch64

For primitive loops used by comparisons, hashing, filters, and gather-like access, I would investigate the following order:

1. A portable algorithmic or representation improvement.
2. Compiler auto-vectorization or an existing shared abstraction.
3. NEON/SVE only when target detection, a generic fallback, and reproducible results justify it.

For scheduling, I am looking for a supported way to experiment with opt-in parallelism or affinity settings. I do not want to change Daft's global defaults based on one server family.

Validation

Before opening a PR, I would provide a small public fixture, a Rust microbenchmark covering both the fast-path trigger and fallback, semantic tests, and an end-to-end check. I would report AArch64 and x86 results separately. The work can be reproduced without upstreaming proprietary data.

Questions

1. Would maintainers prefer separate, narrowly scoped issues for sort, join, and take?
2. Should the dense integer join path be discussed under perf: hash join probe loop has iterator overhead and per-row Growable extends #7076, or in a separate issue?
3. Are there existing SIMD abstractions, benchmark conventions, or architecture-aware scheduling policies that contributors should use first?

[srilman]

Hi @GENG-CHOVYYYY, thanks for opening this up. This is something we've been thinking about for a while, so its great that you brought it up.

• Yes, narrower scopes or each section (sort, join, take, etc) would be preferable
• I do think the dense integer join path can be in another issue, maybe built on the top of the existing perf one.
• Right now:
  • For x86, we maintain 2 versions of daft (regular Daft that supports SIMD up to AVX2 tuned for Skylake and daft-lts to target architectures before that). We don't have a build that supports AVX-512 or newer. To support that, we would want some sort of architecture-aware scheduling policy.
  • For ARM64, we don't have any restrictions. This is something we would need to decide first to see if we just use a minimum SIMD instruction set or some sort of scheduling policy.
• Ideally a SIMD abstraction would help minimize the amount of custom code we would need per architecture. For example, AVX-512 has specific instructions that can be used to heavily optimize take operations, but don't have exact equivalents in AVX2 and below
• For benchmarking, right now we use CodSpeed for microbenchmarking in cases like this. I would like if we relied on it more for profiling specific operations / functions like this.

[GENG-CHOVYYYY]

@srilman Thanks, this is very helpful.

I’ll split the work into narrower issues for sort, join, and take. For the dense integer join path, I’ll open a separate issue and link it back to #7076, rather than extending that thread with several unrelated cases.

For ARM64, I agree that the first question is the portability boundary: whether Daft should target a minimum SIMD baseline or use architecture-aware dispatch and scheduling. I’ll start by profiling the candidate operations and checking for opportunities that fit a shared SIMD abstraction before proposing any architecture-specific implementation.

I’ll also look into adding focused CodSpeed coverage for the trigger and fallback cases. That should make it easier to compare the portable path, existing implementation, and any future architecture-specific path without relying on machine-specific numbers.

I’ll follow up with the first scoped issue once I have the minimal reproduction and benchmark shape ready. Thanks again for the guidance.

[GENG-CHOVYYYY]

Following up on #7389, I opened #7422 to scope the dense integer argsort fast path separately. I have a guarded prototype and paired baseline results ready for review.

[GENG-CHOVYYYY]

Hi @srilman — following up on your guidance in #7389.

The scoped implementation is now available in #7430. It keeps the change limited to a guarded dense-integer argsort path, with explicit fallback coverage for sparse/wide ranges and null-containing inputs. All CI checks and CodSpeed have passed.

Could you please confirm whether #7422 is approved for this narrow scope, or let me know what design changes or additional evidence would be needed before review? I will keep the work limited to this path.