Skip to content

Commit 7627ddc

Browse files
thiagohoraclaude
andauthored
[OPIK-8023] Bind dataset-item JSON sort keys as query parameters (#7935)
* [OPIK-8023] Bind dataset-item JSON sort keys as query parameters Dataset-item sorting by output.*/input.*/metadata.* built the JSONExtractRaw key by interpolating the key text directly into the SQL expression, whereas data.* keys are passed as bound query parameters. Make the JSON sort keys consistent with data.* by binding the key as a ClickHouse parameter (JSONExtractRaw(col, :param)). This also handles keys containing special characters correctly. - FilterQueryBuilder: template uses a bound placeholder; bind key via field.bindKey() - SortingFactoryDatasets: JSON fields go through the standard dynamic bindKeyParam path - SortingQueryBuilder: bind all dynamic keys; remove now-unused overloads - DatasetItemDAO: use single-arg hasDynamicKeys/bindDynamicKeys - Add integration test covering JSON sort keys containing special characters Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [OPIK-8023] Address review feedback on dataset-item sort key binding - SortingField: bindKeyParam feeds a SQL parameter placeholder name, so generate it server-side only -- it is no longer read from request input (READ_ONLY) and is always a safe identifier. This keeps the placeholder name safe for data.* and experiment_scores.* too, which build namespace[:sorting_param_<bindKeyParam>]. - Strengthen the dataset-item sort test: seed distinct values under the requested JSON key and assert the full ordered result for both ASC and DESC, comparing whole objects. - Add SortingFactoryTest coverage asserting bindKeyParam is server-generated, not taken from input. - Fix stale comments in FilterQueryBuilder and SortableFields to match the bound-parameter approach. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [OPIK-8023] Add version-DAO JSON sort coverage and document nested-key behavior - DatasetVersionResourceTest: add a regression test that sorts versioned experiment items by a JSON output key (with special characters) through the push-top-limit path, which was previously uncovered (only total_estimated_cost and usage.total_tokens were tested). - FilterQueryBuilder: note that the JSON sort key is the segment after the first dot, i.e. a single top-level key ("output.a.b" looks up "a.b"); nested traversal is not performed. This matches the previous behavior. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [OPIK-8023] Mark bindKeyParam @JsonIgnore (internal, off the JSON API surface) bindKeyParam is an internal value used only to build the SQL parameter placeholder name. Use @JsonIgnore instead of READ_ONLY so it is neither serialized nor deserialized; it is always (re)generated server-side as a safe identifier in the canonical constructor. Confirmed the frontend and Python/TS SDKs only send field/direction, so this is transparent. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [OPIK-8023] Rely on SortingField constructor for bindKeyParam; annotate @nullable - SortingFactoryDatasets: remove processFields/ensureBindKeyParam. The canonical constructor now guarantees a safe bindKeyParam for dynamic fields, so the factory-side normalization was duplicate work (and left two implementations to maintain). This also aligns the datasets factory with the other sorting factories, which never overrode processFields. - SortingField: annotate bindKeyParam @nullable and document the contract (null for static fields, a server-generated safe identifier for dynamic fields). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [OPIK-8023] Parameterize version-DAO JSON sort test over namespace and direction Replace the single output.* test with a @MethodSource-parameterized test covering output.*, input.* and metadata.* (with special-character keys) for both ASC and DESC, asserting the full expected item order via whole-object comparison. This exercises the separate input/metadata push-top-limit CTE bindings, and uses SortingField.builder() instead of the positional constructor. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent fa282b0 commit 7627ddc

9 files changed

Lines changed: 289 additions & 80 deletions

File tree

apps/opik-backend/src/main/java/com/comet/opik/api/sorting/SortableFields.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,9 @@ public class SortableFields {
5151
public static final String INSTRUCTIONS = "instructions";
5252
public static final String WEBHOOK_URL = "webhook_url";
5353
public static final String DATA = "data.*"; // Truly dynamic - uses Map with parameter binding
54-
public static final String OUTPUT_WILDCARD = "output.*"; // JSON fields - use JSONExtractRaw, not parameter binding
55-
public static final String INPUT_WILDCARD = "input.*"; // JSON fields - use JSONExtractRaw, not parameter binding
56-
public static final String METADATA_WILDCARD = "metadata.*"; // JSON fields - use JSONExtractRaw, not parameter binding (metadata field already exists above)
54+
public static final String OUTPUT_WILDCARD = "output.*"; // JSON fields - use JSONExtractRaw with the key bound as a parameter
55+
public static final String INPUT_WILDCARD = "input.*"; // JSON fields - use JSONExtractRaw with the key bound as a parameter
56+
public static final String METADATA_WILDCARD = "metadata.*"; // JSON fields - use JSONExtractRaw with the key bound as a parameter (metadata field already exists above)
5757
public static final String COMMENTS = "comments";
5858
public static final String EXPERIMENT_ID = "experiment_id";
5959
public static final String PASS_RATE = "pass_rate";

apps/opik-backend/src/main/java/com/comet/opik/api/sorting/SortingFactoryDatasets.java

Lines changed: 2 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
11
package com.comet.opik.api.sorting;
22

33
import org.apache.commons.collections4.CollectionUtils;
4-
import org.apache.commons.lang3.StringUtils;
54

65
import java.util.List;
76
import java.util.Set;
8-
import java.util.UUID;
97

108
import static com.comet.opik.api.sorting.SortableFields.COMMENTS;
119
import static com.comet.opik.api.sorting.SortableFields.CREATED_AT;
@@ -26,9 +24,6 @@
2624
import static com.comet.opik.api.sorting.SortableFields.TAGS;
2725
import static com.comet.opik.api.sorting.SortableFields.TOTAL_ESTIMATED_COST;
2826
import static com.comet.opik.api.sorting.SortableFields.USAGE;
29-
import static com.comet.opik.domain.filter.FilterQueryBuilder.INPUT_FIELD_PREFIX;
30-
import static com.comet.opik.domain.filter.FilterQueryBuilder.METADATA_FIELD_PREFIX;
31-
import static com.comet.opik.domain.filter.FilterQueryBuilder.OUTPUT_FIELD_PREFIX;
3227

3328
public class SortingFactoryDatasets extends SortingFactory {
3429

@@ -88,41 +83,6 @@ private boolean matchesSupported(String field, Set<String> supportedSet) {
8883
return false;
8984
}
9085

91-
@Override
92-
protected List<SortingField> processFields(List<SortingField> sorting) {
93-
// Ensure dynamic fields have bindKeyParam set (needed after JSON deserialization)
94-
return sorting.stream()
95-
.map(this::ensureBindKeyParam)
96-
.toList();
97-
}
98-
99-
private SortingField ensureBindKeyParam(SortingField sortingField) {
100-
String field = sortingField.field();
101-
102-
// JSON fields (output.*, input.*, metadata.*) should NOT be treated as dynamic
103-
// because they use JSONExtractRaw with literal keys in DatasetItemDAO
104-
if (field.startsWith(OUTPUT_FIELD_PREFIX) || field.startsWith(INPUT_FIELD_PREFIX)
105-
|| field.startsWith(METADATA_FIELD_PREFIX)) {
106-
return sortingField.toBuilder()
107-
.bindKeyParam(null)
108-
.build();
109-
}
110-
111-
// Only dynamic fields need bindKeyParam
112-
if (!sortingField.isDynamic()) {
113-
return sortingField;
114-
}
115-
116-
// If bindKeyParam is already set, return as-is
117-
String bindKeyParam = sortingField.bindKeyParam();
118-
if (StringUtils.isNotBlank(bindKeyParam)) {
119-
return sortingField;
120-
}
121-
122-
// Generate UUID for dynamic field
123-
bindKeyParam = UUID.randomUUID().toString().replace("-", "");
124-
return sortingField.toBuilder()
125-
.bindKeyParam(bindKeyParam)
126-
.build();
127-
}
86+
// Note: bindKeyParam is generated in SortingField's canonical constructor (a safe identifier for
87+
// dynamic fields), so no factory-side normalization is needed here.
12888
}

apps/opik-backend/src/main/java/com/comet/opik/api/sorting/SortingField.java

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,36 @@
11
package com.comet.opik.api.sorting;
22

3+
import com.fasterxml.jackson.annotation.JsonIgnore;
34
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
45
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
56
import com.fasterxml.jackson.databind.annotation.JsonNaming;
7+
import jakarta.annotation.Nullable;
68
import jakarta.validation.constraints.NotBlank;
79
import lombok.Builder;
810

911
import java.util.UUID;
12+
import java.util.regex.Pattern;
1013

1114
@Builder(toBuilder = true)
1215
@JsonIgnoreProperties(ignoreUnknown = true)
1316
@JsonNaming(PropertyNamingStrategies.SnakeCaseStrategy.class)
1417
public record SortingField(
1518
@NotBlank String field,
1619
Direction direction,
17-
String bindKeyParam) {
20+
// bindKeyParam feeds a SQL parameter placeholder name (see bindKey()), so it is an internal
21+
// value derived server-side only: @JsonIgnore keeps it off the JSON API surface entirely (neither
22+
// serialized nor deserialized). It is null for static (non-dynamic) fields and a server-generated
23+
// safe identifier for dynamic fields (see the canonical constructor).
24+
@JsonIgnore @Nullable String bindKeyParam) {
1825

19-
// Canonical constructor with auto-generation of bindKeyParam for dynamic fields
26+
private static final Pattern SAFE_BIND_KEY_PARAM = Pattern.compile("[A-Za-z0-9_]+");
27+
28+
// Canonical constructor. bindKeyParam is rendered into a SQL placeholder name, so for dynamic
29+
// fields it must always be a server-generated safe identifier. Regenerate it whenever it is
30+
// missing or is not a safe identifier, so no client-influenced value can reach the SQL.
2031
public SortingField {
21-
// Auto-generate bindKeyParam for dynamic fields if not provided
22-
if (bindKeyParam == null && field != null && field.contains(".")) {
32+
if (field != null && field.contains(".")
33+
&& (bindKeyParam == null || !SAFE_BIND_KEY_PARAM.matcher(bindKeyParam).matches())) {
2334
bindKeyParam = UUID.randomUUID().toString().replace("-", "");
2435
}
2536
}

apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetItemDAO.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1392,8 +1392,7 @@ public Mono<DatasetItemPage> getItems(
13921392
}
13931393

13941394
var hasDynamicKeys = datasetItemSearchCriteria.sortingFields() != null
1395-
&& sortingQueryBuilder.hasDynamicKeys(datasetItemSearchCriteria.sortingFields(),
1396-
itemFieldMapping);
1395+
&& sortingQueryBuilder.hasDynamicKeys(datasetItemSearchCriteria.sortingFields());
13971396

13981397
var selectStatement = connection.createStatement(finalTemplate.render())
13991398
.bind("datasetId", datasetItemSearchCriteria.datasetId())
@@ -1411,7 +1410,7 @@ public Mono<DatasetItemPage> getItems(
14111410
// Bind dynamic sorting keys if present
14121411
if (hasDynamicKeys) {
14131412
selectStatement = sortingQueryBuilder.bindDynamicKeys(selectStatement,
1414-
datasetItemSearchCriteria.sortingFields(), itemFieldMapping);
1413+
datasetItemSearchCriteria.sortingFields());
14151414
}
14161415

14171416
bindSearchCriteria(datasetItemSearchCriteria, selectStatement);

apps/opik-backend/src/main/java/com/comet/opik/domain/filter/FilterQueryBuilder.java

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,10 @@ public class FilterQueryBuilder {
5353

5454
public static final String JSONPATH_ROOT = "$";
5555

56-
private static final String JSON_EXTRACT_RAW_TEMPLATE = "JSONExtractRaw(%s, '%s')";
56+
// The JSON key is passed as a bound query parameter (:sorting_param_xxx) rather than interpolated,
57+
// consistent with the data.* sort path and robust to keys containing special characters.
58+
// See buildDatasetItemFieldMapping.
59+
private static final String JSON_EXTRACT_RAW_TEMPLATE = "JSONExtractRaw(%s, :%s)";
5760
public static final String OUTPUT_FIELD_PREFIX = "output.";
5861
public static final String INPUT_FIELD_PREFIX = "input.";
5962
public static final String METADATA_FIELD_PREFIX = "metadata.";
@@ -1395,8 +1398,8 @@ private static String getKey(Filter filter) {
13951398
/**
13961399
* Builds field mapping for DatasetItem JSON fields (output, input, metadata).
13971400
* These fields are stored as JSON strings in ClickHouse, so we need to use JSONExtractRaw
1398-
* instead of bracket notation. We use literal keys instead of bind parameters
1399-
* to avoid the dynamic field tuple wrapping.
1401+
* instead of bracket notation. The JSON key is bound as a query parameter rather than
1402+
* interpolated, consistent with the data.* sort path.
14001403
* <p>
14011404
* This is used for sorting DatasetItem fields.
14021405
*
@@ -1409,20 +1412,23 @@ public Map<String, String> buildDatasetItemFieldMapping(@NonNull List<SortingFie
14091412
for (SortingField field : sortingFields) {
14101413
String fieldName = field.field();
14111414

1412-
// Check if this is a JSON field (output, input, or metadata)
1413-
// Use literal keys instead of bind parameters to avoid dynamic field handling
1415+
// Check if this is a JSON field (output, input, or metadata).
1416+
// The JSON key is bound as a query parameter (field.bindKey() -> :sorting_param_xxx) rather
1417+
// than interpolated into the SQL. The key VALUE (field.dynamicKey()) is bound later in
1418+
// SortingQueryBuilder.bindDynamicKeys(), consistent with the data.* path, so keys containing
1419+
// special characters are handled correctly.
1420+
// Note: dynamicKey() is everything after the first dot, i.e. a single top-level JSON key
1421+
// ("output.a.b" looks up the key "a.b"); nested traversal is not performed. This matches the
1422+
// previous behavior.
14141423
if (fieldName.startsWith(OUTPUT_FIELD_PREFIX)) {
1415-
String key = fieldName.substring(OUTPUT_FIELD_PREFIX.length());
14161424
fieldMapping.put(fieldName,
1417-
JSON_EXTRACT_RAW_TEMPLATE.formatted("output", key));
1425+
JSON_EXTRACT_RAW_TEMPLATE.formatted("output", field.bindKey()));
14181426
} else if (fieldName.startsWith(INPUT_FIELD_PREFIX)) {
1419-
String key = fieldName.substring(INPUT_FIELD_PREFIX.length());
14201427
fieldMapping.put(fieldName,
1421-
JSON_EXTRACT_RAW_TEMPLATE.formatted("input", key));
1428+
JSON_EXTRACT_RAW_TEMPLATE.formatted("input", field.bindKey()));
14221429
} else if (fieldName.startsWith(METADATA_FIELD_PREFIX)) {
1423-
String key = fieldName.substring(METADATA_FIELD_PREFIX.length());
14241430
fieldMapping.put(fieldName,
1425-
JSON_EXTRACT_RAW_TEMPLATE.formatted("metadata", key));
1431+
JSON_EXTRACT_RAW_TEMPLATE.formatted("metadata", field.bindKey()));
14261432
}
14271433
// For other fields (including feedback_scores, data, etc.), use default dbField()
14281434
}

apps/opik-backend/src/main/java/com/comet/opik/domain/sorting/SortingQueryBuilder.java

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -70,28 +70,18 @@ private String transformNullDirection(SortingField sortingField) {
7070
}
7171

7272
public boolean hasDynamicKeys(@NonNull List<SortingField> sorting) {
73-
return hasDynamicKeys(sorting, null);
74-
}
75-
76-
public boolean hasDynamicKeys(@NonNull List<SortingField> sorting, Map<String, String> fieldMapping) {
77-
// Only fields with bindKeyParam need dynamic binding.
78-
// Fields in the fieldMapping use literal SQL expressions (e.g. JSONExtractRaw), so no bind param exists.
73+
// A dynamic field is bound whenever it carries a bindKeyParam. This covers JSON fields
74+
// (output/input/metadata), whose key is bound via JSONExtractRaw(col, :param), as well as
75+
// map-style fields such as data[:param] and experiment_scores_agg[:param].
7976
return sorting.stream()
8077
.filter(SortingField::isDynamic)
81-
.filter(field -> field.bindKeyParam() != null)
82-
.anyMatch(field -> fieldMapping == null || !fieldMapping.containsKey(field.field()));
83-
}
84-
85-
public Statement bindDynamicKeys(Statement statement, List<SortingField> sorting) {
86-
return bindDynamicKeys(statement, sorting, null);
78+
.anyMatch(field -> field.bindKeyParam() != null);
8779
}
8880

89-
public Statement bindDynamicKeys(Statement statement, List<SortingField> sorting,
90-
Map<String, String> fieldMapping) {
81+
public Statement bindDynamicKeys(Statement statement, @NonNull List<SortingField> sorting) {
9182
sorting.stream()
9283
.filter(SortingField::isDynamic)
9384
.filter(sortingField -> sortingField.bindKeyParam() != null)
94-
.filter(sortingField -> fieldMapping == null || !fieldMapping.containsKey(sortingField.field()))
9585
.forEach(sortingField -> {
9686
try {
9787
statement.bind(sortingField.bindKey(), sortingField.dynamicKey());

apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/DatasetVersionResourceTest.java

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3510,6 +3510,89 @@ void sortByUsageTotalTokens__whenAscendingOrder__thenReturnSortedByAverage() {
35103510
datasetItems.get(0).id());
35113511
}
35123512

3513+
static Stream<Arguments> sortByJsonKeyThroughPushTopLimit() {
3514+
// namespace, JSON key (with special characters), direction, expected item-index order
3515+
return Stream.of(
3516+
Arguments.of("output", "score's", Direction.ASC, List.of(0, 1, 2)),
3517+
Arguments.of("output", "score's", Direction.DESC, List.of(2, 1, 0)),
3518+
Arguments.of("input", "in\"put", Direction.ASC, List.of(0, 1, 2)),
3519+
Arguments.of("input", "in\"put", Direction.DESC, List.of(2, 1, 0)),
3520+
Arguments.of("metadata", "me\\ta", Direction.ASC, List.of(0, 1, 2)),
3521+
Arguments.of("metadata", "me\\ta", Direction.DESC, List.of(2, 1, 0)));
3522+
}
3523+
3524+
@ParameterizedTest
3525+
@MethodSource
3526+
@DisplayName("should sort versioned experiment items by a JSON key (output/input/metadata) through the push-top-limit path")
3527+
void sortByJsonKeyThroughPushTopLimit(String namespace, String jsonKey, Direction direction,
3528+
List<Integer> expectedIndexOrder) {
3529+
var datasetName = UUID.randomUUID().toString();
3530+
var datasetId = createDataset(datasetName);
3531+
int count = 3;
3532+
createDatasetItems(datasetId, count);
3533+
3534+
var version = getLatestVersion(datasetId);
3535+
var datasetItems = datasetResourceClient.getDatasetItems(
3536+
datasetId, 1, 10, DatasetVersionService.LATEST_TAG, API_KEY, TEST_WORKSPACE).content();
3537+
3538+
var projectName = UUID.randomUUID().toString();
3539+
3540+
// One trial per dataset item; each trace carries a distinct value under the requested key (which
3541+
// contains special characters) in the requested namespace. Sorting by <namespace>.<key> routes
3542+
// through the push-top-limit path, whose JSON expression binds the key as a parameter
3543+
// (JSONExtractRaw(argMax(...), :param)).
3544+
var traceIds = IntStream.range(0, count)
3545+
.mapToObj(i -> {
3546+
var value = JsonUtils.valueToTree(Map.of(jsonKey, (i + 1) * 10));
3547+
var builder = factory.manufacturePojo(Trace.class).toBuilder().projectName(projectName);
3548+
switch (namespace) {
3549+
case "output" -> builder.output(value);
3550+
case "input" -> builder.input(value);
3551+
case "metadata" -> builder.metadata(value);
3552+
default -> throw new IllegalStateException("Unexpected namespace: " + namespace);
3553+
}
3554+
var trace = builder.build();
3555+
traceResourceClient.createTrace(trace, API_KEY, TEST_WORKSPACE);
3556+
return trace.id();
3557+
})
3558+
.toList();
3559+
3560+
var experiment = experimentResourceClient.createPartialExperiment()
3561+
.datasetName(datasetName)
3562+
.datasetVersionId(version.id())
3563+
.build();
3564+
var experimentId = experimentResourceClient.create(experiment, API_KEY, TEST_WORKSPACE);
3565+
3566+
IntStream.range(0, count).forEach(i -> {
3567+
var item = factory.manufacturePojo(ExperimentItem.class).toBuilder()
3568+
.experimentId(experimentId)
3569+
.datasetItemId(datasetItems.get(i).id())
3570+
.traceId(traceIds.get(i))
3571+
.build();
3572+
experimentResourceClient.createExperimentItem(Set.of(item), API_KEY, TEST_WORKSPACE);
3573+
});
3574+
3575+
// Baseline fetch (no sorting) captures the full objects as returned; the assertion only tests order.
3576+
var baseline = datasetResourceClient.getDatasetItemsWithExperimentItems(
3577+
datasetId, List.of(experimentId), null, null, null, API_KEY, TEST_WORKSPACE).content();
3578+
assertThat(baseline).hasSize(count);
3579+
3580+
Map<UUID, DatasetItem> baselineById = baseline.stream()
3581+
.collect(Collectors.toMap(DatasetItem::id, item -> item));
3582+
List<DatasetItem> expected = expectedIndexOrder.stream()
3583+
.map(i -> baselineById.get(datasetItems.get(i).id()))
3584+
.toList();
3585+
3586+
var sorting = List.of(SortingField.builder().field(namespace + "." + jsonKey).direction(direction).build());
3587+
var sorted = datasetResourceClient.getDatasetItemsWithExperimentItems(
3588+
datasetId, List.of(experimentId), null, null, sorting, API_KEY, TEST_WORKSPACE);
3589+
3590+
// Compare the whole DatasetItem objects, in order - not just their ids.
3591+
assertThat(sorted.content())
3592+
.usingRecursiveFieldByFieldElementComparatorIgnoringFields(IGNORED_FIELDS_DATA_ITEM)
3593+
.containsExactlyElementsOf(expected);
3594+
}
3595+
35133596
@Test
35143597
@DisplayName("Success: PUT /items without query param returns 204 (backward compatibility)")
35153598
void putItems_whenRespondWithLatestVersionNotSet_thenReturnsNoContent() {

0 commit comments

Comments
 (0)