Add indexed queries

README.md

@@ -15,9 +15,9 @@ Both paths support raw bytes, zstd, and Blosc-zstd and return contiguous,
 DLPack-compatible tensors. CPU builds have no CUDA or nvCOMP dependency.
 
 Source URIs identify concrete Zarr v3 arrays, including arrays inside
-[NGFF](https://ngff.openmicroscopy.org/) multiscale images. Queries currently
-select rectangular regions in array coordinates. Index queries, transformed
-crops, and automatic NGFF level selection are planned extensions.
+[NGFF](https://ngff.openmicroscopy.org/) multiscale images. Queries select
+rectangular regions or ordered index arrays along each dimension. Transformed
+crops and automatic NGFF level selection are planned extensions.
 
 ## CPU quick start
 
@@ -68,6 +68,23 @@ install dependencies with `brew install cmake ninja pkg-config zstd c-blosc`:
 pip install . --config-settings=cmake.define.DAMACY_CUDA=OFF
 ```
 
+## Indexed queries
+
+Use `IndexQuery` to mix index arrays and contiguous slices:
+
+```python
+query = damacy.IndexQuery(
+    uri="/data/image.zarr/0",
+    selection=([7, 2, 7], slice(16, 80), [100, 3, 40, 3]),
+)
+assert query.shape == (3, 64, 4)
+```
+
+Choose `BatchSpec` with that sample shape and pass the query to `pipeline.push`.
+Each axis is independent: all combinations of its selected indices appear in
+the output. Order and duplicates are preserved. Both CPU and CUDA support the
+same queries; see [query semantics](docs/pipeline.md#indexed-queries).
+
 ## CUDA quick start
 
 ```python

bench/main.c

@@ -729,13 +729,15 @@ fill_random_sample(const struct scenario* sc,
 {
   uint32_t z = (uint32_t)rng_range(rng, at->n);
   s->uri = &at->uris[(size_t)z * BENCH_MAX_URI];
-  s->aabb.rank = at->rank;
+  s->rank = at->rank;
   const int64_t* shape = &at->shapes[(size_t)z * DAMACY_MAX_RANK];
   for (uint8_t d = 0; d < at->rank; ++d) {
     int64_t span = shape[d] - sc->sample_shape[d] + 1;
     int64_t beg = (int64_t)rng_range(rng, (uint64_t)span);
-    s->aabb.dims[d].beg = beg;
-    s->aabb.dims[d].end = beg + sc->sample_shape[d];
+    s->axes[d] = (struct damacy_axis_selection){
+      .kind = DAMACY_AXIS_INTERVAL,
+      .interval = { beg, beg + sc->sample_shape[d] }
+    };
   }
 }
 

dev/cpu-pipeline.md

@@ -1,9 +1,9 @@
 # CPU pipeline and query architecture
 
 The CPU milestone introduces a shared preparation stage and separate CPU and
-CUDA executors. Both support the existing rectangular queries. A CPU build
-returns results in ordinary RAM and has no CUDA dependency. Index queries,
-spatial resampling, and NGFF level selection remain later features.
+CUDA executors. Both support rectangular and indexed queries. A CPU build
+returns results in ordinary RAM and has no CUDA dependency. Spatial resampling
+and NGFF level selection remain later features.
 
 The public construction and lifetime contracts are in
 [Pipeline composition](../docs/pipeline.md). The C factory API is
@@ -16,7 +16,7 @@ existing callers.
 | Component | Responsibility | Implementation |
 | --- | --- | --- |
 | Metadata reader/provider | Zarr descriptions, shard indexes, asynchronous metadata I/O, cache capacities. | `FileMetadataReader`, `ZarrMetadata`; active caches in `pipeline/zarr_planner.c`. |
-| Planner | Validate rectangular requests and publish owned source/result plans. | `ChunkPlanner`, `planner/plan_builder.c`. |
+| Planner | Validate requests, enumerate selected chunks, and publish owned source/result plans. | `ChunkPlanner`, `planner/plan_builder.c`. |
 | Executor | Read encoded data, prepare codecs, decode, assemble, manage output storage. | `executor/cpu_executor.c`, `executor/cuda_executor.c`. |
 | Pipeline | Bound preparation, preserve batch order, retry backpressure, publish failures, coordinate shutdown. | `damacy_plan.c`, `damacy_scheduler.c`, `damacy_pop.c`. |
 
@@ -79,10 +79,10 @@ adjacent or overlapping ranges within each shard and interleave merged reads
 across shards. Each unique chunk retains its offset within a merged read; its
 uses still control output assembly. CPU reads use exact encoded byte ranges,
 without CUDA's page alignment. It handles raw bytes, zstd, and C-Blosc zstd
-with no, byte, or bit shuffle. It copies clipped chunk intersections and casts
-supported source types to `f32` or `bf16`, including fill-only chunks. The CUDA
-executor currently retains its per-use decode behavior; cross-sample GPU decode
-reuse can be optimized separately.
+with no, byte, or bit shuffle. It copies clipped chunk intersections, gathers
+indexed positions, and casts supported source types to `f32` or `bf16`,
+including fill-only chunks. The CUDA executor currently retains its per-use
+decode behavior; cross-sample GPU decode reuse can be optimized separately.
 
 ## Resources and ownership
 
@@ -115,20 +115,30 @@ buffer saturation report retriable backpressure, not a storage error. Closing
 a pipeline stops preparation, wakes blocked pops, joins execution, and releases
 queued work before its dependencies can disappear.
 
-## Next: index queries
-
-An index query should accept ordered index vectors along selected dimensions,
-with ordinary contiguous ranges along the others. Preserve caller order and
-repeated indices. Multiple indexed dimensions should use a Cartesian product,
-matching MATLAB's per-dimension indexing; this differs from paired-coordinate
-point queries. A stencil can generate an index vector, so no separate stencil
-query type is needed.
-
-Extend logical result operations with compact index vectors and group chunk
-uses during preparation. Keep output shape independent of the source bounding
-range. Bound owned index storage along with query/plan storage. CPU assembly can
-start with gather operations; CUDA can select an appropriate gather kernel.
-Do not create one planning record per output voxel.
+## Indexed queries
+
+`IndexQuery` supplies ordered index vectors along selected dimensions and
+contiguous slices along the others. Multiple indexed dimensions form a Cartesian
+product. Order and repeated indices survive gathering; output shape is independent
+of the bounding range. A stencil can generate an index vector without a separate
+query type.
+
+`query/selection.c` copies and sorts each vector by source position while retaining
+its output position. Preparation enumerates only touched grid cells, first for
+shard metadata and then for inner chunks. A chunk has one use per sample even if
+many selected positions, including duplicates, fall within it. `PLAN_GATHER`
+regions own sorted index records along with the rest of the plan. Storage grows
+with the sum of index-vector lengths and selected chunk uses.
+
+CPU assembly intersects each axis with a decoded chunk, then writes the Cartesian
+product of selected positions to the specified output positions. Contiguous
+trailing axes retain row copies. CUDA dispatch writes per-chunk selection spans
+to a per-batch array that only indexed samples reference, so rectangular chunk
+records carry no selection data. It uploads that array once per batch with
+compact source/output index pairs. Assembly reuses the existing read, decode,
+fill, shuffle, and cast paths. Index buffers hold the most indices a batch can
+contain and are included in the GPU budget before wave sizing. A nonzero
+`max_index_bytes` must cover them; zero rejects indexed samples at push.
 
 ## Next: spatial resampling and NGFF
 
@@ -158,7 +168,7 @@ the private interfaces leave room to add that resolution stage.
 Resampling needs output tiles with all contributing source chunks available.
 Keep decoded data until dependent tiles finish, and give each output tile one
 writer. The shared plan already separates chunks from their uses, but the
-current chunk-at-a-time rectangular executor will need this additional mode.
+current chunk-at-a-time copy/gather executor will need this additional mode.
 Boundary extension and filter halos must use the selected level's coordinates;
 voxel-center conventions and downsampling quality need explicit tests.
 

docs/api.md

@@ -10,6 +10,8 @@ The full public surface of the `damacy` package.
 
 ::: damacy.Sample
 
+::: damacy.IndexQuery
+
 ::: damacy.Batch
 
 ## Components

docs/budget.md

@@ -27,7 +27,8 @@ configuration — then round up generously:
 pool_reservation       = 2 × samples_per_batch × prod(sample_shape) × dtype_bytes
 one_chunk_per_wave     ≈ max_chunk_uncompressed_bytes × 2   # compressed + decoded buffers, one chunk
 both_waves_one_chunk   ≈ 2 × one_chunk_per_wave             # two GPU waves resident at once
-budget_floor           ≈ pool_reservation + both_waves_one_chunk + scratch_slack
+index_reservation      = 2 × (8 × samples_per_batch × sum(sample_shape) + 16 × 16384 × rank)   # 0 if max_index_bytes = 0
+budget_floor           ≈ pool_reservation + index_reservation + both_waves_one_chunk + scratch_slack
 ```
 
 The `2 ×` in `pool_reservation` is output-side double-buffering.
@@ -99,17 +100,23 @@ budget.
 
 ### Where the budget is spent
 
-Four buckets share `max_gpu_memory_bytes`:
+These allocations share `max_gpu_memory_bytes`:
 
 | Bucket                  | What it holds                                              | Scales with                                              |
 | ----------------------- | ---------------------------------------------------------- | -------------------------------------------------------- |
 | Output batch pool       | The batches you `pop`, double-buffered                     | `samples_per_batch`, `sample_shape`, dtype                      |
 | Wave-resident buffers   | Compressed + decoded chunk bytes for the two in-flight waves | budget headroom, in `max_chunk_uncompressed_bytes` steps |
 | Decoder scratch         | nvcomp's working memory                                    | Peak sub-stream count in the dataset                     |
-| Per-wave metadata       | Pointer/size arrays for the decoder                        | Peak sub-stream count                                    |
-
-The first two are the large ones; the last two are small but
-*depend on the data*. damacy cannot know the sub-stream count of
+| Per-wave metadata       | Pointer/size arrays for the decoder and assembly           | Peak sub-stream and chunk counts                         |
+| Indexed-query data      | Source/output index pairs and per-chunk selection ranges for two batch slots | Output shape and rank; none when `max_index_bytes=0` |
+
+Index storage is reserved at construction from the output shape. A nonzero
+`max_index_bytes` below `8 × samples_per_batch × sum(sample_shape)` raises
+`BudgetExceeded` from `Pipeline(cfg)`. Each slot also reserves 16 bytes per
+axis for up to 16384 chunks, to record which part of each chunk an indexed
+query selects. Setting `max_index_bytes=0` drops both reservations; `push`
+then rejects indexed queries.
+Decoder scratch and fanout storage depend on the data. damacy cannot know the sub-stream count of
 a chunk until it inspects the chunk's header, so damacy picks
 per-wave geometry such that even after adaptive growth to the
 structural ceiling, the total fits inside the cap. Grows commit

docs/pipeline.md

@@ -3,7 +3,7 @@
 A pipeline receives a planner, an executor, an output specification, and queue
 limits. The planner owns metadata preparation; the executor owns bulk reads,
 decoding, and output buffers. CPU and CUDA execution use the same rectangular
-queries and prepared-plan contract.
+and indexed queries and prepared-plan contract.
 
 ## Construct a CPU pipeline
 
@@ -96,10 +96,79 @@ The injected readers inherit the caller's host affinity when their workers
 start. The legacy `Config.numa_strategy` also applies while constructing those
 readers, preserving placement of the complete legacy pipeline.
 
+## Indexed queries
+
+`IndexQuery(uri, selection=...)` accepts one index array or contiguous `slice`
+per source axis. For a `(z, y, x)` array:
+
+```python
+query = damacy.IndexQuery(
+    uri="/data/image.zarr/0",
+    selection=([7, 2, 7], slice(16, 80), [100, 3, 40, 3]),
+)
+output = damacy.BatchSpec(samples=1, shape=query.shape, dtype="f32")
+assert query.shape == (3, 64, 4)
+
+with damacy.Pipeline(
+    planner=planner, executor=executor, output=output,
+    queues=damacy.QueueLimits(lookahead_samples=2),
+) as pipeline:
+    pipeline.push([query])
+    with pipeline.pop() as batch:
+        result = np.from_dlpack(batch)  # CPU executor
+        del result
+```
+
+Index arrays select independently along each dimension: their Cartesian product
+fills the tensor, matching MATLAB's per-dimension indexing or NumPy's `np.ix_`.
+Caller order and repeated indices are preserved. A singleton index array keeps
+its axis. A tuple such as `(7, 2)` is an index array in `IndexQuery`; use a slice
+for a contiguous range. `Sample.aabb` retains its existing interval syntax.
+
+Indices must be nonnegative integers below `2**63 - 1`. Slices use half-open
+bounds, require an explicit stop, and accept a step of one. A `range` object can
+supply a strided or reversed index array. Empty selections are rejected. Source
+bounds are validated when metadata arrives; out-of-bounds selections raise
+`InvalidArgument` from `pop`. Each query's result shape must match `BatchSpec`.
+Rectangular and indexed queries can share a batch when their result shapes match.
+
+Python copies index values into the immutable query at construction. Native
+admission copies accepted queries, and prepared plans own their index data.
+The planner enumerates only selected shards and chunks, including when indices
+span large gaps. `max_chunks` counts each selected chunk once per sample;
+repeated indices inside that chunk do not consume extra chunk entries.
+
+The C API uses `damacy_sample.rank` and one tagged `damacy_axis_selection`
+per axis. C callers must rebuild against the updated headers and explicitly
+set every active axis to `DAMACY_AXIS_INTERVAL` or `DAMACY_AXIS_INDICES`:
+
+```c
+int64_t rows[] = {7, 2, 7};
+struct damacy_sample query = {
+    .uri = "/data/image.zarr/0",
+    .rank = 2,
+    .axes = {
+        {.kind = DAMACY_AXIS_INDICES, .indices = {.values = rows, .count = 3}},
+        {.kind = DAMACY_AXIS_INTERVAL, .interval = {.beg = 4, .end = 8}},
+    },
+};
+```
+
+This requests a `(3, 4)` sample. The union's `kind` selects its active member;
+there is no default kind. A zero or unknown tag, a null or empty index array,
+or an interval with negative, empty, or reversed bounds returns `DAMACY_INVAL`
+from `damacy_push`. Each axis's length must match the configured sample shape.
+Only `axes[0..rank)` are active, in the Zarr array's stored axis order. Preparation
+derives the bounding AABB from these selections.
+
+`damacy_push` copies the URI and index values for its consumed prefix. The
+unconsumed suffix remains caller-owned and can be retried.
+
 ## Limits and backpressure
 
-Sizes are bytes, with positive explicit limits. Zero is invalid for these
-capacities. Defaults come from the Python value objects.
+Sizes are bytes, with positive explicit limits. `CudaLimits.max_index_bytes`
+also accepts zero, which rejects indexed CUDA queries at push. Defaults come
+from the Python value objects.
 
 | Setting | Scope |
 | --- | --- |
@@ -110,13 +179,25 @@ capacities. Defaults come from the Python value objects.
 | `PlanLimits.max_chunks` | Chunk uses per batch, including chunks used by multiple samples. |
 | `PlanLimits.max_chunk_bytes` | Decoded source bytes per chunk, before output conversion. |
 | `PlanLimits.max_shards_per_sample` | Maximum number of shard files touched by one sample. |
-| `PlanLimits.max_plan_bytes` | Owned storage per prepared plan. |
+| `PlanLimits.max_plan_bytes` | Owned storage per prepared plan, including index arrays; also caps copied index data per queued query. |
 | `CpuLimits.max_memory_bytes` | CPU executor's buffers, codec workspace, active read plans, and temporary read-planning scratch. |
 | `CpuLimits.decode_workers` | Total decoding/assembly workers, including the calling scheduler thread. |
 | `CpuLimits.chunks_per_input_buffer` | Chunks each of the two encoded-input buffers holds; from `decode_workers` to 16384. |
 | `FileReader.workers` | Bulk I/O workers, separate from decoding workers. |
 | `FileReader.max_inflight_reads` | Bulk read capacity; execution respects this bound and retries saturation. |
 | `CudaLimits` | GPU memory and execution geometry, plus the CUDA codec-layout cache capacity. |
+| `CudaLimits.max_index_bytes` | Device index storage per batch: eight bytes per index across all indexed axes and samples. Zero, or at least `8 * samples * sum(shape)`. Default 64 MiB. |
+
+CUDA reserves index storage for its two execution slots within
+`max_gpu_memory_bytes`. A nonzero `max_index_bytes` must hold every index a
+batch can contain, `8 * samples * sum(shape)` bytes, so an accepted query never
+runs out of room. A smaller value raises `BudgetExceeded` from `Pipeline`. Each
+slot allocates that amount, plus 16 bytes per axis for up to 16384 chunks to
+record each chunk's selected range. An identical amount of pinned host staging
+is allocated. With `max_index_bytes=0`, `push` raises `BudgetExceeded` for an
+`IndexQuery`; the query is not consumed and the pipeline keeps running.
+`Config.max_index_bytes` provides the same setting through the CUDA
+convenience adapter.
 
 For injected pipelines, define `floor = queues.lookahead_samples + output.samples`.
 The metadata cache requires at least `floor` array entries and
@@ -242,9 +323,8 @@ nvCOMP, and a runtime NVIDIA driver. GDS requires a CUDA build.
 
 ## Future queries
 
-The first milestone implements rectangular copy/cast queries. Index-array
-queries can add ordered selections and repeated indices without a separate
-stencil type. Spatial queries will describe a fixed output tensor, a transform
+Rectangular and indexed queries copy or gather existing voxels. Spatial
+queries will describe a fixed output tensor, a transform
 from output coordinates to source space, and sampler settings including
 interpolation, antialiasing, and boundary handling.