Add gRPC multi-query support (onlineQueryMulti)

[sjmignot]

Summary

• add ChalkClient.onlineQueryMulti, which runs several independent online queries in parallel over a single request, via the chalk.engine.v1.QueryService/OnlineQueryMulti RPC
• add OnlineQueryMultiResult: positional per-query results, request-level errors, shared trace id, one close() for the whole batch
• decompose GRPCClient.onlineQuery into serializeInputs / buildBulkRequest / queryStub / decodeBulkResponse, now shared by both paths so single and multi build byte-identical sub-requests
• fix decodeBulkResponse to rethrow on a decode failure instead of returning a result wrapping already-closed Arrow objects

The generated stubs for this RPC were already vendored in ai.chalk.protos and had never been wired up, so there is no proto regeneration here. Callers with N unrelated queries — different roots, different outputs — previously paid N round trips.

Usage

ChalkClient client = ChalkClient.createGrpc();

OnlineQueryParamsComplete users = OnlineQueryParams.builder()
        .withInput("user.id", Arrays.asList(1, 2, 3))
        .withOutputs("user.email", "user.credit_score")
        .build();

OnlineQueryParamsComplete merchants = OnlineQueryParams.builder()
        .withInput("merchant.id", Arrays.asList("m1", "m2"))
        .withOutputs("merchant.risk_score")
        .build();

try (OnlineQueryMultiResult multi = client.onlineQueryMulti(List.of(users, merchants))) {
    // errors the engine did not attach to one query
    for (ServerError error : multi.getGlobalErrors()) {
        log.warn("multi-query error: {}", error.getMessage());
    }

    // getResults() is 1:1 with the queries, in the order they were passed
    OnlineQueryResult userResult = multi.getResults().get(0);
    OnlineQueryResult merchantResult = multi.getResults().get(1);

    if (userResult.getErrors().length == 0) {
        User[] fetched = userResult.unmarshal(User.class);
    }
}
// closing the wrapper closes every sub-result and its allocator

There is also a varargs overload for the common case:

try (OnlineQueryMultiResult multi = client.onlineQueryMulti(users, merchants)) {
    ...
}

Use onlineQuery for one query over many rows — it is already a bulk operation. onlineQueryMulti is for queries that are unrelated to each other.

Semantics

• Results are positional. getResults() holds one result per query, in the order passed. The engine pre-sizes its response vector, so the counts always match; a mismatch is rejected client-side rather than silently misattributed, since position is the only correlation the proto offers.
• Queries fail independently. A resolver that fails while computing one query is reported on that query's own getErrors() and leaves the others intact.
• One caveat worth knowing: a query that fails to run at all (rate limited, or an internal engine error) is reported in getGlobalErrors() with its 1-based position in the message text, while its slot in getResults() holds an empty result reporting no errors of its own. An empty result is therefore not proof of success — check the global errors too. This is documented on the API and covered by a test.
• Every sub-query is sent as a bulk_request. The engine rejects single_request outright ("OnlineQueryRequest is unsupported; Please use OnlineQueryBulkRequest w/ a single row."), which is also what the reference Python client does.
• Some settings apply to the whole request, which carries one deadline and one set of headers. The queries must agree on environmentId, branch and queryName or the call throws ClientException naming both conflicting indices: the API server resolves one branch to one deployment, and named queries are routed by an exact header match on x-chalk-query-name, so a mixed batch could not be honored. Leaving one of these unset just takes the client's value, which is not a disagreement, and queryNameVersion may differ freely since each sub-request carries its own. The deadline is the longest of the effective per-query timeouts, where a query that sets none is bounded by the client-level timeout; per-query timeouts are not independently enforced.

Allocator ownership

Each sub-result gets its own Arrow child allocator. OnlineQueryResult.close() closes the allocator it is handed, so a shared allocator would be released by the first sub-result closed while the others' tables still held buffers. Closing the OnlineQueryMultiResult closes every sub-result; callers should not close them individually. A failure partway through decoding closes everything already decoded and retains no Arrow memory.

Behavior change

decodeBulkResponse (the extracted decode step, previously inline in onlineQuery) now rethrows. The old code closed its tables and allocator on a decode failure but then fell through to return, handing back an OnlineQueryResult whose Arrow objects were already closed — so any use of it threw anyway. The multi path cannot tolerate this, since a mid-batch failure would also leak the sibling allocators. This turns an obscure downstream IllegalStateException into a ClientException with a message.

