Skip to content

[Draft] Batch updated by key groups #303

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
suddendust wants to merge 4 commits into
base: main
Choose a base branch
from
+344 −58
Open

[Draft] Batch updated by key groups #303

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d869af5
Batch updated by key groups
suddendust Apr 30, 2026
fb9132c
Merge branch 'main' into feat/292
suddendust May 18, 2026
5db4adb
WIP
suddendust May 18, 2026
9e532a2
Merge branch 'main' into feat/292
suddendust May 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
198 changes: 181 additions & 17 deletions ...e/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2290,15 +2290,15 @@ void testSetAllFieldTypes() throws Exception {
SubDocumentUpdate.of("rating", 4.5f),
SubDocumentUpdate.of("weight", 123.456),
// Case 2: Top-level arrays
SubDocumentUpdate.of("tags", new String[] {"tag4", "tag5", "tag6"}),
SubDocumentUpdate.of("numbers", new Integer[] {10, 20, 30}),
SubDocumentUpdate.of("scores", new Double[] {1.1, 2.2, 3.3}),
SubDocumentUpdate.of("flags", new Boolean[] {true, false, true}),
SubDocumentUpdate.of("tags", new String[]{"tag4", "tag5", "tag6"}),
SubDocumentUpdate.of("numbers", new Integer[]{10, 20, 30}),
SubDocumentUpdate.of("scores", new Double[]{1.1, 2.2, 3.3}),
SubDocumentUpdate.of("flags", new Boolean[]{true, false, true}),
// Case 3 & 4: One nested path in JSONB (props) - tests nested primitive
SubDocumentUpdate.of("props.brand", "NewBrand"),
// Use 'sales' JSONB column for nested array test
SubDocumentUpdate.of(
"sales.regions", SubDocumentValue.of(new String[] {"US", "EU", "APAC"})));
"sales.regions", SubDocumentValue.of(new String[]{"US", "EU", "APAC"})));

UpdateOptions options =
UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build();
Expand Down Expand Up @@ -2510,7 +2510,7 @@ void testSetMultipleNestedPathsInSameJsonbColumn() throws Exception {
SubDocumentUpdate.of("props.size", "XL"),
SubDocumentUpdate.of("props.newField", "newValue"),
SubDocumentUpdate.of(
"props.owners", SubDocumentValue.of(new String[] {"owner1", "owner2"})));
"props.owners", SubDocumentValue.of(new String[]{"owner1", "owner2"})));

UpdateOptions options =
UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build();
Expand Down Expand Up @@ -2818,7 +2818,7 @@ void testAddArrayValue() {
SubDocumentUpdate.builder()
.subDocument("price")
.operator(UpdateOperator.ADD)
.subDocumentValue(SubDocumentValue.of(new Integer[] {1, 2, 3}))
.subDocumentValue(SubDocumentValue.of(new Integer[]{1, 2, 3}))
.build());

UpdateOptions options =
Expand Down Expand Up @@ -2865,19 +2865,19 @@ void testAppendToListAllCases() throws Exception {
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.APPEND_TO_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"newTag1", "newTag2"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"newTag1", "newTag2"}))
.build(),
// Nested JSONB array: append to existing props.colors
SubDocumentUpdate.builder()
.subDocument("props.colors")
.operator(UpdateOperator.APPEND_TO_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"green", "yellow"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"green", "yellow"}))
.build(),
// Nested JSONB: append to non-existent array (creates it)
SubDocumentUpdate.builder()
.subDocument("sales.regions")
.operator(UpdateOperator.APPEND_TO_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"US", "EU"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"US", "EU"}))
.build());

UpdateOptions options =
Expand Down Expand Up @@ -2956,13 +2956,13 @@ void testAddToListIfAbsentAllCases() throws Exception {
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT)
.subDocumentValue(SubDocumentValue.of(new String[] {"existing1", "newTag"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"existing1", "newTag"}))
.build(),
// Nested JSONB: 'red' exists, 'green' is new → adds only 'green'
SubDocumentUpdate.builder()
.subDocument("props.colors")
.operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT)
.subDocumentValue(SubDocumentValue.of(new String[] {"red", "green"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"red", "green"}))
.build());

UpdateOptions options =
Expand Down Expand Up @@ -3033,13 +3033,13 @@ void testRemoveAllFromListAllCases() throws Exception {
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.REMOVE_ALL_FROM_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"tag1"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"tag1"}))
.build(),
// Nested JSONB: remove 'red' and 'blue' → leaves green
SubDocumentUpdate.builder()
.subDocument("props.colors")
.operator(UpdateOperator.REMOVE_ALL_FROM_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"red", "blue"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"red", "blue"}))
.build());

