Skip to content

Unnest on nested json arrays in flat collections #244

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

Merged
merged 2 commits into from
Oct 27, 2025
+97 −21
Merged

Unnest on nested json arrays in flat collections #244

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
27b3d33
Support for unnested for json arrays in flat collections
suddendust Oct 22, 2025
ed0aa06
Added testFlatCollectionArrayAnyOnJsonbArray
suddendust Oct 22, 2025
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
65 changes: 65 additions & 0 deletions ...store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3981,6 +3981,71 @@ void testFlatVsNestedCollectionNestedFieldSelections(String dataStoreName) throw
assertDocsAndSizeEqual(
dataStoreName, flatBrandNoAliasIterator, "query/no_alias_response.json", 8);
}

/**
* Tests UNNEST operation on JSONB array fields in flat collections. This validates that
* jsonb_array_elements() is used for JSONB arrays (props.colors) instead of unnest() which is
* only for native arrays (tags).
*/
@ParameterizedTest
@ArgumentsSource(PostgresProvider.class)
void testFlatCollectionUnnestJsonbArray(String dataStoreName) throws IOException {
Datastore datastore = datastoreMap.get(dataStoreName);
Collection flatCollection =
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);

// Test UNNEST on JSONB array field: props.colors
// Expected: Should unnest colors and count distinct items with colors
// Data: id=1 has ["Blue", "Green"], id=3 has ["Black"], id=5 has ["Orange", "Blue"]
// Total: 5 color entries from 3 items
Query unnestJsonbQuery =
Query.builder()
.addSelection(IdentifierExpression.of("item"))
.addSelection(JsonIdentifierExpression.of("props", "colors"))
.addFromClause(
UnnestExpression.of(JsonIdentifierExpression.of("props", "colors"), false))
.build();

Iterator<Document> resultIterator = flatCollection.aggregate(unnestJsonbQuery);

long count = 0;
while (resultIterator.hasNext()) {
resultIterator.next();
count++;
}

// Expecting 5 results: 2 from Soap (Blue, Green), 1 from Shampoo (Black),
// 2 from Lifebuoy (Orange, Blue)
assertEquals(5, count, "Should find 5 color entries after unnesting JSONB arrays");
}

@ParameterizedTest
@ArgumentsSource(PostgresProvider.class)
void testFlatCollectionArrayAnyOnJsonbArray(String dataStoreName) {
Datastore datastore = datastoreMap.get(dataStoreName);
Collection flatCollection =
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);

// Test ArrayRelationalFilterExpression.ANY on JSONB array (props.colors)
// This uses jsonb_array_elements() internally
Query jsonbArrayQuery =
Query.builder()
.addSelection(IdentifierExpression.of("item"))
.setFilter(
ArrayRelationalFilterExpression.builder()
.operator(ArrayOperator.ANY)
.filter(
RelationalExpression.of(
JsonIdentifierExpression.of("props", "colors"),
EQ,
ConstantExpression.of("Blue")))
.build())
.build();

long count = flatCollection.count(jsonbArrayQuery);
// ids 1 and 5 have "Blue" in their colors array
assertEquals(2, count, "Should find 2 items with 'Blue' color (ids 1, 5)");
}
}

