diff --git a/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/schema/OpenSearchSchemaBuilder.java b/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/schema/OpenSearchSchemaBuilder.java index 28aa6cd2f06ea..012ff77ffb84f 100644 --- a/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/schema/OpenSearchSchemaBuilder.java +++ b/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/schema/OpenSearchSchemaBuilder.java @@ -27,6 +27,7 @@ import org.opensearch.core.common.Strings; import org.opensearch.index.IndexNotFoundException; +import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; @@ -352,6 +353,15 @@ private static void addLeafFields( Map nested = (Map) fieldProps.get("properties"); if (nested != null) { addLeafFields(builder, typeFactory, nested, fieldName); + // Also expose the object itself as a struct (ROW) column, so a query can + // address the whole object (`fields nested_metadata`, `stats … by obj`) and + // not just its leaves. The object has no physical storage — the scan reads + // the leaves — so ObjectStructMaterializer strips this column from the scan + // and re-assembles it with make_struct in a project directly above it. + RelDataType structType = buildObjectType(typeFactory, nested, fieldName); + if (structType != null) { + builder.add(fieldName, structType); + } } continue; } @@ -372,6 +382,54 @@ private static void addLeafFields( } } + /** + * Builds the struct (ROW) type for an {@code object} mapping: one struct field per supported + * sub-field, recursing for sub-objects. Field names are the local names (the struct + * nesting already carries the path), while the flat leaf columns added by + * {@link #addLeafFields} keep their dotted paths — that dotted convention is how + * {@code ObjectStructMaterializer} pairs a struct field back to its backing column. + * + *

Returns {@code null} when the object contributes no supported field, so callers omit the + * column entirely rather than declaring an empty struct. + */ + @SuppressWarnings("unchecked") + private static RelDataType buildObjectType(RelDataTypeFactory typeFactory, Map properties, String pathPrefix) { + List types = new ArrayList<>(); + List names = new ArrayList<>(); + for (Map.Entry fieldEntry : properties.entrySet()) { + String localName = fieldEntry.getKey(); + Map fieldProps = (Map) fieldEntry.getValue(); + String fieldType = (String) fieldProps.get("type"); + RelDataType childType; + if (fieldType == null || "object".equals(fieldType)) { + Map nested = (Map) fieldProps.get("properties"); + childType = nested == null ? null : buildObjectType(typeFactory, nested, pathPrefix + "." + localName); + } else if ("nested".equals(fieldType)) { + // Array-of-sub-docs needs LIST + UNNEST; out of scope here. + childType = null; + } else { + // An ordinary scalar leaf (keyword, long, date_nanos, …). Typed exactly as + // addLeafFields types the equivalent flat column — including `scaling_factor`, + // the mapping parameter of `scaled_float` (a float persisted as a scaled long), + // which buildLeafType needs to type that field. parseScalingFactor yields NaN + // when absent, which buildLeafType treats as "not a scaled_float". + double scalingFactor = parseScalingFactor(fieldProps.get(SCALING_FACTOR_FIELD)); + childType = buildLeafType(fieldType, (String) fieldProps.get("format"), scalingFactor, typeFactory); + } + if (childType == null) { + // Unsupported sub-field (geo_point, nested, …) — dropped from the struct for the + // same reason it is dropped from the flat columns. + continue; + } + names.add(localName); + types.add(childType); + } + if (names.isEmpty()) { + return null; + } + return typeFactory.createTypeWithNullability(typeFactory.createStructType(types, names), true); + } + /** * Normalizes a scaling_factor value (Number or String from mapping JSON) to a double. * Returns {@link Double#NaN} when the value is null or unparseable, signaling the caller diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java index 1148b4ddb6a83..5cc88b5a0b70f 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java @@ -144,6 +144,11 @@ public static FieldType fromSqlTypeName(SqlTypeName sqlTypeName) { case BINARY, VARBINARY -> FieldType.BINARY; case ARRAY -> FieldType.ARRAY; case MAP -> FieldType.MAP; + // A Calcite ROW is the struct shape an `object` mapping materializes into (see + // ObjectStructMaterializer). Without this case, any struct-typed expression fails + // capability resolution in OpenSearchProjectRule and surfaces as + // UnsupportedFunctionException rather than dispatching to a backend. + case ROW -> OBJECT; default -> null; }; } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/MakeStructFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/MakeStructFunction.java new file mode 100644 index 0000000000000..cb8841dd062a7 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/MakeStructFunction.java @@ -0,0 +1,95 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.spi; + +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.SqlTypeName; + +import java.util.ArrayList; +import java.util.List; + +/** + * {@code make_struct('name0', value0, 'name1', value1, ...)} → ROW. + * + *

Materializes an OpenSearch {@code object} field as a struct from the flat dotted leaf + * columns the parquet scan produces. The engine stores {@code object} sub-fields as flat + * columns ({@code a.b.c}); this function re-assembles them into the nested shape at query + * time so a projection or aggregate can address the object as a single value. Nested objects + * nest the call: + * + * 

+ * make_struct('top', $1, 'properties', make_struct('name', $2, 'value', $3))
+ *
+ * + *

Executes as DataFusion's {@code named_struct}, which takes the same + * (name, value, name, value, …) calling convention. The call does not reach DataFusion + * through Substrait's function-signature matching: {@code MakeStructCallConverter} intercepts it + * and builds the invocation directly, because isthmus cannot match a variadic function whose + * operands are deliberately of differing types. That is what makes the number of struct fields + * unbounded. + * + *

The return type is always supplied explicitly by the caller via + * {@link #makeCall(RexBuilder, RelDataType, List, List)} — the operand-driven inference is a + * placeholder ({@code ANY}) because the authoritative ROW type comes from the index mapping, + * not from the operands. + * + * @opensearch.internal + */ +public final class MakeStructFunction { + + /** The function name used in Calcite plans and Substrait serialization. */ + public static final String NAME = "make_struct"; + + /** Singleton Calcite SqlFunction: {@code make_struct(VARCHAR, ANY, ...) → ROW}. */ + public static final SqlFunction FUNCTION = new SqlFunction( + NAME, + SqlKind.OTHER_FUNCTION, + // Placeholder: callers always construct with an explicit ROW type (see makeCall). + opBinding -> opBinding.getTypeFactory().createSqlType(SqlTypeName.ANY), + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.USER_DEFINED_FUNCTION + ); + + private MakeStructFunction() {} + + /** + * Builds {@code make_struct('f0', v0, 'f1', v1, ...)} with an explicit ROW return type. + * + * @param rexBuilder builder for the enclosing plan + * @param structType the ROW type this call produces (from the index mapping) + * @param fieldNames struct field names, in order + * @param fieldValues struct field value expressions, positionally paired with {@code fieldNames} + */ + public static RexNode makeCall(RexBuilder rexBuilder, RelDataType structType, List fieldNames, List fieldValues) { + if (fieldNames.size() != fieldValues.size()) { + throw new IllegalArgumentException( + "make_struct requires one value per field name; got " + fieldNames.size() + " names and " + fieldValues.size() + " values" + ); + } + // Field names are VARCHAR, not CHAR. rexBuilder.makeLiteral(String) would produce a + // CHAR(n) literal, which Substrait types as the fixed-width `char`; DataFusion's + // named_struct expects a variable-length string for a field name, so VARCHAR (→ Substrait + // `string`) is the faithful type. It also keeps the literal free of the padding semantics + // CHAR carries. + RelDataType nameType = rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR); + List operands = new ArrayList<>(fieldNames.size() * 2); + for (int i = 0; i < fieldNames.size(); i++) { + operands.add(rexBuilder.makeLiteral(fieldNames.get(i), nameType, true)); + operands.add(fieldValues.get(i)); + } + return rexBuilder.makeCall(structType, FUNCTION, operands); + } +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java index eb74cc1719497..30156d45758db 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java @@ -455,7 +455,16 @@ public enum ScalarFunction { MINSPAN_BUCKET(Category.SCALAR, SqlKind.OTHER_FUNCTION), /** PPL range_bucket(value, data_min, data_max, start_param, end_param). VARCHAR label. */ - RANGE_BUCKET(Category.SCALAR, SqlKind.OTHER_FUNCTION); + RANGE_BUCKET(Category.SCALAR, SqlKind.OTHER_FUNCTION), + + // ── Composite (object / struct) construction ───────────────────── + /** + * {@code make_struct('name0', v0, 'name1', v1, ...)} → ROW. Materializes an OpenSearch + * {@code object} field from the flat dotted leaf columns the parquet scan produces + * (see {@code ObjectStructMaterializer}). Maps to DataFusion's {@code named_struct}. + * Nests for sub-objects. See {@link MakeStructFunction}. + */ + MAKE_STRUCT(Category.SCALAR, SqlKind.OTHER_FUNCTION); /** * Category of scalar function. diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index 1a96960f0e09a..e3747c532b040 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -473,6 +473,18 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP */ private static final Set POLYMORPHIC_RETURN_PROJECT_OPS = Set.of(ScalarFunction.CAST, ScalarFunction.SAFE_CAST); + /** + * Project-side scalar functions whose return type is a struct (Calcite {@code ROW} → + * {@link FieldType#OBJECT}). Registered separately for the same reason as + * {@link #MAP_RETURNING_PROJECT_OPS}: the capability lookup in + * {@code OpenSearchProjectRule.resolveScalarViableBackends} keys on the call's return type. + * + *

{@code MAKE_STRUCT} materializes an OpenSearch {@code object} field from the flat dotted + * leaf columns the parquet scan produces (see {@code ObjectStructMaterializer}); it executes as + * DataFusion's {@code named_struct}, emitted by {@code MakeStructCallConverter}. + */ + private static final Set OBJECT_RETURNING_PROJECT_OPS = Set.of(ScalarFunction.MAKE_STRUCT); + private static final Set AGG_FUNCTIONS = Set.of( AggregateFunction.SUM, AggregateFunction.SUM0, @@ -635,6 +647,9 @@ public Set projectCapabilities() { for (ScalarFunction op : MAP_RETURNING_PROJECT_OPS) { caps.add(new ProjectCapability.Scalar(op, Set.of(FieldType.MAP), formats, true)); } + for (ScalarFunction op : OBJECT_RETURNING_PROJECT_OPS) { + caps.add(new ProjectCapability.Scalar(op, Set.of(FieldType.OBJECT), formats, true)); + } for (ScalarFunction op : POLYMORPHIC_RETURN_PROJECT_OPS) { for (FieldType ft : FieldType.values()) { caps.add(new ProjectCapability.Scalar(op, Set.of(ft), formats, true)); @@ -647,6 +662,18 @@ public Set projectCapabilities() { public Set aggregateCapabilities() { Set formats = Set.copyOf(plugin.getSupportedFormats()); Set caps = new HashSet<>(); + // Aggregates over a struct-typed (object) column. Registered separately because + // FieldType.OBJECT is deliberately NOT in SUPPORTED_FIELD_TYPES — most aggregates + // are meaningless on a struct (no ordering ⇒ MIN/MAX, no arithmetic ⇒ SUM/AVG), + // and adding OBJECT there would also wrongly claim filter/sort support. + // + // COUNT is the meaningful one (non-null count) and is what the SQL plugin's + // Lucene path supports today, i.e. `stats count()`. Without this the + // capability lookup for (COUNT, OBJECT) misses and the query fails at plan time + // with "Function [COUNT] is not currently supported as an aggregate function". + for (FieldType objectType : Set.of(FieldType.OBJECT)) { + caps.add(new AggregateCapability(AggregateFunction.COUNT, Set.of(objectType), formats)); + } for (AggregateFunction func : AGG_FUNCTIONS) { for (FieldType type : SUPPORTED_FIELD_TYPES) { // 3-arg constructor leaves decomposition=null so the diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 56a911e309abe..84a1add3195ce 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java @@ -57,6 +57,7 @@ import java.math.MathContext; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.function.Function; @@ -577,7 +578,16 @@ private byte[] convertToSubstrait(RelNode fragment) { throw new IllegalStateException("Substrait conversion rejected the plan: " + e.getMessage(), e); } - List fieldNames = root.fields.stream().map(field -> field.getValue()).toList(); + // Output column names Substrait attaches to the plan root (Plan.Root.names); DataFusion + // names the result schema from them. Flattened so struct columns also contribute their + // nested field names — see flattenNamesForSubstrait. + // + // Each root field carries its own index into the row type, and that index is NOT always + // the field's position in this list (RelRoot may project or permute), so the type must be + // looked up by field.getKey() rather than by list position — otherwise a name gets paired + // with the wrong type and a mispaired struct injects nested names that do not belong, + // corrupting the schema. + List fieldNames = flattenNamesForSubstrait(root.fields, preprocessed.getRowType()); Plan.Root substraitRoot = Plan.Root.builder().input(substraitRel).names(fieldNames).build(); Plan plan = Plan.builder().addRoots(substraitRoot).build(); @@ -610,7 +620,73 @@ static Plan rewire(Plan inner, Rel wrapper, List wrapperNames) { /** Wrapper's output column names from its Calcite row type. */ private static List fieldNames(RelNode fragment) { - return fragment.getRowType().getFieldList().stream().map(RelDataTypeField::getName).toList(); + return flattenNamesForSubstrait(fragment.getRowType()); + } + + /** + * Flattens column names for a Substrait {@code Plan.Root} / {@code NamedStruct}. + * + *

Substrait carries schema names as ONE depth-first list covering every field at every + * nesting level — a struct contributes its own name followed by its children's names, + * recursively. A Substrait struct expression is positional (it holds only field + * values, no names), so the names live exclusively in this list; that is why it has to be + * complete. + * + *

Example — for the output row type + * 

+     * id   INTEGER
+     * meta ROW(top VARCHAR, props ROW(name VARCHAR, value VARCHAR))
+     *
+ * Substrait expects six names, depth-first: + * 
+     * ["id", "meta", "top", "props", "name", "value"]
+     *          └────┬────┘   └──────┬──────┘
+     *         meta's children   props' children
+     *
+ * Emitting only the top-level {@code ["id", "meta"]} makes DataFusion's substrait consumer + * reject the plan with {@code "Named schema must contain names for all fields"} — which is + * exactly what happened for an {@code object} materialized by + * {@code ObjectStructMaterializer} before this flattening was added. + * + *

Names are taken from {@code rootFields} rather than from {@code rowType} because a + * {@link RelRoot} may alias its output columns; only the nested names come from the type. Each + * root field also carries its own index into {@code rowType}, which is not necessarily its + * position in the list ({@code RelRoot} may project or permute), so the type is looked up by + * {@link java.util.Map.Entry#getKey()}. + */ + private static List flattenNamesForSubstrait(List> rootFields, RelDataType rowType) { + List fields = rowType.getFieldList(); + List flattened = new ArrayList<>(rootFields.size()); + for (Map.Entry rootField : rootFields) { + flattened.add(rootField.getValue()); + int index = rootField.getKey(); + if (index >= 0 && index < fields.size() && fields.get(index).getType().isStruct()) { + appendNestedNames(flattened, fields.get(index).getType()); + } + } + return flattened; + } + + /** Row-type-driven overload: names come from the row type itself, so positions align by construction. */ + private static List flattenNamesForSubstrait(RelDataType rowType) { + List flattened = new ArrayList<>(rowType.getFieldCount()); + for (RelDataTypeField field : rowType.getFieldList()) { + flattened.add(field.getName()); + if (field.getType().isStruct()) { + appendNestedNames(flattened, field.getType()); + } + } + return flattened; + } + + /** Appends a struct's field names depth-first (children before the next sibling). */ + private static void appendNestedNames(List out, RelDataType structType) { + for (RelDataTypeField child : structType.getFieldList()) { + out.add(child.getName()); + if (child.getType().isStruct()) { + appendNestedNames(out, child.getType()); + } + } } private static Rel replaceInput(Rel wrapper, Rel newInput) { @@ -798,7 +874,20 @@ protected ImmutableList getSigs() { aggConverter, windowConverter, typeConverter - ); + ) { + @Override + public List getCallConverters() { + // Offer struct construction to MakeStructCallConverter before function matching. + // RexExpressionConverter.visitCall walks this list and only throws "Unable to + // convert call ..." once every converter declines, so building the invocation + // there sidesteps isthmus' SingularArgumentMatcher — which is what removes the + // arity ceiling a declared named_struct signature would impose. + List converters = new ArrayList<>(); + converters.add(new MakeStructCallConverter(extensions, typeConverter)); + converters.addAll(super.getCallConverters()); + return converters; + } + }; return new SubstraitRelVisitor(converterProvider) { @Override public Rel visit(org.apache.calcite.rel.core.Aggregate aggregate) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index e05d6127c9c0d..662729360b705 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -738,6 +738,7 @@ private static SimpleExtension.ExtensionCollection loadSubstraitExtensions() { SimpleExtension.ExtensionCollection delegationExtensions = SimpleExtension.load(List.of("/delegation_functions.yaml")); SimpleExtension.ExtensionCollection scalarExtensions = SimpleExtension.load(List.of("/opensearch_scalar_functions.yaml")); SimpleExtension.ExtensionCollection arrayExtensions = SimpleExtension.load(List.of("/opensearch_array_functions.yaml")); + SimpleExtension.ExtensionCollection structExtensions = SimpleExtension.load(List.of("/opensearch_struct_functions.yaml")); SimpleExtension.ExtensionCollection aggregateExtensions = SimpleExtension.load(List.of("/opensearch_aggregate_functions.yaml")); SimpleExtension.ExtensionCollection windowExtensions = SimpleExtension.load(List.of("/opensearch_window_functions.yaml")); SimpleExtension.ExtensionCollection arithmeticOverloads = SimpleExtension.load( @@ -746,6 +747,7 @@ private static SimpleExtension.ExtensionCollection loadSubstraitExtensions() { return DefaultExtensionCatalog.DEFAULT_COLLECTION.merge(delegationExtensions) .merge(scalarExtensions) .merge(arrayExtensions) + .merge(structExtensions) .merge(aggregateExtensions) .merge(windowExtensions) .merge(arithmeticOverloads); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MakeStructCallConverter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MakeStructCallConverter.java new file mode 100644 index 0000000000000..e5ef5a183cb85 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MakeStructCallConverter.java @@ -0,0 +1,117 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion; + +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.opensearch.analytics.spi.MakeStructFunction; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.function.Function; + +import io.substrait.expression.Expression; +import io.substrait.expression.FunctionArg; +import io.substrait.extension.SimpleExtension; +import io.substrait.isthmus.CallConverter; +import io.substrait.isthmus.TypeConverter; + +/** + * Converts a {@code make_struct} / {@code named_struct} call into a Substrait + * {@link Expression.ScalarFunctionInvocation} built directly, bypassing isthmus' + * function-signature matching. This is what lifts the arity ceiling on struct construction. + * + *

Why the signature matcher can't do this

+ * + *
    + *
  • A variadic declaration never matches. isthmus binds variadic functions through a + * {@code SingularArgumentMatcher} — it derives ONE type every operand must satisfy. + * {@code named_struct} interleaves {@code string} field names with values of unrelated + * types, so no such type exists. Verified against isthmus 0.89.1: neither + * {@code value: any1} nor unconstrained {@code value: any} (both with + * {@code parameterConsistency: INCONSISTENT}) matches — each still fails with + * "Unable to convert call named_struct(...)". Same reason {@link MakeArrayAdapter} must + * widen {@code make_array}'s operands to a common element type first.
    • + *
  • Fixed-arity enumeration is bounded. One impl per field count does match, but an + * object's width is data-dependent: an OTel span's {@code attributes} carries ~55 + * sub-fields and grows as new attribute keys appear.
    • + *
+ * + *

{@code RexExpressionConverter.visitCall} offers a call to every registered + * {@link CallConverter} and throws "Unable to convert call ..." only once all of them decline, so + * constructing the invocation here skips the matcher entirely. The declaration is used purely as + * the extension anchor (name + URN) that the consumer resolves by name; the argument list + * we attach is the real one, so its length is unconstrained by the declared variant. + * + *

Why not a nested-struct expression

+ * + * Substrait models struct construction natively as {@code Expression.NestedStruct}, which would be + * the cleaner representation, but DataFusion's substrait consumer rejects it at execution: + * {@code "This feature is not implemented: Nested struct expressions are not yet supported"}. + * DataFusion does implement {@code named_struct} as a scalar function at any arity, so the + * function-invocation form is the one that actually executes. Revisit if that gap is closed + * upstream. + * + *

Operands are forwarded unchanged — including the field-name literals — because DataFusion's + * {@code named_struct} takes the interleaved {@code (name, value, …)} form. + * + * @opensearch.internal + */ +class MakeStructCallConverter implements CallConverter { + + /** DataFusion's native struct constructor — the name the consumer resolves. */ + static final String NAMED_STRUCT = "named_struct"; + + private final SimpleExtension.ExtensionCollection extensions; + private final TypeConverter typeConverter; + + MakeStructCallConverter(SimpleExtension.ExtensionCollection extensions, TypeConverter typeConverter) { + this.extensions = extensions; + this.typeConverter = typeConverter; + } + + @Override + public Optional convert(RexCall call, Function topLevelConverter) { + String operator = call.getOperator().getName(); + // The engine emits `make_struct`; `named_struct` is accepted too so the converter stays + // correct if a rename ever reaches it first. + if (!MakeStructFunction.NAME.equalsIgnoreCase(operator) && !NAMED_STRUCT.equalsIgnoreCase(operator)) { + return Optional.empty(); + } + + Optional declaration = findNamedStructDeclaration(); + if (declaration.isEmpty()) { + // No anchor to reference — decline so the failure surfaces as isthmus' normal + // "Unable to convert call" rather than an NPE deep in proto serialization. + return Optional.empty(); + } + + List arguments = new ArrayList<>(call.getOperands().size()); + for (RexNode operand : call.getOperands()) { + arguments.add(topLevelConverter.apply(operand)); + } + + return Optional.of( + Expression.ScalarFunctionInvocation.builder() + .declaration(declaration.get()) + .addAllArguments(arguments) + .outputType(typeConverter.toSubstrait(call.getType())) + .build() + ); + } + + /** + * Any declared {@code named_struct} variant works as the anchor — the consumer resolves the + * function by name, and the arity we attach is independent of the variant's declared arity. + */ + private Optional findNamedStructDeclaration() { + return extensions.scalarFunctions().stream().filter(variant -> NAMED_STRUCT.equals(variant.name())).findFirst(); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_struct_functions.yaml b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_struct_functions.yaml new file mode 100644 index 0000000000000..1d2e397fac9f8 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_struct_functions.yaml @@ -0,0 +1,38 @@ +%YAML 1.2 +--- +# Substrait extension declaring DataFusion's native struct constructor. +# +# This file exists ONLY to give the function an extension anchor (URN + name) that the consumer +# resolves by name. The argument list is attached by MakeStructCallConverter, which builds the +# Expression.ScalarFunctionInvocation directly and therefore never goes through isthmus' +# signature matcher — so the single impl below does NOT bound the number of struct fields. +# +# WHY THE MATCHER IS BYPASSED (and why this is not an enumeration of arities): +# isthmus binds a variadic function via a SingularArgumentMatcher — it derives ONE type that every +# variadic operand must satisfy (the same reason MakeArrayAdapter widens make_array's operands to a +# common element type first). named_struct interleaves `string` field names with values of +# unrelated types, so no such type exists. Verified against isthmus 0.89.1: a variadic impl matches +# neither as `value: any1` nor as unconstrained `value: any`, with +# `parameterConsistency: INCONSISTENT` in both cases. Enumerating one impl per field count DOES +# match but is bounded, and an object's width is data-dependent (an OTel span's `attributes` +# carries ~55 sub-fields and grows), so enumeration is not a shippable answer. +# +# Substrait models struct construction natively as Expression.NestedStruct, which would need no +# declaration at all, but DataFusion's substrait consumer rejects it at execution: +# "This feature is not implemented: Nested struct expressions are not yet supported". +# DataFusion does implement `named_struct` as a scalar function at any arity, so the +# function-invocation form is the one that executes. +urn: extension:org.opensearch:struct_functions +scalar_functions: + - name: named_struct + description: >- + Construct a struct from interleaved (field-name, value) pairs: + named_struct('f0', v0, 'f1', v1, ...). Anchor declaration only — see the header. The real + arguments are attached by MakeStructCallConverter and are not limited to this arity. + impls: + - args: + - value: string + name: "name" + - value: any1 + name: "value" + return: any diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/MakeStructCallConverterTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/MakeStructCallConverterTests.java new file mode 100644 index 0000000000000..2d6010bbbc671 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/MakeStructCallConverterTests.java @@ -0,0 +1,163 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.be.datafusion; + +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.analytics.spi.MakeStructFunction; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.function.Function; + +import io.substrait.expression.Expression; +import io.substrait.extension.SimpleExtension; +import io.substrait.isthmus.TypeConverter; + +/** + * Substrait-serialization tests for {@link MakeStructCallConverter}. + * + *

This is the layer both hard struct bugs lived in — an arity ceiling on + * {@code named_struct} and a schema whose nested names were missing — and neither was reachable + * from a plan-shape unit test (those stop before Substrait) nor cheap to catch in a REST IT (needs + * a live cluster plus the Rust native library). The width test below is specifically the regression + * guard for the ceiling: it fails immediately if struct construction is ever routed back through + * isthmus' function-signature matching, which cannot bind a variadic call whose operands differ in + * type and therefore has to enumerate arities. + */ +public class MakeStructCallConverterTests extends OpenSearchTestCase { + + private RelDataTypeFactory typeFactory; + private RexBuilder rexBuilder; + private MakeStructCallConverter converter; + + /** Converts operands the way isthmus would; enough for asserting arity and shape. */ + private Function operandConverter; + + @Override + public void setUp() throws Exception { + super.setUp(); + typeFactory = new JavaTypeFactoryImpl(); + rexBuilder = new RexBuilder(typeFactory); + SimpleExtension.ExtensionCollection extensions = SimpleExtension.load(List.of("/opensearch_struct_functions.yaml")); + converter = new MakeStructCallConverter(extensions, TypeConverter.DEFAULT); + // Stand-in for isthmus' recursive conversion: every operand becomes a distinct literal, so + // the assertions can count arguments without depending on the real converter graph. + operandConverter = rex -> Expression.StrLiteral.builder().value(rex.toString()).build(); + } + + /** + * A 100-field object serializes to a single invocation carrying all 200 operands. No arity is + * declared anywhere, which is the whole point: an OTel span's {@code attributes} already has + * ~55 sub-fields and grows with the data, so any fixed ceiling is a latent production failure. + */ + public void testWideStructHasNoArityCeiling() { + int fieldCount = 100; + Optional converted = converter.convert(makeStructCall(fieldCount), operandConverter); + + assertTrue("converter must handle a 100-field struct", converted.isPresent()); + Expression.ScalarFunctionInvocation invocation = asInvocation(converted.get()); + assertEquals("one argument per name and per value", fieldCount * 2, invocation.arguments().size()); + assertEquals("named_struct", invocation.declaration().name()); + } + + /** The common narrow case still works, and keeps the interleaved (name, value) ordering. */ + public void testNarrowStructKeepsInterleavedOperands() { + Optional converted = converter.convert(makeStructCall(2), operandConverter); + + assertTrue(converted.isPresent()); + assertEquals(4, asInvocation(converted.get()).arguments().size()); + } + + /** Unrelated operators must fall through so other converters (and function matching) still run. */ + public void testDeclinesUnrelatedOperator() { + RexNode left = rexBuilder.makeLiteral("a"); + RexNode right = rexBuilder.makeLiteral("b"); + RexCall unrelated = (RexCall) rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, left, right); + + assertFalse("EQUALS must not be claimed by the struct converter", converter.convert(unrelated, operandConverter).isPresent()); + } + + /** + * The mis-paired case is rejected at construction, which is why the converter's own odd-operand + * guard is unreachable from our emitter: {@code MakeStructFunction.makeCall} refuses a call + * whose name and value lists differ in length, so a struct with a dangling field can never be + * built in the first place. (Calcite's own operand assertions make a hand-rolled odd-arity + * {@code RexCall} un-constructible under {@code -ea}, so this asserts the real contract + * instead.) + */ + public void testEmitterRejectsMismatchedNameAndValueCounts() { + List twoNames = List.of("f0", "f1"); + List oneValue = List.of(rexBuilder.makeLiteral("v0")); + + IllegalArgumentException thrown = expectThrows( + IllegalArgumentException.class, + () -> MakeStructFunction.makeCall(rexBuilder, structTypeOf(2), twoNames, oneValue) + ); + assertTrue( + "message should name the mismatch, got: " + thrown.getMessage(), + thrown.getMessage().contains("one value per field name") + ); + } + + /** + * Without a {@code named_struct} declaration to anchor to there is no extension reference to + * emit, so the converter declines and the failure surfaces as isthmus' ordinary + * "Unable to convert call" rather than an NPE during proto serialization. Guards the coupling + * to {@code opensearch_struct_functions.yaml} being on the classpath. + */ + public void testDeclinesWhenAnchorDeclarationMissing() throws Exception { + MakeStructCallConverter withoutAnchor = new MakeStructCallConverter( + SimpleExtension.ExtensionCollection.builder().build(), + TypeConverter.DEFAULT + ); + + assertFalse(withoutAnchor.convert(makeStructCall(2), operandConverter).isPresent()); + } + + // ── helpers ────────────────────────────────────────────────────────────────────── + + /** {@code make_struct('f0', 'v0', 'f1', 'v1', …)} with {@code fieldCount} pairs. */ + private RexCall makeStructCall(int fieldCount) { + List names = new ArrayList<>(fieldCount); + List values = new ArrayList<>(fieldCount); + for (int i = 0; i < fieldCount; i++) { + names.add("f" + i); + values.add(rexBuilder.makeLiteral("v" + i)); + } + return (RexCall) MakeStructFunction.makeCall(rexBuilder, structTypeOf(fieldCount), names, values); + } + + /** ROW type with {@code fieldCount} VARCHAR fields. */ + private RelDataType structTypeOf(int fieldCount) { + List types = new ArrayList<>(fieldCount); + List names = new ArrayList<>(fieldCount); + for (int i = 0; i < fieldCount; i++) { + types.add(typeFactory.createSqlType(SqlTypeName.VARCHAR)); + names.add("f" + i); + } + return typeFactory.createStructType(types, names); + } + + private static Expression.ScalarFunctionInvocation asInvocation(Expression expression) { + assertTrue( + "expected a ScalarFunctionInvocation, got " + expression.getClass().getSimpleName(), + expression instanceof Expression.ScalarFunctionInvocation + ); + return (Expression.ScalarFunctionInvocation) expression; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ObjectStructMaterializer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ObjectStructMaterializer.java new file mode 100644 index 0000000000000..815a9ca4c66ce --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ObjectStructMaterializer.java @@ -0,0 +1,224 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.planner; + +import org.apache.calcite.plan.RelOptAbstractTable; +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelShuttleImpl; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.logical.LogicalTableScan; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.opensearch.analytics.spi.MakeStructFunction; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * Materializes OpenSearch {@code object} fields into structs in a project directly above the + * table scan. + * + *

The engine stores an {@code object} mapping as flat dotted leaf columns + * ({@code nested_metadata.top}, {@code nested_metadata.properties.name}, …); the schema also + * exposes the object itself as a struct-typed (ROW) column so a query can name it — see + * {@code OpenSearchSchemaBuilder.addLeafFields}. The object has no physical column + * though: {@code FieldStorageResolver.populateFromProperties} recurses past object parents + * precisely because "object fields themselves have no storage". So the scan must not be asked + * to produce it. + * + *

This rewriter therefore does two things at once — it removes struct columns from the scan + * and re-assembles them in a project above it, at their original positions: + * + * 

+ * LogicalProject(id=[$0], nested_metadata.top=[$1], nested_metadata.properties.name=[$2],
+ *                nested_metadata.properties.value=[$3],
+ *                nested_metadata=[make_struct('top', $1,
+ *                                   'properties', make_struct('name', $2, 'value', $3))])
+ *   LogicalTableScan(table=[[t]])      // row type: leaves only, no nested_metadata
+ *
+ * + *

The project reproduces the scan's original row type exactly — same field names, + * order, and types — so every {@code RexInputRef} in the operators above stays valid and no + * upstream rewriting is needed. Consequently an aggregate over an object + * ({@code stats … by nested_metadata}) receives an already-materialized struct, and a projection + * of the object returns the whole object. + * + *

A sub-object nests another {@code make_struct} over its own leaves, so arbitrarily deep + * {@code object} trees materialize in one pass. + * + *

Runs as a one-shot top-down pass (like {@code OpenSearchTopKRewriter}) rather than a HEP + * rule: a rule matching {@code TableScan} and producing {@code Project(TableScan)} would re-match + * its own output and never reach fixpoint. + * + *

Placement matters, and it is before {@code trimFields} — see the call site in + * {@code PlannerImpl.runAllOptimizations}. This pass emits one {@code make_struct} per object + * column the scan declares, and the field trimmer then removes the ones a given query never + * references; without that, a query filtering only on a leaf would pay for (and could fail on) an + * unrelated object. Leaf predicates still push down afterwards, because the leaves remain in this + * project's output and {@code FILTER_PROJECT_TRANSPOSE} moves filters through it. + * + *

If predicates authored against the struct itself ever need to push down, the companion + * rewrite is {@code GET_FIELD(struct, 'x') → leaf ref}. + * + * @opensearch.internal + */ +public final class ObjectStructMaterializer { + + private ObjectStructMaterializer() {} + + /** + * Rewrites scans that expose struct-typed columns into a leaf-only scan plus a + * struct-materializing project. + * + * @return the rewritten plan, or {@link Optional#empty()} when the plan has no object + * columns (callers keep the original plan unchanged) + */ + public static Optional rewrite(RelNode root) { + Materializer materializer = new Materializer(); + RelNode rewritten = root.accept(materializer); + return materializer.changed ? Optional.of(rewritten) : Optional.empty(); + } + + private static final class Materializer extends RelShuttleImpl { + + private boolean changed = false; + + @Override + public RelNode visit(TableScan scan) { + RelDataType originalRowType = scan.getRowType(); + List originalFields = originalRowType.getFieldList(); + if (originalFields.stream().noneMatch(f -> f.getType().isStruct())) { + return scan; + } + + // The scan keeps only physically-stored (non-struct) columns. + RelDataTypeFactory typeFactory = scan.getCluster().getTypeFactory(); + RelDataTypeFactory.Builder leafTypeBuilder = typeFactory.builder(); + Map leafIndexByName = new HashMap<>(); + for (RelDataTypeField field : originalFields) { + if (field.getType().isStruct()) { + continue; + } + leafIndexByName.put(field.getName(), leafTypeBuilder.getFieldCount()); + leafTypeBuilder.add(field.getName(), field.getType()); + } + RelDataType leafRowType = leafTypeBuilder.build(); + if (leafRowType.getFieldCount() == 0) { + // Nothing physical to read (an object-only projection). Leave the plan alone + // rather than emit a scan with an empty row type. + return scan; + } + + RelOptTable leafTable = new LeafOnlyTable(scan.getTable(), leafRowType); + RelNode leafScan = LogicalTableScan.create(scan.getCluster(), leafTable, scan.getHints()); + + // Rebuild the ORIGINAL row type above the trimmed scan: pass leaves through and + // assemble each struct in place, so parent input refs keep their meaning. + RexBuilder rexBuilder = scan.getCluster().getRexBuilder(); + List projects = new ArrayList<>(originalFields.size()); + List names = new ArrayList<>(originalFields.size()); + for (RelDataTypeField field : originalFields) { + names.add(field.getName()); + if (field.getType().isStruct()) { + RexNode struct = buildStruct(rexBuilder, leafScan, field.getName(), field.getType(), leafIndexByName); + if (struct == null) { + // A backing leaf is absent (e.g. an unsupported sub-field type was dropped + // from the schema). Emit a typed NULL rather than a partial struct, keeping + // the row type stable for the operators above. + projects.add(rexBuilder.makeNullLiteral(field.getType())); + } else { + projects.add(struct); + } + } else { + projects.add(rexBuilder.makeInputRef(leafScan, leafIndexByName.get(field.getName()))); + } + } + + changed = true; + return LogicalProject.create(leafScan, List.of(), projects, names); + } + + /** + * Recursively builds {@code make_struct} for {@code structType}, resolving each leaf to the + * trimmed scan's column named {@code path + "." + fieldName}. Returns {@code null} when any + * leaf is missing, signaling the caller to skip materialization for this column. + */ + private static RexNode buildStruct( + RexBuilder rexBuilder, + RelNode leafScan, + String path, + RelDataType structType, + Map leafIndexByName + ) { + List fieldNames = new ArrayList<>(); + List fieldValues = new ArrayList<>(); + for (RelDataTypeField child : structType.getFieldList()) { + String childPath = path + "." + child.getName(); + RexNode value; + if (child.getType().isStruct()) { + // A child that is itself an object nests another make_struct over its own leaves. + value = buildStruct(rexBuilder, leafScan, childPath, child.getType(), leafIndexByName); + } else { + Integer leafIndex = leafIndexByName.get(childPath); + value = leafIndex == null ? null : rexBuilder.makeInputRef(leafScan, leafIndex); + } + if (value == null) { + return null; + } + fieldNames.add(child.getName()); + fieldValues.add(value); + } + if (fieldNames.isEmpty()) { + return null; + } + return MakeStructFunction.makeCall(rexBuilder, structType, fieldNames, fieldValues); + } + } + + /** + * Wraps the scanned table with a row type stripped of struct columns, so downstream physical + * resolution ({@code FieldStorageResolver}) only ever sees fields that actually have storage. + * Mirrors the {@code IndexNameTable} wrapper in {@code OpenSearchTableScanRule}. + */ + private static final class LeafOnlyTable extends RelOptAbstractTable { + + private final RelOptTable delegate; + + LeafOnlyTable(RelOptTable delegate, RelDataType leafRowType) { + super(delegate.getRelOptSchema(), delegate.getQualifiedName().getLast(), leafRowType); + this.delegate = delegate; + } + + @Override + public List getQualifiedName() { + // Preserve the original qualified name: OpenSearchTableScanRule resolves the index + // from it, and RelOptAbstractTable would otherwise report a single-segment name. + return delegate.getQualifiedName(); + } + + @Override + public double getRowCount() { + return delegate.getRowCount(); + } + + @Override + public T unwrap(Class clazz) { + T unwrapped = super.unwrap(clazz); + return unwrapped != null ? unwrapped : delegate.unwrap(clazz); + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java index fa6b73dc7ee2c..c5bbc205b5bb9 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java @@ -41,6 +41,7 @@ import org.opensearch.analytics.planner.rules.OpenSearchAggregateReduceRule; import org.opensearch.analytics.planner.rules.OpenSearchAggregateRule; import org.opensearch.analytics.planner.rules.OpenSearchAggregateSplitRule; +import org.opensearch.analytics.planner.rules.OpenSearchAggregateStructKeyRule; import org.opensearch.analytics.planner.rules.OpenSearchBroadcastJoinSplitRule; import org.opensearch.analytics.planner.rules.OpenSearchCheckedLongSumRule; import org.opensearch.analytics.planner.rules.OpenSearchCheckedLongSumWindowRule; @@ -128,6 +129,17 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con RelNode modifiedRelNode = rawRelNode; modifiedRelNode = removeSubQueries(modifiedRelNode, listener); + // Materialize `object` fields into structs in a project directly above each scan. + // Runs BEFORE trimFields on purpose: the pass emits one make_struct per object column + // the scan declares, and the field trimmer then deletes the ones this query never + // references (a query filtering only on a leaf must not pay for — or fail on — an + // unrelated object). Leaf predicates still push down afterwards, since the leaves stay + // in the project's output and FILTER_PROJECT_TRANSPOSE moves filters through it. + Optional objectStructs = ObjectStructMaterializer.rewrite(modifiedRelNode); + if (objectStructs.isPresent()) { + modifiedRelNode = objectStructs.get(); + RelNodeUtils.logPlan(LOGGER, "After object-struct materialization", modifiedRelNode); + } modifiedRelNode = trimFields(modifiedRelNode); modifiedRelNode = extractLiteralAgg(modifiedRelNode, listener); modifiedRelNode = reduceExpressions(modifiedRelNode, listener); @@ -438,6 +450,9 @@ private static RelNode pushdownRules(RelNode input, RuleProfilingListener listen * Runs before {@link OpenSearchAggregateRule} marks the aggregate so the marking phase, the * Volcano split rule, and the {@code DistributedAggregateRewriter} see the rewritten shape: *

    + *
  • {@link OpenSearchAggregateStructKeyRule} — {@code GROUP BY } → + * {@code GROUP BY }, rebuilding the struct above the aggregate; + * grouping on a ROW value is ~10x slower than on the equivalent scalars. *
  • {@link OpenSearchCheckedLongSumRule} and {@link OpenSearchCheckedLongSumWindowRule} — * PPL's reflective {@code CHECKED_LONG_SUM} marker → Calcite's canonical {@code SUM}.
    • *
  • {@link OpenSearchDistinctCountRule} — single-arg {@code COUNT(DISTINCT x)} → @@ -451,6 +466,8 @@ private static RelNode pushdownRules(RelNode input, RuleProfilingListener listen private static RelNode decomposeAggregates(RelNode input, RuleProfilingListener listener) { return HepPhase.named("aggregate-decompose") .bottomUp() + // Runs first so the rules below see scalar group keys rather than a ROW value. + .addRuleInstance(new OpenSearchAggregateStructKeyRule()) .addRuleInstance(new OpenSearchCheckedLongSumRule()) .addRuleInstance(new OpenSearchCheckedLongSumWindowRule()) .addRuleInstance(new OpenSearchDistinctCountRule()) diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateStructKeyRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateStructKeyRule.java new file mode 100644 index 0000000000000..b6c2a7f2b23b9 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateStructKeyRule.java @@ -0,0 +1,295 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.planner.rules; + +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.util.ImmutableBitSet; +import org.opensearch.analytics.spi.MakeStructFunction; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Rewrites {@code GROUP BY } into {@code GROUP BY }, rebuilding + * the struct in a {@link LogicalProject} above the aggregate. + * + * 
    + * before:  LogicalProject(city=[make_struct('name', $1, 'pop', $2)])
+ *            LogicalAggregate(group=[{0}], count()=[COUNT()])          // grouping on a ROW value
+ *
+ * after:   LogicalProject(city=[make_struct('name', $0, 'pop', $1)], count()=[$2])
+ *            LogicalAggregate(group=[{n, n+1}], count()=[COUNT()])     // grouping on scalars
+ *              LogicalProject(..., $leaf refs appended...)
+ *
    + * + *

    Why. Grouping on a struct value is dramatically slower than grouping on the equivalent + * scalar columns — measured at ~10x on a 55-field object (586ms vs 58ms over 100k rows, with the + * result collapsed to a single row so response serialization can't account for it), and ~1.8x on a + * 4-field object. The struct carries no information its leaves don't, so the group keys are + * interchangeable: two rows produce equal structs exactly when all their leaves are equal. Note + * {@code named_struct} builds its {@code StructArray} with no validity buffer, so a materialized + * object is never itself NULL — there is no "NULL struct vs struct-of-all-NULLs" distinction for + * leaf grouping to lose. + * + *

    Placement. Runs in the {@code aggregate-decompose} phase, which exists for pre-marking + * rewrites on a plain {@code LogicalAggregate}. Two reasons it has to be here: {@code PROJECT_MERGE} + * has already run in the pushdown phase, so the {@code make_struct} is visible as an expression in + * the aggregate's input project rather than hidden behind an intervening projection; and the marking + * phase plus the distributed-aggregate rewriter must see the final group-key shape. + * + *

    The expanded struct expression is dropped when it becomes dead. Leaves are appended to + * the input project; the original {@code make_struct} is then removed unless an aggregate call still + * reads it, and the surviving indices are remapped. Leaving it in costs a struct construction per + * input row for a value nothing consumes — negligible on a narrow object but the dominant cost on a + * wide one, since it scales with field count: measured at parity with grouping raw leaves up to ~200 + * fields, and 1.87x at 500 fields before this removal. + * + *

    Dropping is skipped (the expression stays, which is correct but slower) when any aggregate call + * carries a distinct-key set or a within-group collation, since those hold input indices too and are + * not worth remapping for the gain. The remap never reaches beyond the aggregate this rule rebuilds, + * because the project added above restores the original row type. + * + *

    Bails out (leaving the plan untouched) on {@code GROUPING SETS} / {@code CUBE} / {@code ROLLUP}, + * and on any group key that is not literally a {@code make_struct} call — an {@code inputRef} to a + * struct produced further down is not expanded, since this rule reads the call's operands to find + * the leaves. + * + * @opensearch.internal + */ +public class OpenSearchAggregateStructKeyRule extends RelOptRule { + + public OpenSearchAggregateStructKeyRule() { + super(operand(LogicalAggregate.class, operand(LogicalProject.class, any())), "OpenSearchAggregateStructKeyRule"); + } + + @Override + public boolean matches(RelOptRuleCall ruleCall) { + LogicalAggregate agg = ruleCall.rel(0); + if (agg.getGroupType() != Aggregate.Group.SIMPLE || agg.getGroupSets().size() != 1) { + return false; + } + LogicalProject project = ruleCall.rel(1); + List exprs = project.getProjects(); + for (int key : agg.getGroupSet()) { + if (key < exprs.size() && isMakeStruct(exprs.get(key))) { + return true; + } + } + return false; + } + + @Override + public void onMatch(RelOptRuleCall ruleCall) { + LogicalAggregate agg = ruleCall.rel(0); + LogicalProject project = ruleCall.rel(1); + RexBuilder rexBuilder = agg.getCluster().getRexBuilder(); + List exprs = project.getProjects(); + List names = project.getRowType().getFieldNames(); + + // Group keys whose expression is a make_struct we can expand into its leaves. + Map expanded = new LinkedHashMap<>(); + for (int key : agg.getGroupSet()) { + if (isMakeStruct(exprs.get(key))) { + expanded.put(key, (RexCall) exprs.get(key)); + } + } + if (expanded.isEmpty()) { + return; + } + + // A struct expression is dead once its group key is expanded — unless an aggregate call + // still reads it. Dropping shifts every later index, so it is only attempted when the + // remap is confined to argument lists and filter arguments: a distinct-key set or a + // within-group collation also carries input indices, and rather than remap those we keep + // the expression (correct, just not as fast). Nothing downstream is affected either way, + // since the project added above restores the aggregate's original row type. + boolean remappable = agg.getAggCallList() + .stream() + .allMatch(c -> c.distinctKeys == null && c.collation.getFieldCollations().isEmpty()); + ImmutableBitSet.Builder referenced = ImmutableBitSet.builder(); + for (AggregateCall c : agg.getAggCallList()) { + c.getArgList().forEach(referenced::set); + if (c.filterArg >= 0) { + referenced.set(c.filterArg); + } + } + ImmutableBitSet aggRefs = referenced.build(); + Set dropped = new HashSet<>(); + if (remappable) { + for (int key : expanded.keySet()) { + if (aggRefs.get(key) == false) { + dropped.add(key); + } + } + } + + // Rebuild the input project with the leaves FIRST. Placing them at [0..N) makes the new + // group set exactly ImmutableBitSet.range(N) — a prefix by construction — which is what + // OpenSearchAggregateSplitRule.shouldSkipPartialFinalSplit requires to emit PARTIAL/FINAL. + // Appending instead leaves the keys after whatever else the project carries (an aggregate's + // argument column, say), and a key at index >= groupCount lands on PARTIAL's agg-output slot, + // so the gate refuses the split and the aggregate degrades to coordinator-gather. + Map leafIndex = new LinkedHashMap<>(); + List newExprs = new ArrayList<>(exprs.size()); + List newNames = new ArrayList<>(exprs.size()); + for (Map.Entry entry : expanded.entrySet()) { + String base = names.get(entry.getKey()); + for (RexNode leaf : leaves(entry.getValue())) { + leafIndex.computeIfAbsent(leaf, l -> { + newExprs.add(l); + newNames.add(uniquify(base + "$" + newExprs.size(), newNames)); + return newExprs.size() - 1; + }); + } + } + // Surviving originals follow, shifted up by the number of leaves. + int[] remap = new int[exprs.size()]; + for (int i = 0; i < exprs.size(); i++) { + if (dropped.contains(i)) { + remap[i] = -1; + continue; + } + remap[i] = newExprs.size(); + newExprs.add(exprs.get(i)); + newNames.add(uniquify(names.get(i), newNames)); + } + + LogicalProject widened = (LogicalProject) LogicalProject.create( + project.getInput(), + project.getHints(), + newExprs, + newNames, + project.getVariablesSet() + ); + + ImmutableBitSet.Builder groupBuilder = ImmutableBitSet.builder(); + for (int key : agg.getGroupSet()) { + RexCall structCall = expanded.get(key); + if (structCall == null) { + groupBuilder.set(remap[key]); + } else { + for (RexNode leaf : leaves(structCall)) { + groupBuilder.set(leafIndex.get(leaf)); + } + } + } + ImmutableBitSet newGroupSet = groupBuilder.build(); + + // Aggregate calls only ever reference kept expressions (that is what `dropped` guarantees), + // so remapping argument lists and filter arguments is sufficient. + List newCalls = new ArrayList<>(agg.getAggCallList().size()); + for (AggregateCall c : agg.getAggCallList()) { + if (dropped.isEmpty()) { + newCalls.add(c); + continue; + } + List args = new ArrayList<>(c.getArgList().size()); + for (int a : c.getArgList()) { + assert remap[a] >= 0 : "aggregate call references a dropped expression"; + args.add(remap[a]); + } + newCalls.add(c.copy(args, c.filterArg < 0 ? -1 : remap[c.filterArg])); + } + + LogicalAggregate newAgg = (LogicalAggregate) agg.copy(agg.getTraitSet(), widened, newGroupSet, List.of(newGroupSet), newCalls); + + // Restore the original output row type: struct keys are rebuilt from the grouped leaves, + // everything else is a passthrough reference. + int newGroupCount = newGroupSet.cardinality(); + List originalKeys = agg.getGroupSet().asList(); + List restored = new ArrayList<>(agg.getRowType().getFieldCount()); + for (int pos = 0; pos < agg.getRowType().getFieldCount(); pos++) { + if (pos < originalKeys.size()) { + int key = originalKeys.get(pos); + RexCall structCall = expanded.get(key); + if (structCall == null) { + restored.add(rexBuilder.makeInputRef(newAgg, newGroupSet.indexOf(remap[key]))); + } else { + restored.add(rebuild(rexBuilder, newAgg, structCall, newGroupSet, leafIndex)); + } + } else { + restored.add(rexBuilder.makeInputRef(newAgg, newGroupCount + (pos - originalKeys.size()))); + } + } + + ruleCall.transformTo(LogicalProject.create(newAgg, List.of(), restored, agg.getRowType().getFieldNames(), agg.getVariablesSet())); + } + + /** Appends a counter until {@code candidate} is absent from {@code taken}. */ + private static String uniquify(String candidate, List taken) { + String name = candidate; + for (int i = 0; taken.contains(name); i++) { + name = candidate + "_" + i; + } + return name; + } + + private static boolean isMakeStruct(RexNode node) { + return node instanceof RexCall call && call.getOperator() == MakeStructFunction.FUNCTION; + } + + /** + * The scalar value operands of a {@code make_struct} call, in order, descending through nested + * {@code make_struct} calls. Name literals (the even operands) are skipped. + */ + private static List leaves(RexCall structCall) { + List out = new ArrayList<>(); + collectLeaves(structCall, out); + return out; + } + + private static void collectLeaves(RexCall structCall, List out) { + List operands = structCall.getOperands(); + for (int i = 1; i < operands.size(); i += 2) { + RexNode value = operands.get(i); + if (isMakeStruct(value)) { + collectLeaves((RexCall) value, out); + } else { + out.add(value); + } + } + } + + /** Rebuilds {@code structCall}'s shape over the aggregate's grouped leaf outputs. */ + private static RexNode rebuild( + RexBuilder rexBuilder, + LogicalAggregate newAgg, + RexCall structCall, + ImmutableBitSet newGroupSet, + Map leafIndex + ) { + List operands = structCall.getOperands(); + List names = new ArrayList<>(operands.size() / 2); + List values = new ArrayList<>(operands.size() / 2); + for (int i = 0; i < operands.size(); i += 2) { + RelDataTypeField field = structCall.getType().getFieldList().get(i / 2); + names.add(field.getName()); + RexNode value = operands.get(i + 1); + if (isMakeStruct(value)) { + values.add(rebuild(rexBuilder, newAgg, (RexCall) value, newGroupSet, leafIndex)); + } else { + values.add(rexBuilder.makeInputRef(newAgg, newGroupSet.indexOf(leafIndex.get(value)))); + } + } + return MakeStructFunction.makeCall(rexBuilder, structCall.getType(), names, values); + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/engine/OpenSearchSchemaBuilderTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/engine/OpenSearchSchemaBuilderTests.java index 0cd11566a34e3..ea3c6dd251914 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/engine/OpenSearchSchemaBuilderTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/engine/OpenSearchSchemaBuilderTests.java @@ -259,10 +259,22 @@ public void testNestedObjectWithUnsupportedLeafDropsOnlyThatLeaf() throws Except assertNotNull(table); RelDataType rowType = table.getRowType(new org.apache.calcite.jdbc.JavaTypeFactoryImpl()); - assertEquals("Only 2 supported nested leaves should remain", 2, rowType.getFieldCount()); + // 2 supported leaves + the struct-typed `customer` parent column (the object itself is + // now addressable; ObjectStructMaterializer assembles it from these leaves). + assertEquals("2 supported nested leaves plus the object parent", 3, rowType.getFieldCount()); assertFieldType(rowType, "customer.id", SqlTypeName.VARCHAR); assertFieldType(rowType, "customer.age", SqlTypeName.INTEGER); assertNull("nested geo_point leaf must be dropped", rowType.getField("customer.home", true, false)); + + // The unsupported sub-field is dropped from the struct too, for the same reason it is + // dropped from the flat columns. + RelDataTypeField parent = rowType.getField("customer", true, false); + assertNotNull("object parent must be exposed as a column", parent); + assertEquals(SqlTypeName.ROW, parent.getType().getSqlTypeName()); + assertEquals("struct carries only the supported sub-fields", 2, parent.getType().getFieldCount()); + assertNotNull(parent.getType().getField("id", true, false)); + assertNotNull(parent.getType().getField("age", true, false)); + assertNull("geo_point sub-field must be dropped from the struct", parent.getType().getField("home", true, false)); } /** diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java index 1eb99a6e97403..c0b17a631e6cb 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java @@ -253,7 +253,7 @@ protected RelOptTable mockNullableTable(String tableName, String... fieldNames) return mockTable(tableName, rowTypeBuilder.build()); } - private RelOptTable mockTable(String tableName, RelDataType rowType) { + protected RelOptTable mockTable(String tableName, RelDataType rowType) { RelOptTable table = mock(RelOptTable.class); when(table.getQualifiedName()).thenReturn(List.of(tableName)); when(table.getRowType()).thenReturn(rowType); diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java index 52ca31f44d8c8..ff5e1faa3e33e 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java @@ -53,6 +53,9 @@ public class MockDataFusionBackend extends MockBackend implements SearchBackEndP SUPPORTED_TYPES.addAll(FieldType.keyword()); SUPPORTED_TYPES.addAll(FieldType.date()); SUPPORTED_TYPES.add(FieldType.BOOLEAN); + // Calcite ROW maps to FieldType.OBJECT — needed for make_struct's return type to + // resolve a capable backend in OpenSearchProjectRule. + SUPPORTED_TYPES.add(FieldType.OBJECT); } private static final Set STANDARD_OPS = Set.of( @@ -206,7 +209,9 @@ protected Set scanCapabilities() { ScalarFunction.CONCAT, ScalarFunction.UPPER, ScalarFunction.SIN, - ScalarFunction.ABS + ScalarFunction.ABS, + // Object materialization: make_struct over the flat dotted leaves (ObjectStructMaterializer). + ScalarFunction.MAKE_STRUCT ); private static final Set PROJECT_CAPS; diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ObjectStructPlanShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ObjectStructPlanShapeTests.java new file mode 100644 index 0000000000000..af5ba10504194 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ObjectStructPlanShapeTests.java @@ -0,0 +1,282 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.planner; + +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * End-to-end plan-shape coverage for {@link ObjectStructMaterializer}: an OpenSearch + * {@code object} field is re-assembled from its flat dotted leaf columns into a + * {@code make_struct} call in a project directly above the scan. + * + *

    The scan carries only leaf columns — an object parent has no physical storage + * ({@code FieldStorageResolver.populateFromProperties} recurses past object parents), so the + * struct is appended by the project, never read from the scan: + * + * 

    + * id                               INTEGER
+ * nested_metadata.top              VARCHAR
+ * nested_metadata.properties.name  VARCHAR
+ * nested_metadata.properties.value VARCHAR
+ *
    + */ +public class ObjectStructPlanShapeTests extends PlanShapeTestBase { + + /** Leaf columns backing {@code nested_metadata} in {@link #objectTable()}: top, name, value. */ + private static final int LEAF_COUNT = 3; + + /** + * Table as the schema builder now produces it for an {@code object} mapping: flat dotted leaf + * columns PLUS the struct-typed parent column. + */ + private RelOptTable objectTable() { + RelDataType varchar = typeFactory.createSqlType(SqlTypeName.VARCHAR); + RelDataType properties = typeFactory.createStructType(List.of(varchar, varchar), List.of("name", "value")); + RelDataType meta = typeFactory.createStructType(List.of(varchar, properties), List.of("top", "properties")); + + RelDataTypeFactory.Builder builder = typeFactory.builder(); + builder.add("id", typeFactory.createSqlType(SqlTypeName.INTEGER)); + builder.add("nested_metadata.top", varchar); + builder.add("nested_metadata.properties.name", varchar); + builder.add("nested_metadata.properties.value", varchar); + builder.add("nested_metadata", meta); + return mockTable("test_index", builder.build()); + } + + /** Field mappings for the leaf columns — the object parent is intentionally absent. */ + private Map> leafFieldMappings() { + return Map.of( + "id", + Map.of("type", "integer"), + "nested_metadata.top", + Map.of("type", "keyword"), + "nested_metadata.properties.name", + Map.of("type", "keyword"), + "nested_metadata.properties.value", + Map.of("type", "keyword") + ); + } + + /** + * The core rewrite: a project above the scan passes the leaves through and appends the + * object, nesting a second {@code make_struct} for the {@code properties} sub-object. + */ + public void testMaterializerAppendsNestedStructProjectAboveScan() { + RelNode scan = stubScan(objectTable()); + + Optional rewritten = ObjectStructMaterializer.rewrite(scan); + + assertTrue("materializer should fire when an object's leaves are present", rewritten.isPresent()); + assertPlanShape( + """ + LogicalProject(id=[$0], nested_metadata.top=[$1], nested_metadata.properties.name=[$2], nested_metadata.properties.value=[$3], nested_metadata=[make_struct('top':VARCHAR, $1, 'properties':VARCHAR, make_struct('name':VARCHAR, $2, 'value':VARCHAR, $3))]) + LogicalTableScan(table=[[test_index]]) + """, + rewritten.get() + ); + } + + /** No object spec ⇒ no rewrite, so plans without objects are untouched. */ + public void testMaterializerNoOpWithoutObjectSpec() { + RelNode scan = stubScan(mockTable("test_index", "status", "size")); + assertFalse(ObjectStructMaterializer.rewrite(scan).isPresent()); + } + + /** + * A struct column whose backing leaves are absent from the scan (e.g. an unsupported sub-field + * type was dropped from the schema) yields a typed NULL, never a partially-filled struct. + */ + public void testMaterializerEmitsNullForObjectWithMissingLeaf() { + RelDataType varchar = typeFactory.createSqlType(SqlTypeName.VARCHAR); + RelDataType meta = typeFactory.createStructType(List.of(varchar), List.of("top")); + RelDataTypeFactory.Builder builder = typeFactory.builder(); + builder.add("id", typeFactory.createSqlType(SqlTypeName.INTEGER)); + // NOTE: no "nested_metadata.top" leaf column — the struct cannot be assembled. + builder.add("nested_metadata", meta); + RelNode scan = stubScan(mockTable("test_index", builder.build())); + + RelNode rewritten = ObjectStructMaterializer.rewrite(scan).orElseThrow(); + + String shape = RelOptUtil.toString(rewritten); + assertFalse("no partial struct should be emitted, got:\n" + shape, shape.contains("make_struct")); + assertTrue("expected a typed NULL for the unassemblable object, got:\n" + shape, shape.contains("null:RecordType")); + } + + /** + * Projecting the object returns the whole object: the materialized struct survives the full + * planner (marking, CBO) and sits above the scan. + */ + public void testProjectOnObjectReturnsWholeObjectThroughPlanner() { + RelNode scan = stubScan(objectTable()); + RelNode materialized = ObjectStructMaterializer.rewrite(scan).orElseThrow(); + int objectIndex = materialized.getRowType().getFieldCount() - 1; + RelNode plan = LogicalProject.create( + materialized, + List.of(), + List.of(rexBuilder.makeInputRef(materialized, 0), rexBuilder.makeInputRef(materialized, objectIndex)), + List.of("id", "nested_metadata") + ); + + RelNode result = runPlanner(plan, buildContext("parquet", 1, leafFieldMappings())); + + String shape = RelOptUtil.toString(result); + assertTrue("expected a materialized struct, got:\n" + shape, shape.contains("make_struct")); + assertTrue( + "sub-object must nest a second make_struct, got:\n" + shape, + shape.indexOf("make_struct") != shape.lastIndexOf("make_struct") + ); + assertTrue("struct assembly must sit above the scan, got:\n" + shape, shape.indexOf("make_struct") < shape.indexOf("TableScan")); + } + + /** + * {@code stats count() by nested_metadata} — {@link org.opensearch.analytics.planner.rules.OpenSearchAggregateStructKeyRule} expands the + * struct group key into its leaf columns and rebuilds the object in a project above the + * aggregate. Grouping on a ROW value is ~10x slower than on the equivalent scalars, and the two + * are interchangeable because equal leaves imply equal structs. + * + *

    Grouping (rather than {@code count(nested_metadata)}) is the meaningful probe: counting a + * non-nullable column is equivalent to {@code count(*)}, so Calcite drops the column reference + * and trims the project away — correctly, but it proves nothing about the object. A group key + * genuinely needs the struct's value. + */ + public void testAggregateOnObjectGroupsByLeavesAndRebuildsAbove() { + RelNode scan = stubScan(objectTable()); + RelNode materialized = ObjectStructMaterializer.rewrite(scan).orElseThrow(); + int objectIndex = materialized.getRowType().getFieldCount() - 1; + AggregateCall count = AggregateCall.create( + SqlStdOperatorTable.COUNT, + false, + List.of(), + -1, + materialized, + typeFactory.createSqlType(SqlTypeName.BIGINT), + "cnt" + ); + RelNode plan = makeAggregate(materialized, ImmutableBitSet.of(objectIndex), count); + + RelNode result = runPlanner(plan, buildContext("parquet", 1, leafFieldMappings())); + + String shape = RelOptUtil.toString(result); + assertTrue("expected the object to be rebuilt, got:\n" + shape, shape.contains("make_struct")); + + // The aggregate must group on scalars, not on the ROW column. The struct sits at the last + // index of the materialized row type, so a group set still containing it means the rule + // didn't fire. + int aggLine = shape.indexOf("Aggregate(group=[{"); + assertTrue("expected an aggregate, got:\n" + shape, aggLine >= 0); + String groupSet = shape.substring(aggLine, shape.indexOf("]", aggLine) + 1); + assertFalse( + "aggregate must not group on the struct column (index " + objectIndex + "), got: " + groupSet, + groupSet.contains("{" + objectIndex + "}") + ); + assertEquals("expected one group key per leaf, got: " + groupSet, LEAF_COUNT, groupSet.split(",").length); + + // ...and the rebuild must sit ABOVE the aggregate now, which is the inverse of the old shape. + assertTrue("struct must be rebuilt above the aggregate, got:\n" + shape, shape.indexOf("make_struct") < aggLine); + + // The struct expression below the aggregate is dead once the group key is expanded, and the + // rule drops it. Leaving it in would build the object per input row for a value nothing + // reads — negligible on a narrow object, dominant on a wide one. Plan text is top-down, so + // "no make_struct at or after the aggregate line" means none survived underneath it. + assertTrue( + "the dead struct expression below the aggregate must be dropped, got:\n" + shape, + shape.lastIndexOf("make_struct") < aggLine + ); + } + + /** + * Multi-shard counterpart: the expanded group keys must stay a PREFIX of the input row type, or + * {@code OpenSearchAggregateSplitRule.shouldSkipPartialFinalSplit} refuses the PARTIAL/FINAL + * split and the aggregate silently degrades to coordinator-gather. + * + *

    That gate compares, for any group key at index {@code k >= groupCount}, the input column's + * {@link org.apache.calcite.sql.type.SqlTypeFamily} against the aggregate result that PARTIAL + * puts in that slot. Appending leaves while leaving the dead struct expression in place yields + * {@code group=[{1,2,3}]} with {@code groupCount=3} — non-prefix, VARCHAR leaf vs BIGINT COUNT, + * families differ, split skipped. Dropping the dead expression makes the keys {@code {0,1,2}} + * and the split proceeds. This test is the guard for that. + */ + public void testAggregateOnObject_2shard() { + RelNode scan = stubScan(objectTable()); + RelNode materialized = ObjectStructMaterializer.rewrite(scan).orElseThrow(); + int objectIndex = materialized.getRowType().getFieldCount() - 1; + AggregateCall count = AggregateCall.create( + SqlStdOperatorTable.COUNT, + false, + List.of(), + -1, + materialized, + typeFactory.createSqlType(SqlTypeName.BIGINT), + "cnt" + ); + RelNode plan = makeAggregate(materialized, ImmutableBitSet.of(objectIndex), count); + + RelNode result = runPlanner(plan, buildContext("parquet", 2, leafFieldMappings())); + + assertPlanShape( + """ + OpenSearchProject(nested_metadata=[ANNOTATED_PROJECT_EXPR(id=2, backends=[mock-parquet], make_struct('top':VARCHAR, $0, 'properties':VARCHAR, ANNOTATED_PROJECT_EXPR(id=1, backends=[mock-parquet], make_struct('name':VARCHAR, $1, 'value':VARCHAR, $2))))], cnt=[$3], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0, 1, 2}], cnt=[SUM($3)], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[], partitionCount=0]]) + OpenSearchAggregate(group=[{0, 1, 2}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchProject(nested_metadata$1=[$1], nested_metadata$2=[$2], nested_metadata$3=[$3], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + """, + result + ); + } + + /** + * The prefix property must hold even when the project carries columns besides the object — + * here {@code id}, needed as MAX's argument. If the expanded leaves land after it, the group set + * is non-prefix again and the PARTIAL/FINAL split is refused just as in the append-only case. + */ + public void testAggregateWithLeafArgGroupedByObject_2shard() { + RelNode scan = stubScan(objectTable()); + RelNode materialized = ObjectStructMaterializer.rewrite(scan).orElseThrow(); + int objectIndex = materialized.getRowType().getFieldCount() - 1; + AggregateCall max = AggregateCall.create( + SqlStdOperatorTable.MAX, + false, + List.of(0), + -1, + materialized, + typeFactory.createSqlType(SqlTypeName.INTEGER), + "max_id" + ); + RelNode plan = makeAggregate(materialized, ImmutableBitSet.of(objectIndex), max); + + RelNode result = runPlanner(plan, buildContext("parquet", 2, leafFieldMappings())); + + assertPlanShape( + """ + OpenSearchProject(nested_metadata=[ANNOTATED_PROJECT_EXPR(id=2, backends=[mock-parquet], make_struct('top':VARCHAR, $0, 'properties':VARCHAR, ANNOTATED_PROJECT_EXPR(id=1, backends=[mock-parquet], make_struct('name':VARCHAR, $1, 'value':VARCHAR, $2))))], max_id=[$3], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0, 1, 2}], max_id=[MAX($3)], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[], partitionCount=0]]) + OpenSearchAggregate(group=[{0, 1, 2}], max_id=[MAX($3)], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchProject(nested_metadata$1=[$1], nested_metadata$2=[$2], nested_metadata$3=[$3], id=[$0], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + """, + result + ); + } +} diff --git a/sandbox/qa/analytics-engine-rest/build.gradle b/sandbox/qa/analytics-engine-rest/build.gradle index 60c02bfe279b9..876363c9a515f 100644 --- a/sandbox/qa/analytics-engine-rest/build.gradle +++ b/sandbox/qa/analytics-engine-rest/build.gradle @@ -129,6 +129,31 @@ testClusters.integTest { setting 'analytics.delegation.lucene.blocked_predicates', '[]' } +// ── Interactive cluster for manual queries ────────────────────────────────── +// Foreground 1-node cluster with the SAME plugin set / feature flags / native-lib wiring the +// ITs use, so hand-run queries behave identically to the test suite: +// +// ./gradlew -Dsandbox.enabled=true :sandbox:qa:analytics-engine-rest:runAnalytics +// +// Then POST PPL to http://localhost:9200/_plugins/_ppl. Ctrl-C to stop. +// Requires the Rust native lib: :sandbox:libs:dataformat-native:buildRustLibrary (PROTOC must +// point at protoc 3+; the system protoc may be 2.x and fails on proto3 syntax). +testClusters { + runAnalytics { + // 1 node is the default; setting it explicitly throws "Cannot shrink cluster". + configureAnalyticsCluster(delegate) + setting 'analytics.delegation.lucene.blocked_predicates', '[]' + // NOTE: RunTask hardcodes http 9200 / transport 9300 and testclusters rejects an + // 'http.port' override, so port 9200 must be free before starting. + } +} + +task runAnalytics(type: org.opensearch.gradle.testclusters.RunTask) { + description = 'Runs a 1-node analytics-engine cluster in the foreground for manual queries' + group = 'Verification' + useCluster testClusters.runAnalytics +} + integTest { systemProperty 'tests.security.manager', 'false' // Forward plan-shape harness props to the forked test JVM (see PlanShapeGoldenTestBase): diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ApmServiceMapObjectIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ApmServiceMapObjectIT.java new file mode 100644 index 0000000000000..be392d313efc8 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ApmServiceMapObjectIT.java @@ -0,0 +1,233 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.qa; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +/** + * Regression contract for the dashboards-observability APM service-map experience. + * + *

    Pins the eight PPL queries in + * {@code dashboards-observability/public/components/apm/query_services/query_requests/ppl_queries.ts}. + * Six of them project parent object identifiers ({@code sourceNode.keyAttributes}, + * {@code targetNode.groupByAttributes}, …) and used to fail on the analytics engine with + * {@code IllegalArgumentException: Field [sourceNode.keyAttributes] not found} — thrown by + * {@code QualifiedNameResolver} because an {@code object} was not a column in the Calcite row + * type. Service-map topology, service detail panels, and operation/dependency lists went blank. + * Only the two {@code distinct_count} widgets (leaf scalars) worked. + * + *

    Now the schema exposes each object as a struct column and {@code ObjectStructMaterializer} + * assembles it with {@code make_struct} above the scan, so a parent-object identifier resolves + * for every operator — not just projection. {@link #testOperatorMatrixOnParentObject} covers the + * operator matrix from the bug report ({@code fields} / {@code eval} / {@code dedup} / + * {@code sort} / {@code isnotnull}), since the original report noted the failure was broader + * than projection. + * + *

    Fixture topology (3 connections): frontend→checkout, frontend→payment, checkout→inventory. + */ +public class ApmServiceMapObjectIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("apm_service_map", "otel-apm-service-map"); + + /** Back-quoted because the index name contains hyphens, exactly as the dashboards plugin sends it. */ + private static final String SOURCE = "source = `otel-apm-service-map`"; + + /** Time clause matching {@code buildTimeFilterClause}'s 'YYYY-MM-DD HH:mm:ss.SSS' rendering. */ + private static final String TIME = " | where timestamp >= '2026-08-26 00:00:00.000'" + + " and timestamp <= '2026-08-27 00:00:00.000'"; + + private static boolean dataProvisioned = false; + + @Override + protected void onBeforeQuery() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + // ── The six queries that projected parent objects (previously all HTTP 500) ──────── + + /** {@code getQueryListServices} — projects all four parent objects at once. */ + public void testListServices() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | dedup nodeConnectionHash" + + " | fields sourceNode.keyAttributes, sourceNode.groupByAttributes," + + " targetNode.keyAttributes, targetNode.groupByAttributes" + ); + assertEquals("one row per connection", 3, rows.size()); + // Every projected cell must be a materialized object, not a scalar or null. + for (List row : rows) { + assertEquals(4, row.size()); + for (Object cell : row) { + assertTrue("expected an object, got: " + cell, cell instanceof Map); + } + } + } + + /** {@code getQueryGetService} — parent objects plus the two leaf equality filters. */ + public void testGetService() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | where sourceNode.keyAttributes.environment = 'prod'" + + " | where sourceNode.keyAttributes.name = 'frontend'" + + " | dedup nodeConnectionHash" + + " | fields sourceNode.keyAttributes, sourceNode.groupByAttributes" + ); + // Two frontend-sourced connections (h1, h2) survive dedup on the hash. + assertEquals(2, rows.size()); + assertEquals("frontend", keyAttribute(rows.get(0).get(0), "name")); + assertEquals("js", keyAttribute(rows.get(0).get(1), "telemetry_sdk_language")); + } + + /** {@code getQueryServiceAttributes} — parent objects with {@code sort - timestamp | head 1}. */ + public void testServiceAttributes() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | where sourceNode.keyAttributes.environment = 'prod'" + + " | where sourceNode.keyAttributes.name = 'frontend'" + + " | fields sourceNode.keyAttributes, sourceNode.groupByAttributes, timestamp" + + " | sort - timestamp" + + " | head 1" + ); + assertEquals(1, rows.size()); + assertEquals("frontend", keyAttribute(rows.get(0).get(0), "name")); + // Descending sort picks the later of the two frontend rows (10:05, not 10:00). + assertTrue("expected the newest row, got: " + rows.get(0).get(2), rows.get(0).get(2).toString().contains("10:05")); + } + + /** + * {@code getQueryListServiceOperations} — mixes parent objects with leaf scalars in one + * projection. {@code getQueryListServiceDependencies} builds a byte-for-byte identical + * pipeline, so this covers both. + */ + public void testListServiceOperationsAndDependencies() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | dedup operationConnectionHash" + + " | fields sourceNode.keyAttributes, sourceOperation.name," + + " targetNode.keyAttributes, targetOperation.name" + ); + assertEquals(3, rows.size()); + for (List row : rows) { + assertTrue("col0 must be an object", row.get(0) instanceof Map); + assertTrue("col1 must be a scalar operation name", row.get(1) instanceof String); + assertTrue("col2 must be an object", row.get(2) instanceof Map); + assertTrue("col3 must be a scalar operation name", row.get(3) instanceof String); + } + } + + /** {@code getQueryGetServiceMap} — the topology query; four parent objects, interleaved order. */ + public void testGetServiceMap() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | dedup nodeConnectionHash" + + " | fields sourceNode.keyAttributes, targetNode.keyAttributes," + + " sourceNode.groupByAttributes, targetNode.groupByAttributes" + ); + assertEquals(3, rows.size()); + for (List row : rows) { + assertEquals(4, row.size()); + for (Object cell : row) { + assertTrue("expected an object, got: " + cell, cell instanceof Map); + } + } + } + + // ── The two count widgets that already worked — guard against regression ────────── + + /** {@code getQueryOperationDependenciesCount} — frontend's 'GET /cart' fans out to 2 services. */ + public void testOperationDependenciesCount() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | where sourceNode.keyAttributes.environment = 'prod'" + + " | where sourceNode.keyAttributes.name = 'frontend'" + + " | where sourceOperation.name = 'GET /cart'" + + " | stats distinct_count(targetNode.keyAttributes.name) as dependency_count" + ); + assertEquals(1, rows.size()); + assertEquals(2, ((Number) rows.get(0).get(0)).intValue()); + } + + /** {@code getQueryDependencyDownstreamCount} — checkout has 1 downstream dependency. */ + public void testDependencyDownstreamCount() throws IOException { + List> rows = rows( + SOURCE + + TIME + + " | where sourceNode.keyAttributes.environment = 'prod'" + + " | where sourceNode.keyAttributes.name = 'checkout'" + + " | stats distinct_count(targetNode.keyAttributes.name) as dependency_count" + ); + assertEquals(1, rows.size()); + assertEquals(1, ((Number) rows.get(0).get(0)).intValue()); + } + + // ── Operator matrix from the bug report ────────────────────────────────────────── + + /** + * Every operator that takes a parent-object identifier as an argument. The original report + * called out that the failure was not projection-specific: {@code eval}, {@code dedup}, + * {@code sort} and {@code isnotnull} threw the same "Field [...] not found" because the + * object was absent from the row type. One schema change fixes all of them, so all are + * pinned here. + */ + public void testOperatorMatrixOnParentObject() throws IOException { + // fields on a top-level object + assertFalse(rows(SOURCE + " | fields sourceNode | head 1").isEmpty()); + // fields on a one-level-deep object + assertFalse(rows(SOURCE + " | fields sourceNode.keyAttributes | head 1").isEmpty()); + // eval with an object on the RHS + assertFalse(rows(SOURCE + " | eval ka = sourceNode.keyAttributes | fields ka | head 1").isEmpty()); + // dedup keyed on an object + assertFalse(rows(SOURCE + " | dedup sourceNode.keyAttributes | head 1").isEmpty()); + // sort keyed on an object + assertFalse(rows(SOURCE + " | sort sourceNode.keyAttributes | head 1").isEmpty()); + // null-test on an object + assertFalse(rows(SOURCE + " | where isnotnull(sourceNode.keyAttributes) | fields nodeConnectionHash | head 1").isEmpty()); + } + + /** Control rows from the matrix: scalar-leaf filter and leaf-only aggregation still work. */ + public void testOperatorMatrixLeafControls() throws IOException { + List> filtered = rows( + SOURCE + " | where sourceNode.keyAttributes.name = 'frontend' | fields nodeConnectionHash" + ); + assertEquals(2, filtered.size()); + + List> aggregated = rows( + SOURCE + " | stats count() by sourceNode.keyAttributes.name" + ); + assertEquals("two distinct source services", 2, aggregated.size()); + } + + // ── helpers ────────────────────────────────────────────────────────────────────── + + @SuppressWarnings("unchecked") + private List> rows(String ppl) throws IOException { + Map response = executePpl(ppl); + List> rows = (List>) response.get("datarows"); + assertNotNull("Response missing 'datarows' for query: " + ppl, rows); + return rows; + } + + /** Reads a sub-field out of a materialized object cell. */ + @SuppressWarnings("unchecked") + private static Object keyAttribute(Object objectCell, String field) { + assertTrue("expected an object cell, got: " + objectCell, objectCell instanceof Map); + return ((Map) objectCell).get(field); + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldIT.java index 0aa3bad9318e4..9d83dd28123f1 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldIT.java @@ -102,18 +102,14 @@ public void testFilterOnDeeplyNestedObjectField() throws IOException { ); } - // ── Object-parent projection (gated on query-then-fetch) ────────────────── + // ── Object-parent projection ─────────────────────────────────────────────── // // Projecting an object parent (top-level "city" or intermediate "city.location") - // returns a nested JSON value reconstructed from _source. Analytics-engine emits - // only flat leaves into the Calcite row type today, so parent references fall - // through QualifiedNameResolver and throw "Field [city.location] not found". - // - // Support requires query-then-fetch (QTF): coordinator returns docIds post-filter, - // a fetch stage pulls the doc from the shard, and the parent sub-object is - // reconstructed from _source or from parquet rows. QTF is tracked separately. + // returns the nested object. No query-then-fetch / _source read is needed: the + // schema exposes the object as a struct (ROW) column and ObjectStructMaterializer + // re-assembles it with make_struct over the flat leaf columns the scan already + // produces, in a project directly above the scan. - @AwaitsFix(bugUrl = "Object parent projection requires query-then-fetch (QTF) for source-based materialization") public void testSelectIntermediateObjectField() throws IOException { assertRowsEqual( "source=" + DATASET.indexName + " | fields city.location | head 1", @@ -121,7 +117,6 @@ public void testSelectIntermediateObjectField() throws IOException { ); } - @AwaitsFix(bugUrl = "Object parent projection requires query-then-fetch (QTF) for source-based materialization") public void testSelectTopLevelObjectField() throws IOException { assertRowsEqual( "source=" + DATASET.indexName + " | fields city | head 1", @@ -129,7 +124,6 @@ public void testSelectTopLevelObjectField() throws IOException { ); } - @AwaitsFix(bugUrl = "Object parent projection requires query-then-fetch (QTF) for source-based materialization") public void testSelectTopLevelObjectFieldWithSiblings() throws IOException { assertRowsEqual( "source=" + DATASET.indexName + " | fields city, account | head 1", @@ -140,7 +134,6 @@ public void testSelectTopLevelObjectFieldWithSiblings() throws IOException { ); } - @AwaitsFix(bugUrl = "Object parent projection requires query-then-fetch (QTF) for source-based materialization") public void testSelectParentAndLeafMixed() throws IOException { assertRowsEqual( "source=" + DATASET.indexName + " | fields city.name, city.location | head 1", @@ -148,8 +141,38 @@ public void testSelectParentAndLeafMixed() throws IOException { ); } + // ── Aggregation involving object fields ─────────────────────────────────── + // + // Leaf aggregations (min/max/sum on city.population, city.location.latitude, …) are covered + // above. These cover aggregating on the OBJECT VALUE itself — the group key is a struct + // materialized by ObjectStructMaterializer, so the aggregate receives an assembled object. + + /** Group by an intermediate object ({@code city.location}) — 3 distinct locations. */ + public void testGroupByIntermediateObjectField() throws IOException { + assertRowCount("source=" + DATASET.indexName + " | stats count() by city.location", 3); + } + + /** Group by a top-level object ({@code city}) — 3 distinct cities. */ + public void testGroupByTopLevelObjectField() throws IOException { + assertRowCount("source=" + DATASET.indexName + " | stats count() by city", 3); + } + + /** Aggregate a leaf while grouping by an object value. */ + public void testAggregateLeafGroupedByObjectField() throws IOException { + assertRowCount("source=" + DATASET.indexName + " | stats max(city.population) by city.location", 3); + } + // ── helpers (mirrored from FieldsCommandIT) ──────────────────────────────── + /** Asserts only the row count — group order is not deterministic for a struct key. */ + private void assertRowCount(String ppl, int expected) throws IOException { + Map response = executePpl(ppl); + @SuppressWarnings("unchecked") + List> actualRows = (List>) response.get("datarows"); + assertNotNull("Response missing 'datarows' for query: " + ppl, actualRows); + assertEquals("Row count mismatch for query: " + ppl, expected, actualRows.size()); + } + private static List row(Object... values) { return Arrays.asList(values); } @@ -173,4 +196,70 @@ private final void assertRowsEqual(String ppl, List... expected) throws } + + // ── select * (no explicit field list) ───────────────────────────────────────────── + // + // The default query shape, and the one most likely to regress: nothing here names an object, + // so coverage depends entirely on how `*` expands. Verified against the legacy (lucene-only) + // engine on the same mapping — legacy returns the three top-level fields with the objects as + // nested JSON, and these assert we match it. + // + // Note the flat dotted leaves the scan actually reads (`city.name`, + // `city.location.latitude`, …) must NOT appear: they are addressable by name but are not part + // of `*`, so an object's data is returned once, as a struct, not twice. + + /** {@code source=idx} with no field list: objects come back as whole nested values. */ + public void testSelectStarReturnsObjectsAsNestedStructs() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | head 1"); + assertStarShape(response, "source=... | head 1"); + } + + /** Explicit {@code fields *} must behave identically to the implicit form above. */ + public void testFieldsStarReturnsObjectsAsNestedStructs() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | fields * | head 1"); + assertStarShape(response, "source=... | fields * | head 1"); + } + + /** + * Asserts the star-expansion contract: exactly the top-level fields (no dotted leaves), with + * each object materialized as a nested map. Column order is not asserted — it is not part of + * the contract and differs from legacy — so the row is checked by column name. + */ + private void assertStarShape(Map response, String context) { + List columns = extractColumnNames(response); + assertEquals( + "star expansion must yield only top-level fields (no dotted leaves) for " + context, + List.of("account", "city", "id"), + columns.stream().sorted().toList() + ); + + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("datarows"); + assertNotNull("missing datarows for " + context, rows); + assertEquals("expected a single row for " + context, 1, rows.size()); + Map row = new java.util.HashMap<>(); + for (int i = 0; i < columns.size(); i++) { + row.put(columns.get(i), rows.get(0).get(i)); + } + + assertEquals("id for " + context, "1", row.get("id")); + assertEquals( + "account must be a whole nested object for " + context, + Map.of("owner", "alice", "balance", 1000.5), + row.get("account") + ); + // Nested sub-object arrives nested, not flattened to a dotted key. + assertEquals( + "city must nest location for " + context, + Map.of( + "name", + "Seattle", + "population", + 750000, + "location", + Map.of("latitude", 47.6062, "longitude", -122.3321) + ), + row.get("city") + ); + } } diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldMultiShardIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldMultiShardIT.java new file mode 100644 index 0000000000000..4e2a46df524db --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ObjectFieldMultiShardIT.java @@ -0,0 +1,99 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.qa; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +/** + * Multi-shard coverage for grouping on an {@code object}. {@link ObjectFieldIT} exercises the same + * queries at 1 shard, which cannot reach the PARTIAL/FINAL reduce path. + * + *

    Why this matters: {@code OpenSearchAggregateStructKeyRule} expands a struct group key into the + * object's leaf columns. If those keys are not a PREFIX of the aggregate's input row type, + * {@code OpenSearchAggregateSplitRule.shouldSkipPartialFinalSplit} refuses the split (a group key at + * index {@code k >= groupCount} lands on PARTIAL's agg-output slot, and a family mismatch — VARCHAR + * leaf vs BIGINT COUNT — trips the gate). The aggregate then silently degrades to coordinator-gather: + * still correct, but no longer distributed, and invisible at 1 shard. The rule keeps the keys a prefix + * by dropping the struct expression once it is dead; + * {@code ObjectStructPlanShapeTests#testAggregateOnObject_2shard} pins the resulting plan, and this + * pins the results it produces. + * + *

    Reuses {@code object_fields}' mapping and bulk data under a distinct index name so the 1-shard + * index in {@link ObjectFieldIT} is untouched. + */ +public class ObjectFieldMultiShardIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("object_fields", "object_fields_multishard"); + + /** Seattle, Portland, Austin — one group per city in the dataset. */ + private static final int DISTINCT_CITIES = 3; + + private static boolean provisioned = false; + + @Override + protected void onBeforeQuery() throws IOException { + if (provisioned == false) { + DatasetProvisioner.provision(client(), DATASET, 2); + provisioned = true; + } + } + + /** + * Grouping on the whole object across shards: one group per city, every group key returned as a + * nested object (not flattened), and the per-shard counts reduced to the full document count. + */ + @SuppressWarnings("unchecked") + public void testGroupByTopLevelObjectFieldAtTwoShards() throws IOException { + List> rows = rowsOf("source=" + DATASET.indexName + " | stats count() by city"); + assertEquals("one group per city", DISTINCT_CITIES, rows.size()); + + long total = 0; + for (List row : rows) { + // stats output is [count, groupKey]; locate the map rather than assuming a position. + Map city = null; + for (Object cell : row) { + if (cell instanceof Map map) { + city = (Map) map; + } else if (cell instanceof Number n) { + total += n.longValue(); + } + } + assertNotNull("group key must be a nested object, got row: " + row, city); + assertTrue("object must carry its leaves, got: " + city, city.containsKey("name") && city.containsKey("population")); + assertTrue("sub-object must stay nested, got: " + city, city.get("location") instanceof Map); + } + assertEquals("reduced counts must cover every document", DISTINCT_CITIES, total); + } + + /** An intermediate object as the group key — the nested make_struct path, across shards. */ + public void testGroupByIntermediateObjectFieldAtTwoShards() throws IOException { + assertEquals( + DISTINCT_CITIES, + rowsOf("source=" + DATASET.indexName + " | stats count() by city.location").size() + ); + } + + /** A leaf aggregate grouped by an object: exercises FINAL's agg-call remap alongside struct keys. */ + public void testAggregateLeafGroupedByObjectFieldAtTwoShards() throws IOException { + assertEquals( + DISTINCT_CITIES, + rowsOf("source=" + DATASET.indexName + " | stats max(city.population) by city.location").size() + ); + } + + @SuppressWarnings("unchecked") + private List> rowsOf(String ppl) throws IOException { + Map response = executePpl(ppl); + List> rows = (List>) response.get("datarows"); + assertNotNull("Response missing 'datarows' for query: " + ppl, rows); + return rows; + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/apm_service_map/bulk.json b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/apm_service_map/bulk.json new file mode 100644 index 0000000000000..1a4b90b7f6620 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/apm_service_map/bulk.json @@ -0,0 +1,6 @@ +{"index": {}} +{"timestamp": "2026-08-26T10:00:00.000Z", "nodeConnectionHash": "h1", "operationConnectionHash": "o1", "sourceNode": {"keyAttributes": {"name": "frontend", "environment": "prod", "type": "Service"}, "groupByAttributes": {"telemetry_sdk_language": "js", "deployment_env": "prod"}}, "targetNode": {"keyAttributes": {"name": "checkout", "environment": "prod", "type": "Service"}, "groupByAttributes": {"telemetry_sdk_language": "java", "deployment_env": "prod"}}, "sourceOperation": {"name": "GET /cart"}, "targetOperation": {"name": "POST /checkout"}} +{"index": {}} +{"timestamp": "2026-08-26T10:05:00.000Z", "nodeConnectionHash": "h2", "operationConnectionHash": "o2", "sourceNode": {"keyAttributes": {"name": "frontend", "environment": "prod", "type": "Service"}, "groupByAttributes": {"telemetry_sdk_language": "js", "deployment_env": "prod"}}, "targetNode": {"keyAttributes": {"name": "payment", "environment": "prod", "type": "Service"}, "groupByAttributes": {"telemetry_sdk_language": "go", "deployment_env": "prod"}}, "sourceOperation": {"name": "GET /cart"}, "targetOperation": {"name": "POST /pay"}} +{"index": {}} +{"timestamp": "2026-08-26T10:10:00.000Z", "nodeConnectionHash": "h3", "operationConnectionHash": "o3", "sourceNode": {"keyAttributes": {"name": "checkout", "environment": "prod", "type": "Service"}, "groupByAttributes": {"telemetry_sdk_language": "java", "deployment_env": "prod"}}, "targetNode": {"keyAttributes": {"name": "inventory", "environment": "prod", "type": "Service"}, "groupByAttributes": {"telemetry_sdk_language": "java", "deployment_env": "prod"}}, "sourceOperation": {"name": "POST /checkout"}, "targetOperation": {"name": "GET /stock"}} diff --git a/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/apm_service_map/mapping.json b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/apm_service_map/mapping.json new file mode 100644 index 0000000000000..9c224af5deb08 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/resources/datasets/apm_service_map/mapping.json @@ -0,0 +1,63 @@ +{ + "settings": { + "number_of_shards": 1, + "number_of_replicas": 0 + }, + "mappings": { + "properties": { + "timestamp": { + "type": "date" + }, + "nodeConnectionHash": { + "type": "keyword" + }, + "operationConnectionHash": { + "type": "keyword" + }, + "sourceNode": { + "properties": { + "keyAttributes": { + "properties": { + "name": { "type": "keyword" }, + "environment": { "type": "keyword" }, + "type": { "type": "keyword" } + } + }, + "groupByAttributes": { + "properties": { + "telemetry_sdk_language": { "type": "keyword" }, + "deployment_env": { "type": "keyword" } + } + } + } + }, + "targetNode": { + "properties": { + "keyAttributes": { + "properties": { + "name": { "type": "keyword" }, + "environment": { "type": "keyword" }, + "type": { "type": "keyword" } + } + }, + "groupByAttributes": { + "properties": { + "telemetry_sdk_language": { "type": "keyword" }, + "deployment_env": { "type": "keyword" } + } + } + } + }, + "sourceOperation": { + "properties": { + "name": { "type": "keyword" } + } + }, + "targetOperation": { + "properties": { + "name": { "type": "keyword" } + } + } + } + } +}