UpdateOptions options =
Expand Down Expand Up @@ -3512,7 +3512,7 @@ void testBulkUpdateAllOperatorTypes() throws Exception {
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.APPEND_TO_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"newTag1", "newTag2"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"newTag1", "newTag2"}))
.build()));

updates.put(
Expand All @@ -3521,7 +3521,7 @@ void testBulkUpdateAllOperatorTypes() throws Exception {
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT)
.subDocumentValue(SubDocumentValue.of(new String[] {"hygiene", "uniqueTag"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"hygiene", "uniqueTag"}))
.build()));

updates.put(
Expand All @@ -3530,7 +3530,7 @@ void testBulkUpdateAllOperatorTypes() throws Exception {
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.REMOVE_ALL_FROM_LIST)
.subDocumentValue(SubDocumentValue.of(new String[] {"plastic"}))
.subDocumentValue(SubDocumentValue.of(new String[]{"plastic"}))
.build()));

BulkUpdateResult result = flatCollection.bulkUpdate(updates, UpdateOptions.builder().build());
Expand Down Expand Up @@ -3576,6 +3576,170 @@ void testBulkUpdateAllOperatorTypes() throws Exception {
}
}

@Test
@DisplayName(
"Should efficiently batch updates across multiple key groups with complex operations")
void testBulkUpdateMultipleGroupsComplexOperations() throws Exception {
Map<Key, java.util.Collection<SubDocumentUpdate>> updates = new LinkedHashMap<>();

// ===== Group 1: Top-level primitive + top-level array (3 keys: 1, 5, 8) =====
// All have item="Soap" - these should be batched together
// This tests: SET on primitive field, APPEND_TO_LIST on array field
List<SubDocumentUpdate> group1Updates =
List.of(
SubDocumentUpdate.of("price", 99), // SET operator (top-level primitive)
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.APPEND_TO_LIST)
.subDocumentValue(SubDocumentValue.of(new String[]{"updated-tag", "batch-test"}))
.build()); // APPEND_TO_LIST on top-level array

updates.put(rawKey("1"), group1Updates);
updates.put(rawKey("5"), group1Updates);
updates.put(rawKey("8"), group1Updates);

// ===== Group 2: Nested JSONB updates (2 keys: 3, 7) =====
// Both have props - these should be batched together
// This tests: SET on nested JSONB fields
List<SubDocumentUpdate> group2Updates =
List.of(
SubDocumentUpdate.builder()
.subDocument("props.brand")
.operator(UpdateOperator.SET)
.subDocumentValue(SubDocumentValue.of("PremiumBrand"))
.build(), // SET on nested JSONB primitive
SubDocumentUpdate.builder()
.subDocument("props.size")
.operator(UpdateOperator.SET)
.subDocumentValue(SubDocumentValue.of("XL"))
.build()); // SET on another nested field

updates.put(rawKey("3"), group2Updates);
updates.put(rawKey("7"), group2Updates);

// ===== Group 3: ADD operator + REMOVE_ALL_FROM_LIST (2 keys: 2, 6) =====
// Both have quantity and tags - these should be batched together
// This tests: ADD on numeric field, REMOVE_ALL_FROM_LIST on array
List<SubDocumentUpdate> group3Updates =
List.of(
SubDocumentUpdate.builder()
.subDocument("quantity")
.operator(UpdateOperator.ADD)
.subDocumentValue(SubDocumentValue.of(100))
.build(), // ADD to numeric field
SubDocumentUpdate.builder()
.subDocument("tags")
.operator(UpdateOperator.REMOVE_ALL_FROM_LIST)
.subDocumentValue(SubDocumentValue.of(new String[]{"glass", "plastic"}))
.build()); // REMOVE_ALL_FROM_LIST

updates.put(rawKey("2"), group3Updates);
updates.put(rawKey("6"), group3Updates);

// Execute bulk update - should have 3 groups with 2-3 keys each
BulkUpdateResult result = flatCollection.bulkUpdate(updates, UpdateOptions.builder().build());

// Total unique keys: 1, 2, 3, 5, 6, 7, 8 = 7 keys
assertEquals(7, result.getUpdatedCount(), "Should update 7 rows");