@Nested
Expand Down
31 changes: 18 additions & 13 deletions ...ace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import org.hypertrace.core.documentstore.expression.impl.ArrayRelationalFilterExpression;
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
import org.hypertrace.core.documentstore.expression.impl.DocumentArrayFilterExpression;
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
import org.hypertrace.core.documentstore.expression.impl.KeyExpression;
import org.hypertrace.core.documentstore.expression.impl.LogicalExpression;
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
Expand Down Expand Up @@ -169,22 +170,24 @@ private String getFilterStringForAnyOperator(final ArrayRelationalFilterExpressi
boolean isFlatCollection =
postgresQueryParser.getPgColTransformer().getDocumentType() == DocumentType.FLAT;

boolean isJsonbArray = expression.getArraySource() instanceof JsonIdentifierExpression;

// Extract the field name
final String identifierName =
expression
.getArraySource()
.accept(new PostgresIdentifierExpressionVisitor(postgresQueryParser));

final String parsedLhs;
if (isFlatCollection) {
// For flat collections, assume all arrays are native PostgreSQL arrays
if (isFlatCollection && !isJsonbArray) {
// For flat collections with native arrays, use direct column reference
parsedLhs = postgresQueryParser.transformField(identifierName).getPgColumn();
} else {
// For nested collections, use JSONB path accessor
// For nested collections OR JSONB arrays in flat collections, use JSONB path accessor
// Convert 'elements' to planets->'elements' where planets could be an alias for an upper
// level array filter
// For the first time (if 'elements' was not under any nested array, say a top-level field),
// use the field identifier visitor to make it document->'elements'
// use the field identifier visitor to make it document->'elements' or props->'colors'
final PostgresIdentifierExpressionVisitor identifierVisitor =
new PostgresIdentifierExpressionVisitor(postgresQueryParser);
final PostgresSelectTypeExpressionVisitor arrayPathVisitor =
Expand All @@ -206,18 +209,18 @@ private String getFilterStringForAnyOperator(final ArrayRelationalFilterExpressi
.getFilter()
.accept(new PostgresFilterTypeExpressionVisitor(postgresQueryParser, visitorProvider));

if (isFlatCollection) {
if (isFlatCollection && !isJsonbArray) {
// todo: For array filters, UNNEST is not the most optimal way as it won't use the index.
// Perhaps, we should use ANY or @> ARRAY operator

// For flat collections, assume all arrays are native and use unnest()
// For flat collections with native arrays (e.g., tags), use unnest()
// Infer array type from filter to properly cast empty array
String arrayTypeCast = inferArrayTypeCastFromFilter(expression.getFilter());
return String.format(
"EXISTS (SELECT 1 FROM unnest(COALESCE(%s, ARRAY[]%s)) AS \"%s\" WHERE %s)",
parsedLhs, arrayTypeCast, alias, parsedFilter);
} else {
// For nested collections with JSONB arrays, use jsonb_array_elements()
// For nested collections OR JSONB arrays in flat collections, use jsonb_array_elements()
return String.format(
"EXISTS (SELECT 1 FROM jsonb_array_elements(COALESCE(%s, '[]'::jsonb)) AS \"%s\" WHERE %s)",
parsedLhs, alias, parsedFilter);
Expand Down Expand Up @@ -284,18 +287,20 @@ private String getFilterStringForAnyOperator(final DocumentArrayFilterExpression
boolean isFlatCollection =
postgresQueryParser.getPgColTransformer().getDocumentType() == DocumentType.FLAT;

boolean isJsonbArray = expression.getArraySource() instanceof JsonIdentifierExpression;

// Extract the field name
final String identifierName =
expression
.getArraySource()
.accept(new PostgresIdentifierExpressionVisitor(postgresQueryParser));

final String parsedLhs;
if (isFlatCollection) {
// For flat collections, assume all arrays are native PostgreSQL arrays
// Use direct column reference with double quotes
if (isFlatCollection && !isJsonbArray) {
// For flat collections with native arrays, use direct column reference with double quotes
parsedLhs = postgresQueryParser.transformField(identifierName).getPgColumn();
} else {
// For nested collections OR JSONB arrays in flat collections, use JSONB path accessor
final PostgresIdentifierExpressionVisitor identifierVisitor =
new PostgresIdentifierExpressionVisitor(postgresQueryParser);
final PostgresSelectTypeExpressionVisitor arrayPathVisitor =
Expand All @@ -316,16 +321,16 @@ private String getFilterStringForAnyOperator(final DocumentArrayFilterExpression
.getFilter()
.accept(new PostgresFilterTypeExpressionVisitor(postgresQueryParser, wrapper));

if (isFlatCollection) {
// For flat collections, assume all arrays are native and use unnest()
if (isFlatCollection && !isJsonbArray) {
// For flat collections with native arrays, use unnest()
// Note: DocumentArrayFilterExpression typically works with JSONB arrays containing objects
// For simplicity, we default to text[] type cast, though this may need refinement
String arrayTypeCast = "::text[]";
return String.format(
"EXISTS (SELECT 1 FROM unnest(COALESCE(%s, ARRAY[]%s)) AS \"%s\" WHERE %s)",
parsedLhs, arrayTypeCast, alias, parsedFilter);
} else {
// For nested collections with JSONB arrays, use jsonb_array_elements()
// For nested collections OR JSONB arrays in flat collections, use jsonb_array_elements()
return String.format(
"EXISTS (SELECT 1 FROM jsonb_array_elements(COALESCE(%s, '[]'::jsonb)) AS \"%s\" WHERE %s)",
parsedLhs, alias, parsedFilter);
Expand Down
22 changes: 14 additions & 8 deletions ...trace/core/documentstore/postgres/query/v1/vistors/PostgresFromTypeExpressionVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import java.util.stream.Collectors;
import lombok.Getter;
import org.hypertrace.core.documentstore.DocumentType;
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
import org.hypertrace.core.documentstore.expression.impl.SubQueryJoinExpression;
import org.hypertrace.core.documentstore.expression.impl.UnnestExpression;
import org.hypertrace.core.documentstore.parser.FromTypeExpressionVisitor;
Expand Down Expand Up @@ -47,11 +48,14 @@ public String visit(UnnestExpression unnestExpression) {
boolean isFlatCollection =
postgresQueryParser.getPgColTransformer().getDocumentType() == DocumentType.FLAT;

boolean isJsonbArray =
unnestExpression.getIdentifierExpression() instanceof JsonIdentifierExpression;

String transformedFieldName;
String unnestFunction;

if (isFlatCollection) {
// For flat collections, assume all unnested fields are native PostgreSQL arrays
if (isFlatCollection && !isJsonbArray) {
// For flat collections with native arrays (e.g., tags), use unnest()
// Use the transformer to get the proper column name (handles quotes and naming)
transformedFieldName = postgresQueryParser.transformField(orgFieldName).getPgColumn();
// Use native unnest() for PostgreSQL array columns
Expand All @@ -60,7 +64,7 @@ public String visit(UnnestExpression unnestExpression) {
// e.g., unnest("tags") p1(tags_unnested) instead of p1(tags)
pgColumnName = pgColumnName + "_unnested";
} else {
// For nested collections, use JSONB path accessor
// For nested collections OR JSONB arrays in flat collections, use jsonb_array_elements()
transformedFieldName =
unnestExpression
.getIdentifierExpression()
Expand All @@ -78,8 +82,12 @@ public String visit(UnnestExpression unnestExpression) {
String tableAlias = "t" + preIndex;
String unwindExpr = String.format(unnestFunction, transformedFieldName);

// we'll quote the col name to prevent folding to lower case for top-level array fields
String unwindExprAlias =
String.format(UNWIND_EXP_ALIAS_FMT, nextIndex, getColName(isFlatCollection, pgColumnName));
String.format(
UNWIND_EXP_ALIAS_FMT,
nextIndex,
getColName(isFlatCollection && !isJsonbArray, pgColumnName));

String fmt =
unnestExpression.isPreserveNullAndEmptyArrays()
Expand Down Expand Up @@ -144,9 +152,7 @@ private static String prepareTable0Query(PostgresQueryParser postgresQueryParser
/*
Returns the column name with double quotes if the collection is flat to prevent folding to lower-case by PG
*/
private String getColName(boolean isFlatCollection, String pgColumnName) {
return isFlatCollection
? PostgresUtils.wrapFieldNamesWithDoubleQuotes(pgColumnName)
: pgColumnName;
private String getColName(boolean shouldQuote, String pgColumnName) {
return shouldQuote ? PostgresUtils.wrapFieldNamesWithDoubleQuotes(pgColumnName) : pgColumnName;
}
}
Loading