diff --git a/src/main/java/ai/chalk/client/ChalkClientImpl.java b/src/main/java/ai/chalk/client/ChalkClientImpl.java index 58a0adc..df3c2ae 100644 --- a/src/main/java/ai/chalk/client/ChalkClientImpl.java +++ b/src/main/java/ai/chalk/client/ChalkClientImpl.java @@ -3,6 +3,7 @@ import ai.chalk.exceptions.ChalkException; import ai.chalk.exceptions.ClientException; +import ai.chalk.exceptions.ServerError; import ai.chalk.internal.arrow.FeatherProcessor; import ai.chalk.internal.bytes.BytesProducer; import ai.chalk.internal.config.Loader; @@ -13,11 +14,20 @@ import ai.chalk.internal.request.models.SendRequestParams; import ai.chalk.models.OnlineQueryParamsComplete; import ai.chalk.models.OnlineQueryResult; +import ai.chalk.models.QueryMeta; import ai.chalk.models.UploadFeaturesParams; import ai.chalk.models.UploadFeaturesResult; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.PropertyNamingStrategies; +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; +import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.*; +import org.apache.arrow.vector.table.Table; import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; import java.util.*; public class ChalkClientImpl implements ChalkClient { @@ -58,6 +68,13 @@ public ChalkClientImpl(BuilderImpl config) throws ChalkException { } public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws ChalkException { + if (BytesProducer.hasHasManyInputs(params)) { + return onlineQueryJson(params); + } + return onlineQueryFeather(params); + } + + private OnlineQueryResult onlineQueryFeather(OnlineQueryParamsComplete params) throws ChalkException { byte[] bodyBytes; try ( var childAllocator = allocator.newChildAllocator( @@ -88,11 +105,115 @@ public OnlineQueryResult onlineQuery(OnlineQueryParamsComplete params) throws Ch 0, FeatherProcessor.ALLOCATOR_SIZE_RESPONSE ); - // ignore the warning here, because we don't want to free the memory yet var bulkResponse = OnlineQueryBulkResponse.fromBytes(response.body(), allocator); return bulkResponse.toResult(); } + private OnlineQueryResult onlineQueryJson(OnlineQueryParamsComplete params) throws ChalkException { + Map jsonBody = BytesProducer.buildJsonOnlineQueryBody(params); + + SendRequestParams request = new SendRequestParams.Builder(params.getTimeout()) + .path("/v1/query/online") + .body(jsonBody) + .method("POST") + .branch(params.getBranch()) + .isEngineRequest(true) + .previewDeploymentId(params.getPreviewDeploymentId()) + .environmentOverride(params.getEnvironmentId()) + .queryName(params.getQueryName()) + .build(); + + HttpResponse response = this.handler.sendRequest(request); + + ObjectMapper mapper = new ObjectMapper(); + mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); + mapper.setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE); + mapper.registerModule(new JavaTimeModule()); + + Map responseMap; + try { + responseMap = mapper.readValue(response.body(), Map.class); + } catch (Exception e) { + throw new ClientException("Failed to parse JSON query response", e); + } + + ServerError[] errors = null; + if (responseMap.get("errors") instanceof List errorsList && !errorsList.isEmpty()) { + errors = new ServerError[errorsList.size()]; + for (int i = 0; i < errorsList.size(); i++) { + errors[i] = mapper.convertValue(errorsList.get(i), ServerError.class); + } + } + + QueryMeta meta = null; + if (responseMap.get("meta") != null) { + meta = mapper.convertValue(responseMap.get("meta"), QueryMeta.class); + } + + BufferAllocator childAllocator = this.allocator.newChildAllocator( + "online_query_json_response", 0, FeatherProcessor.ALLOCATOR_SIZE_RESPONSE + ); + + Table scalarTable = null; + if (responseMap.get("data") instanceof List dataList && !dataList.isEmpty()) { + List vectors = new ArrayList<>(); + try { + for (Object item : dataList) { + var dataItem = (Map) item; + String fieldName = (String) dataItem.get("field"); + Object value = dataItem.get("value"); + if (fieldName == null) continue; + + FieldVector vector = createVectorForValue(fieldName, value, childAllocator); + if (vector != null) vectors.add(vector); + } + scalarTable = new Table(vectors); + } catch (Exception e) { + for (FieldVector v : vectors) v.close(); + childAllocator.close(); + throw new ClientException("Failed to build Arrow table from JSON response", e); + } + } + + return new OnlineQueryResult(scalarTable, new HashMap<>(), errors, meta, childAllocator); + } + + private static FieldVector createVectorForValue(String name, Object value, BufferAllocator allocator) { + if (value == null) { + return new NullVector(name, 1); + } else if (value instanceof Integer intVal) { + BigIntVector v = new BigIntVector(name, allocator); + v.allocateNew(1); + v.set(0, intVal.longValue()); + v.setValueCount(1); + return v; + } else if (value instanceof Long longVal) { + BigIntVector v = new BigIntVector(name, allocator); + v.allocateNew(1); + v.set(0, longVal); + v.setValueCount(1); + return v; + } else if (value instanceof Double doubleVal) { + Float8Vector v = new Float8Vector(name, allocator); + v.allocateNew(1); + v.set(0, doubleVal); + v.setValueCount(1); + return v; + } else if (value instanceof Boolean boolVal) { + BitVector v = new BitVector(name, allocator); + v.allocateNew(1); + v.set(0, boolVal ? 1 : 0); + v.setValueCount(1); + return v; + } else { + VarCharVector v = new VarCharVector(name, allocator); + v.allocateNew(1); + v.set(0, value.toString().getBytes(StandardCharsets.UTF_8)); + v.setValueCount(1); + return v; + } + } + public UploadFeaturesResult uploadFeatures(UploadFeaturesParams params) throws ChalkException { byte[] tableBytes; try { diff --git a/src/main/java/ai/chalk/internal/bytes/BytesProducer.java b/src/main/java/ai/chalk/internal/bytes/BytesProducer.java index 6338cf8..d2e2f3c 100644 --- a/src/main/java/ai/chalk/internal/bytes/BytesProducer.java +++ b/src/main/java/ai/chalk/internal/bytes/BytesProducer.java @@ -84,9 +84,95 @@ public static Map buildJsonHeader(OnlineQueryParams params) { } jsonHeader.put("context", onlineQueryContext); + Map hints = buildInputSchemaHint(params); + if (!hints.isEmpty()) { + jsonHeader.put("input_schema_hint", hints); + } + return jsonHeader; } + /** + * Builds the input_schema_hint map by merging auto-detected has-many inputs + * with any user-provided hints. User-provided hints take precedence. + */ + static Map buildInputSchemaHint(OnlineQueryParams params) { + Map hints = new HashMap<>(); + + if (params.getInputs() != null) { + for (var entry : params.getInputs().entrySet()) { + String fqn = entry.getKey(); + List values = entry.getValue(); + String hint = detectHasManyHint(fqn, values); + if (hint != null) { + hints.put(fqn, hint); + } + } + } + + if (params.getInputSchemaHint() != null) { + hints.putAll(params.getInputSchemaHint()); + } + + return hints; + } + + /** + * Detects whether an input value represents a has-many (list of structs/maps) + * and generates the bracketed projection hint string. + * Returns null if the input is not a has-many. + */ + private static String detectHasManyHint(String fqn, List values) { + if (values == null || values.isEmpty()) { + return null; + } + + // Scan all rows to find a non-empty inner list with Map elements + for (Object v : values) { + if (!(v instanceof List innerList) || innerList.isEmpty()) { + continue; + } + for (Object item : innerList) { + if (item instanceof Map mapItem) { + var sb = new StringBuilder(fqn).append('['); + boolean first = true; + for (Object key : mapItem.keySet()) { + if (!first) sb.append(','); + sb.append(key); + first = false; + } + sb.append(']'); + return sb.toString(); + } + } + } + return null; + } + + public static Map buildJsonOnlineQueryBody(OnlineQueryParamsComplete params) { + Map body = buildJsonHeader(params); + + Map singleRowInputs = new HashMap<>(); + for (var entry : params.getInputs().entrySet()) { + List values = entry.getValue(); + singleRowInputs.put(entry.getKey(), values.get(0)); + } + body.put("inputs", singleRowInputs); + body.put("encoding_options", Map.of("encode_structs_as_objects", true)); + + return body; + } + + public static boolean hasHasManyInputs(OnlineQueryParams params) { + if (params.getInputs() == null) return false; + for (List values : params.getInputs().values()) { + if (values == null || values.isEmpty()) continue; + Object first = values.get(0); + if (first instanceof List) return true; + } + return false; + } + public static byte[] convertOnlineQueryParamsToBytes(OnlineQueryParamsComplete params, BufferAllocator allocator) throws Exception { byte[] arrowBytes; if (params.getInputs() == null) { diff --git a/src/main/java/ai/chalk/internal/request/RequestHandler.java b/src/main/java/ai/chalk/internal/request/RequestHandler.java index 28201a9..f160c36 100644 --- a/src/main/java/ai/chalk/internal/request/RequestHandler.java +++ b/src/main/java/ai/chalk/internal/request/RequestHandler.java @@ -142,6 +142,8 @@ private byte[] getBodyBytes(Object body) throws Exception { bodyBytes = (byte[]) body; } else { ObjectMapper objectMapper = new ObjectMapper(); + objectMapper.registerModule(new JavaTimeModule()); + objectMapper.configure(com.fasterxml.jackson.databind.SerializationFeature.WRITE_DATES_AS_TIMESTAMPS, false); objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); objectMapper.setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE); bodyBytes = objectMapper.writeValueAsBytes(body); diff --git a/src/main/java/ai/chalk/models/OnlineQueryParams.java b/src/main/java/ai/chalk/models/OnlineQueryParams.java index 70b6933..c3ea4b0 100644 --- a/src/main/java/ai/chalk/models/OnlineQueryParams.java +++ b/src/main/java/ai/chalk/models/OnlineQueryParams.java @@ -139,6 +139,15 @@ public class OnlineQueryParams { */ private Duration timeout; + /** + * Schema hints for has-many input features. Maps a has-many feature FQN to its + * bracketed column projection string, e.g. + * {@code "user.txns" -> "user.txns[user.txns.txn_id,user.txns.amount]"}. + * Required when passing empty has-many inputs so the server knows the expected schema. + * Auto-generated from non-empty has-many input data when not explicitly set. + */ + private Map inputSchemaHint; + @AllArgsConstructor @NoArgsConstructor @@ -161,6 +170,7 @@ public static class Builder> { protected List requiredResolverTags; protected Map plannerOptions; protected Duration timeout; + protected Map inputSchemaHint; protected T _withInput(String fqn, List values) { if (this.inputs == null) { @@ -296,6 +306,22 @@ public T withTimeout(Duration timeout) { return (T) this; } + public T withInputSchemaHint(Map inputSchemaHint) { + if (this.inputSchemaHint == null) { + this.inputSchemaHint = new HashMap<>(); + } + this.inputSchemaHint.putAll(inputSchemaHint); + return (T) this; + } + + public T withInputSchemaHint(String fqn, String bracketedProjection) { + if (this.inputSchemaHint == null) { + this.inputSchemaHint = new HashMap<>(); + } + this.inputSchemaHint.put(fqn, bracketedProjection); + return (T) this; + } + // withNow takes a list of ZonedDateTimes and adds them to the now list public T withNow(List now) { if (this.now == null) { @@ -378,7 +404,8 @@ public OnlineQueryParams build() { now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } } @@ -402,7 +429,8 @@ public BuilderComplete( List now, List requiredResolverTags, Map plannerOptions, - Duration timeout + Duration timeout, + Map inputSchemaHint ) { super( inputs, @@ -422,7 +450,8 @@ public BuilderComplete( now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -486,7 +515,8 @@ public OnlineQueryParamsComplete build() { now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } } @@ -510,7 +540,8 @@ public BuilderWithInputs( List now, List requiredResolverTags, Map plannerOptions, - Duration timeout + Duration timeout, + Map inputSchemaHint ) { super( inputs, @@ -530,7 +561,8 @@ public BuilderWithInputs( now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -553,7 +585,8 @@ private BuilderComplete newBuilderComplete() { now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -618,7 +651,8 @@ public BuilderWithOutputs( List now, List requiredResolverTags, Map plannerOptions, - Duration timeout + Duration timeout, + Map inputSchemaHint ) { super( inputs, @@ -638,7 +672,8 @@ public BuilderWithOutputs( now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -662,7 +697,8 @@ public BuilderComplete newBuilderComplete() { now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -729,7 +765,8 @@ public BuilderSeed( List now, List requiredResolverTags, Map plannerOptions, - Duration timeout + Duration timeout, + Map inputSchemaHint ) { super( inputs, @@ -749,7 +786,8 @@ public BuilderSeed( now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -772,7 +810,8 @@ public BuilderWithInputs newBuilderWithInputs() { now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -795,7 +834,8 @@ public BuilderWithOutputs newBuilderWithOutputs() { now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } @@ -896,6 +936,19 @@ public final SingleRowInput withInputs(Map inputs) { return this; } + /** + * Add a has-many input: a list of child rows (each row a Map of column FQN to value) + * for a single-row query. The child rows represent the has-many relationship data. + * + * @param featureFqn The fully qualified name of the has-many feature (e.g. "user.transactions") + * @param rows The list of child row maps + * @return This SingleRowInput for method chaining + */ + public final SingleRowInput withHasManyInput(String featureFqn, List> rows) { + this.inputs.put(featureFqn, List.of(rows)); + return this; + } + /** * Get the internal inputs map. Used by builder classes to retrieve the input data. * diff --git a/src/main/java/ai/chalk/models/OnlineQueryParamsComplete.java b/src/main/java/ai/chalk/models/OnlineQueryParamsComplete.java index a81f002..1e08c45 100644 --- a/src/main/java/ai/chalk/models/OnlineQueryParamsComplete.java +++ b/src/main/java/ai/chalk/models/OnlineQueryParamsComplete.java @@ -49,7 +49,8 @@ public OnlineQueryParamsComplete( List now, List requiredResolverTags, Map plannerOptions, - Duration timeout + Duration timeout, + Map inputSchemaHint ) { super( inputs, @@ -69,7 +70,8 @@ public OnlineQueryParamsComplete( now, requiredResolverTags, plannerOptions, - timeout + timeout, + inputSchemaHint ); } } diff --git a/src/test/java/ai/chalk/client/TestOnlineQueryParams.java b/src/test/java/ai/chalk/client/TestOnlineQueryParams.java index 6a91317..ac3092a 100644 --- a/src/test/java/ai/chalk/client/TestOnlineQueryParams.java +++ b/src/test/java/ai/chalk/client/TestOnlineQueryParams.java @@ -886,6 +886,96 @@ public void testSingleRowInput() throws Exception { BytesProducer.convertOnlineQueryParamsToBytes(params3, allocator); } + @Test + public void testHasManyInputSchemaHintAutoGenerated() throws Exception { + Map row1 = new HashMap<>(); + row1.put("user.sms.sms_id", "S1"); + row1.put("user.sms.text", "hello"); + + List>> hasManyInput = new ArrayList<>(); + hasManyInput.add(Arrays.asList(row1)); + + OnlineQueryParamsComplete params = OnlineQueryParams.builder() + .withInput("user.id", List.of("U1")) + .withInput("user.sms.sms", hasManyInput) + .withOutputs("user.num_sms_events") + .build(); + + Map jsonHeader = BytesProducer.buildJsonHeader(params); + assertTrue(jsonHeader.containsKey("input_schema_hint")); + @SuppressWarnings("unchecked") + Map hints = (Map) jsonHeader.get("input_schema_hint"); + assertTrue(hints.containsKey("user.sms.sms")); + String hintValue = hints.get("user.sms.sms"); + // The hint should contain the FQN and bracketed column names + assertTrue(hintValue.startsWith("user.sms.sms[")); + assertTrue(hintValue.contains("user.sms.sms_id")); + assertTrue(hintValue.contains("user.sms.text")); + assertTrue(hintValue.endsWith("]")); + } + + @Test + public void testHasManyInputSchemaHintExplicit() throws Exception { + List>> emptyHasMany = new ArrayList<>(); + emptyHasMany.add(new ArrayList<>()); + + OnlineQueryParamsComplete params = OnlineQueryParams.builder() + .withInput("user.id", List.of("U1")) + .withInput("user.sms.sms", emptyHasMany) + .withOutputs("user.num_sms_events") + .withInputSchemaHint("user.sms.sms", "user.sms.sms[user.sms.sms_id,user.sms.text]") + .build(); + + Map jsonHeader = BytesProducer.buildJsonHeader(params); + assertTrue(jsonHeader.containsKey("input_schema_hint")); + @SuppressWarnings("unchecked") + Map hints = (Map) jsonHeader.get("input_schema_hint"); + assertEquals("user.sms.sms[user.sms.sms_id,user.sms.text]", hints.get("user.sms.sms")); + } + + @Test + public void testNoInputSchemaHintForScalars() throws Exception { + OnlineQueryParamsComplete params = OnlineQueryParams.builder() + .withInput("user.id", List.of("U1")) + .withInput("user.score", List.of(95.5)) + .withOutputs("user.today") + .build(); + + Map jsonHeader = BytesProducer.buildJsonHeader(params); + // No has-many inputs, so no hint should be generated + assertTrue(!jsonHeader.containsKey("input_schema_hint")); + } + + @Test + public void testSingleRowInputWithHasMany() throws Exception { + Map smsRow1 = new HashMap<>(); + smsRow1.put("user.sms.sms_id", "S1"); + smsRow1.put("user.sms.text", "hello"); + + Map smsRow2 = new HashMap<>(); + smsRow2.put("user.sms.sms_id", "S2"); + smsRow2.put("user.sms.text", "world"); + + var singleRowInput = new OnlineQueryParams.SingleRowInput() + .withInput("user.id", "U1") + .withHasManyInput("user.sms.sms", Arrays.asList(smsRow1, smsRow2)); + + var params = OnlineQueryParams.builder() + .withInput(singleRowInput) + .withOutputs("user.num_sms_events") + .build(); + + // The has-many should be wrapped as List.of(List.of(smsRow1, smsRow2)) + List hasManyValues = params.getInputs().get("user.sms.sms"); + assertEquals(1, hasManyValues.size()); + assertTrue(hasManyValues.get(0) instanceof List); + @SuppressWarnings("unchecked") + List> innerList = (List>) hasManyValues.get(0); + assertEquals(2, innerList.size()); + + BytesProducer.convertOnlineQueryParamsToBytes(params, allocator); + } + @Test public void testNUllValue() throws Exception { Map firstMap = new HashMap<>();