Merge pull request #8107 from dolthub/nicktobey/json-set

Optimize JSON_SET and JSON_REPLACE on `IndexedJsonDocument`
dolthub · Jul 10, 2024 · 00add71 · 00add71
2 parents 586ef39 + d7cb874
commit 00add71
Show file tree

Hide file tree

Showing 3 changed files with 139 additions and 10 deletions.
diff --git a/go/store/prolly/tree/json_chunker.go b/go/store/prolly/tree/json_chunker.go
@@ -130,7 +130,7 @@ func (j *JsonChunker) Done(ctx context.Context) (Node, error) {
 	// When inserting into the beginning of an object or array, we need to add an extra comma.
 	// We could track then in the chunker, but it's easier to just check the next part of JSON to determine
 	// whether we need the comma.
-	if j.jScanner.currentPath.getScannerState() == endOfValue && jsonBytes[0] != '}' && jsonBytes[0] != ']' && jsonBytes[0] != ',' {
+	if j.jScanner.currentPath.getScannerState() == endOfValue && len(jsonBytes) > 0 && jsonBytes[0] != '}' && jsonBytes[0] != ']' && jsonBytes[0] != ',' {
 		j.appendJsonToBuffer([]byte(","))
 	}
 	// Append the rest of the JsonCursor, then continue until we either exhaust the cursor, or we coincide with a boundary from the original tree.

diff --git a/go/store/prolly/tree/json_indexed_document.go b/go/store/prolly/tree/json_indexed_document.go
@@ -144,7 +144,6 @@ func (i IndexedJsonDocument) Lookup(ctx context.Context, pathString string) (res
 }
 
 func (i IndexedJsonDocument) tryLookup(ctx context.Context, pathString string) (sql.JSONWrapper, error) {
-
 	path, err := jsonPathElementsFromMySQLJsonPath([]byte(pathString))
 	if err != nil {
 		return nil, err
@@ -386,22 +385,130 @@ func (i IndexedJsonDocument) tryRemove(ctx context.Context, path string) (types.
 	return NewIndexedJsonDocument(ctx, newRoot, i.m.NodeStore), true, nil
 }
 
-// Set is not yet implemented, so we call it on a types.JSONDocument instead.
-func (i IndexedJsonDocument) Set(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) {
-	v, err := i.ToInterface()
+// Set implements types.MutableJSON
+func (i IndexedJsonDocument) Set(ctx context.Context, path string, val sql.JSONWrapper) (result types.MutableJSON, changed bool, err error) {
+	err = tryWithFallback(
+		ctx,
+		i,
+		func() error {
+			result, changed, err = i.trySet(ctx, path, val)
+			return err
+		},
+		func(jsonDocument types.JSONDocument) error {
+			result, changed, err = jsonDocument.Set(ctx, path, val)
+			return err
+		})
+	return result, changed, err
+}
+
+func (i IndexedJsonDocument) trySet(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) {
+	keyPath, err := jsonPathElementsFromMySQLJsonPath([]byte(path))
 	if err != nil {
 		return nil, false, err
 	}
-	return types.JSONDocument{Val: v}.Set(ctx, path, val)
+
+	jsonCursor, found, err := newJsonCursor(ctx, i.m.NodeStore, i.m.Root, keyPath, false)
+	if err != nil {
+		return nil, false, err
+	}
+
+	// The supplied path may be 0-indexing into a scalar, which is the same as referencing the scalar. Remove
+	// the index and try again.
+	for !found && keyPath.size() > jsonCursor.jsonScanner.currentPath.size() {
+		lastKeyPathElement := keyPath.getLastPathElement()
+		if !lastKeyPathElement.isArrayIndex || lastKeyPathElement.getArrayIndex() != 0 {
+			// The key does not exist in the document.
+			break
+		}
+
+		keyPath.pop()
+		found = compareJsonLocations(keyPath, jsonCursor.jsonScanner.currentPath) == 0
+	}
+
+	if found {
+		return i.replaceIntoCursor(ctx, keyPath, jsonCursor, val)
+	} else {
+		return i.insertIntoCursor(ctx, keyPath, jsonCursor, val)
+	}
 }
 
-// Replace is not yet implemented, so we call it on a types.JSONDocument instead.
-func (i IndexedJsonDocument) Replace(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) {
-	v, err := i.ToInterface()
+// Replace implements types.MutableJSON
+func (i IndexedJsonDocument) Replace(ctx context.Context, path string, val sql.JSONWrapper) (result types.MutableJSON, changed bool, err error) {
+	err = tryWithFallback(
+		ctx,
+		i,
+		func() error {
+			result, changed, err = i.tryReplace(ctx, path, val)
+			return err
+		},
+		func(jsonDocument types.JSONDocument) error {
+			result, changed, err = jsonDocument.Replace(ctx, path, val)
+			return err
+		})
+	return result, changed, err
+}
+
+func (i IndexedJsonDocument) tryReplace(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) {
+	keyPath, err := jsonPathElementsFromMySQLJsonPath([]byte(path))
+	if err != nil {
+		return nil, false, err
+	}
+
+	jsonCursor, found, err := newJsonCursor(ctx, i.m.NodeStore, i.m.Root, keyPath, false)
+	if err != nil {
+		return nil, false, err
+	}
+
+	// The supplied path may be 0-indexing into a scalar, which is the same as referencing the scalar. Remove
+	// the index and try again.
+	for !found && keyPath.size() > jsonCursor.jsonScanner.currentPath.size() {
+		lastKeyPathElement := keyPath.getLastPathElement()
+		if !lastKeyPathElement.isArrayIndex || lastKeyPathElement.getArrayIndex() != 0 {
+			// The key does not exist in the document.
+			return i, false, nil
+		}
+
+		keyPath.pop()
+		found = compareJsonLocations(keyPath, jsonCursor.jsonScanner.currentPath) == 0
+	}
+
+	if !found {
+		// The key does not exist in the document.
+		return i, false, nil
+	}
+
+	return i.replaceIntoCursor(ctx, keyPath, jsonCursor, val)
+}
+
+func (i IndexedJsonDocument) replaceIntoCursor(ctx context.Context, keyPath jsonLocation, jsonCursor *JsonCursor, val sql.JSONWrapper) (types.MutableJSON, bool, error) {
+
+	// The cursor is now pointing to the start of the value being replaced.
+	jsonChunker, err := newJsonChunker(ctx, jsonCursor, i.m.NodeStore)
 	if err != nil {
 		return nil, false, err
 	}
-	return types.JSONDocument{Val: v}.Replace(ctx, path, val)
+
+	// Advance the cursor to the end of the value being removed.
+	keyPath.setScannerState(endOfValue)
+	_, err = jsonCursor.AdvanceToLocation(ctx, keyPath, false)
+	if err != nil {
+		return nil, false, err
+	}
+
+	insertedValueBytes, err := types.MarshallJson(val)
+	if err != nil {
+		return nil, false, err
+	}
+
+	jsonChunker.appendJsonToBuffer(insertedValueBytes)
+	jsonChunker.processBuffer(ctx)
+
+	newRoot, err := jsonChunker.Done(ctx)
+	if err != nil {
+		return nil, false, err
+	}
+
+	return NewIndexedJsonDocument(ctx, newRoot, i.m.NodeStore), true, nil
 }
 
 // ArrayInsert is not yet implemented, so we call it on a types.JSONDocument instead.

diff --git a/go/store/prolly/tree/json_indexed_document_test.go b/go/store/prolly/tree/json_indexed_document_test.go
@@ -259,6 +259,28 @@ func TestIndexedJsonDocument_Extract(t *testing.T) {
 	jsontests.RunJsonTests(t, testCases)
 }
 
+func TestIndexedJsonDocument_Replace(t *testing.T) {
+	ctx := context.Background()
+	ns := NewTestNodeStore()
+	convertToIndexedJsonDocument := func(t *testing.T, s interface{}) interface{} {
+		return newIndexedJsonDocumentFromValue(t, ctx, ns, s)
+	}
+
+	testCases := jsontests.JsonReplaceTestCases(t, convertToIndexedJsonDocument)
+	jsontests.RunJsonTests(t, testCases)
+}
+
+func TestIndexedJsonDocument_Set(t *testing.T) {
+	ctx := context.Background()
+	ns := NewTestNodeStore()
+	convertToIndexedJsonDocument := func(t *testing.T, s interface{}) interface{} {
+		return newIndexedJsonDocumentFromValue(t, ctx, ns, s)
+	}
+
+	testCases := jsontests.JsonSetTestCases(t, convertToIndexedJsonDocument)
+	jsontests.RunJsonTests(t, testCases)
+}
+
 func TestIndexedJsonDocument_Value(t *testing.T) {
 	ctx := context.Background()
 	ns := NewTestNodeStore()