// Verify keys 1, 5, 8 have Group 1 updates (top-level primitive + array)
for (String id : List.of("1", "5", "8")) {
try (CloseableIterator<Document> iter = flatCollection.find(queryById(id))) {
assertTrue(iter.hasNext());
JsonNode json = OBJECT_MAPPER.readTree(iter.next().toJson());
assertEquals(99, json.get("price").asInt(), "Key " + id + " price should be 99");
JsonNode tags = json.get("tags");
List<String> tagList = new ArrayList<>();
tags.forEach(t -> tagList.add(t.asText()));
assertTrue(
tagList.contains("updated-tag"), "Key " + id + " should contain 'updated-tag'");
assertTrue(tagList.contains("batch-test"), "Key " + id + " should contain 'batch-test'");
}
}

// Verify keys 3, 7 have Group 2 updates (nested JSONB)
for (String id : List.of("3", "7")) {
try (CloseableIterator<Document> iter = flatCollection.find(queryById(id))) {
assertTrue(iter.hasNext());
JsonNode json = OBJECT_MAPPER.readTree(iter.next().toJson());
JsonNode props = json.get("props");
assertNotNull(props, "Key " + id + " should have props");
assertEquals(
"PremiumBrand",
props.get("brand").asText(),
"Key " + id + " brand should be updated");
assertEquals("XL", props.get("size").asText(), "Key " + id + " size should be XL");
}
}

// Verify keys 2, 6 have Group 3 updates (ADD + REMOVE_ALL_FROM_LIST)
try (CloseableIterator<Document> iter = flatCollection.find(queryById("2"))) {
assertTrue(iter.hasNext());
JsonNode json = OBJECT_MAPPER.readTree(iter.next().toJson());
assertEquals(101, json.get("quantity").asInt()); // 1 + 100
JsonNode tags = json.get("tags");
List<String> tagList = new ArrayList<>();
tags.forEach(t -> tagList.add(t.asText()));
assertFalse(tagList.contains("glass"), "Key 2 should not have 'glass' tag");
}

try (CloseableIterator<Document> iter = flatCollection.find(queryById("6"))) {
assertTrue(iter.hasNext());
JsonNode json = OBJECT_MAPPER.readTree(iter.next().toJson());
assertEquals(105, json.get("quantity").asInt()); // 5 + 100
JsonNode tags = json.get("tags");
List<String> tagList = new ArrayList<>();
tags.forEach(t -> tagList.add(t.asText()));
assertFalse(tagList.contains("plastic"), "Key 6 should not have 'plastic' tag");
}
}

@Test
@DisplayName(
"Should batch keys whose update shape matches by column:operator:path but whose value "
+ "arrays differ in length (nested JSONB REMOVE_ALL_FROM_LIST)")
void testBulkUpdateSameShapeDifferentParamCardinality() throws Exception {
Map<Key, java.util.Collection<SubDocumentUpdate>> updates = new LinkedHashMap<>();

updates.put(
rawKey("1"),
List.of(
SubDocumentUpdate.builder()
.subDocument("props.colors")
.operator(UpdateOperator.REMOVE_ALL_FROM_LIST)
.subDocumentValue(SubDocumentValue.of(new String[]{"Blue"}))
.build()));

updates.put(
rawKey("5"),
List.of(
SubDocumentUpdate.builder()
.subDocument("props.colors")
.operator(UpdateOperator.REMOVE_ALL_FROM_LIST)
.subDocumentValue(SubDocumentValue.of(new String[]{"Orange", "Blue"}))
.build()));

BulkUpdateResult result = flatCollection.bulkUpdate(updates, UpdateOptions.builder().build());

assertEquals(2, result.getUpdatedCount());

try (CloseableIterator<Document> iter = flatCollection.find(queryById("1"))) {
assertTrue(iter.hasNext());
JsonNode json = OBJECT_MAPPER.readTree(iter.next().toJson());
JsonNode colors = json.get("props").get("colors");
List<String> colorList = new ArrayList<>();
colors.forEach(c -> colorList.add(c.asText()));
assertEquals(List.of("Green"), colorList);
}

try (CloseableIterator<Document> iter = flatCollection.find(queryById("5"))) {
assertTrue(iter.hasNext());
JsonNode json = OBJECT_MAPPER.readTree(iter.next().toJson());
JsonNode colors = json.get("props").get("colors");
assertEquals(0, colors.size());
}
}

@Test
@DisplayName("Should handle edge cases: empty map, null map, non-existent keys")
void testBulkUpdateEdgeCases() throws Exception {
Expand Down
Loading
Loading