Review fixes (second commit)

A review pass over the settings resolvers turned up several problems, all fixed in 0450ccb:

• resolveDeadline dropped the client-level timeout as soon as any one query set its own. A client built with .withTimeout(5m) batching a 1s query alongside a query with no timeout deadlined the whole request at 1s, killing a query that would have had 5 minutes on its own — the opposite of the "no sub-query is cancelled before the time its caller allowed" invariant the method claimed. An unset per-query timeout means the client-level timeout, so it now takes part in the maximum, and a query that is effectively unbounded yields no deadline at all.
• resolveEnvironment compared raw values, so naming the client's own environment on one query and leaving it unset on another was rejected as a conflict even though both resolve identically. It now compares the effective value, as resolveBranch already did.
• Dropped the queryNameVersion agreement requirement. The routing justification does not apply to the version: there is no x-chalk-query-name-version header, and each sub-request carries its own version in its OnlineQueryContext, so queries differing only by version are satisfiable. (Agreement on the name stays, since Envoy route-matches on it.)
• unmarshal() NPE'd on a null scalars table — exactly what the empty slot above leaves behind — so the documented trap made this PR's own README pattern fail with an opaque "Failed to unmarshal" wrapping an NPE. It now names the cause and points at getGlobalErrors(), and the README example guards the case explicitly.
• Minor: the varargs overload now reports the same ClientException as the List overload for a null array instead of a raw NPE, and a repeated format argument that made the count-mismatch message tautological is gone.

Verification

• TestGrpcMultiQuery (new, 23 cases, no engine required) — all passed
• full offline suite — 94 passed, 0 failed (TestGrpcMultiQuery, TestOnlineQueryParams, ai.chalk.arrow.*, ai.chalk.internal.*, ai.chalk.flink.*)
• compileJava, compileTestJava, javadoc, build, shadowJar — all succeed; javadoc warnings unchanged from main
• leak safety is asserted directly: testDecodeMultiResponseDecodeFailureLeaksNothing and testDecodeBulkResponseReleasesAllocatorOnFailure force an Arrow decode failure and assert getAllocatedMemory() == 0. The first fails against the pre-fix decode path.
• the golden-request test asserts buildBulkRequest field by field, which is what guards the single-query path against drift during the refactor
• not run locally: the 6 new integration tests in TestGrpcClient and the REST-rejection test in TestAllClients. They compile but need a live deployment, and there were no CHALK_* credentials in this environment. CI supplies them.

Follow-ups (filed, not fixed here)

• withBranch() is inert on the gRPC path. RequestHeaderInterceptor never sends X-Chalk-Branch-Id, and proxyForRequest reads branch only from that header — OnlineQueryContext.branch_id is [deprecated = true] ("set in the header. This field is not set anymore") and is read by neither the API server nor the engine. This affects single queries equally, so it is out of scope here, but it means onlineQueryMulti validates branch agreement without the branch being honored. Called out in the Javadoc and README. Same story for X-Chalk-Preview-Deployment.
• staleness and storePlanStages are never sent on the gRPC path, despite OnlineQueryBulkRequest.staleness existing. Needs a Duration → bucket-string encoding decision; affects single and multi.
• OnlineQueryResultFeather.close() closes a shared allocator once per map entry — latent in the REST path, masked only because toResult() reads the single key "0". Must be fixed before REST multi-query; a TODO now sits at that call site.
• Requiring agreement on queryName is stricter than the Python gRPC client, which allows mixed names and sends no query-name header at all. Java's gRPC client does send it, and Envoy route-matches on it, hence the stricter rule. Relaxing a validation later is backward-compatible, so this starts strict — but it is the decision most worth a second opinion, since "queries rooted in different feature classes" often means differently-named queries in production.
• validateMultiParams rejects a queryNameVersion set without a queryName, which plain onlineQuery currently accepts and forwards. This is a deliberate port of the reference client's rule (resolve_multi_query_query_name), but it does mean a params object that works via onlineQuery is refused here; applying it to onlineQuery too would be a behavior change needing a release note.
• There is no programmatic way to map a global error back to its sub-query — the engine encodes the position in the message text only. A structured per-slot "did not execute" signal would be a proto change.
• A batch's response is the sum of its sub-responses against the same 500MB maxInboundMessageSize, so very large batches will hit that ceiling sooner than the equivalent separate queries.

[github-actions]

Auto approved automated PR due to label