util/json: add AllPathsWithDepth

Previously, `AllPaths` within the `json` package would search the entire
JSON document until a leaf is found. This PR adds another function,
`AllPathsWithDepth` which limits how deep into the JSON document to
search. This will be helpful for `jsonb_path_query`, as wildcard
accessors (`[*]`) need to search for all paths one level deep.

Epic: None
Release note: None

pkg/util/json/encoded.go

@@ -797,12 +797,12 @@ func (j *jsonEncoded) numInvertedIndexEntries() (int, error) {
 	return decoded.numInvertedIndexEntries()
 }
 
-func (j *jsonEncoded) allPaths() ([]JSON, error) {
+func (j *jsonEncoded) allPathsWithDepth(depth int) ([]JSON, error) {
 	decoded, err := j.decode()
 	if err != nil {
 		return nil, err
 	}
-	return decoded.allPaths()
+	return decoded.allPathsWithDepth(depth)
 }
 
 // HasContainerLeaf implements the JSON interface.

pkg/util/json/json.go

@@ -136,10 +136,12 @@ type JSON interface {
 	// produced if this JSON gets included in an inverted index.
 	numInvertedIndexEntries() (int, error)
 
-	// allPaths returns a slice of new JSON documents, each a path to a leaf
-	// through the receiver. Note that leaves include the empty object and array
-	// in addition to scalars.
-	allPaths() ([]JSON, error)
+	// allPathsWithDepth returns a slice of new JSON documents, each a path
+	// through the receiver. The depth parameter specifies the maximum depth of
+	// the paths to return. If the depth is negative, all paths of any depth are
+	// returned. If the depth is 0, the receiver itself is returned. Note that
+	// leaves include the empty object and array in addition to scalars.
+	allPathsWithDepth(depth int) ([]JSON, error)
 
 	// FetchValKey implements the `->` operator for strings, returning nil if the
 	// key is not found.
@@ -1717,36 +1719,45 @@ func (j jsonObject) numInvertedIndexEntries() (int, error) {
 // through the input. Note that leaves include the empty object and array
 // in addition to scalars.
 func AllPaths(j JSON) ([]JSON, error) {
-	return j.allPaths()
+	return j.allPathsWithDepth(-1)
 }
 
-func (j jsonNull) allPaths() ([]JSON, error) {
+// AllPathsWithDepth returns a slice of new JSON documents, each a path
+// through the receiver. The depth parameter specifies the maximum depth of
+// the paths to return. If the depth is negative, all paths of any depth are
+// returned. If the depth is 0, the receiver itself is returned. Note that
+// leaves include the empty object and array in addition to scalars.
+func AllPathsWithDepth(j JSON, depth int) ([]JSON, error) {
+	return j.allPathsWithDepth(depth)
+}
+
+func (j jsonNull) allPathsWithDepth(depth int) ([]JSON, error) {
 	return []JSON{j}, nil
 }
 
-func (j jsonTrue) allPaths() ([]JSON, error) {
+func (j jsonTrue) allPathsWithDepth(depth int) ([]JSON, error) {
 	return []JSON{j}, nil
 }
 
-func (j jsonFalse) allPaths() ([]JSON, error) {
+func (j jsonFalse) allPathsWithDepth(depth int) ([]JSON, error) {
 	return []JSON{j}, nil
 }
 
-func (j jsonString) allPaths() ([]JSON, error) {
+func (j jsonString) allPathsWithDepth(depth int) ([]JSON, error) {
 	return []JSON{j}, nil
 }
 
-func (j jsonNumber) allPaths() ([]JSON, error) {
+func (j jsonNumber) allPathsWithDepth(depth int) ([]JSON, error) {
 	return []JSON{j}, nil
 }
 
-func (j jsonArray) allPaths() ([]JSON, error) {
-	if len(j) == 0 {
+func (j jsonArray) allPathsWithDepth(depth int) ([]JSON, error) {
+	if len(j) == 0 || depth == 0 {
 		return []JSON{j}, nil
 	}
 	ret := make([]JSON, 0, len(j))
 	for i := range j {
-		paths, err := j[i].allPaths()
+		paths, err := j[i].allPathsWithDepth(depth - 1)
 		if err != nil {
 			return nil, err
 		}
@@ -1757,13 +1768,13 @@ func (j jsonArray) allPaths() ([]JSON, error) {
 	return ret, nil
 }
 
-func (j jsonObject) allPaths() ([]JSON, error) {
-	if len(j) == 0 {
+func (j jsonObject) allPathsWithDepth(depth int) ([]JSON, error) {
+	if len(j) == 0 || depth == 0 {
 		return []JSON{j}, nil
 	}
 	ret := make([]JSON, 0, len(j))
 	for i := range j {
-		paths, err := j[i].v.allPaths()
+		paths, err := j[i].v.allPathsWithDepth(depth - 1)
 		if err != nil {
 			return nil, err
 		}