From f86ab424e64b63dbb1ea1297e1e2f92af8648bed Mon Sep 17 00:00:00 2001 From: Bharathwaj G Date: Fri, 3 Apr 2026 23:19:48 +0530 Subject: [PATCH 1/6] Include protoc in codeql (#21108) * Fix protoc in codeql Signed-off-by: bharath-techie --- .github/workflows/codeql-analysis.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml index 2186c166e19a6..e9552e0537f9b 100644 --- a/.github/workflows/codeql-analysis.yml +++ b/.github/workflows/codeql-analysis.yml @@ -51,6 +51,9 @@ jobs: # Prefix the list here with "+" to use these queries and those in the config file. # queries: ./path/to/local/query, your-org/your-repo/queries@main + - name: Set up protoc + uses: arduino/setup-protoc@v3 + # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). # If this step fails, then you should remove it and run the build manually (see below) - name: Autobuild From 5de15074693c5668f007271b46addb7417baef20 Mon Sep 17 00:00:00 2001 From: Neetika Singhal Date: Fri, 3 Apr 2026 13:38:55 -0700 Subject: [PATCH 2/6] Change access modifier for convert::FieldValueFetcher to allow other Field Fetchers to override the method (#21005) Signed-off-by: Neetika Singhal --- .../java/org/opensearch/index/mapper/FieldValueFetcher.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/index/mapper/FieldValueFetcher.java b/server/src/main/java/org/opensearch/index/mapper/FieldValueFetcher.java index aeacf235591ad..82df8a4b9d08f 100644 --- a/server/src/main/java/org/opensearch/index/mapper/FieldValueFetcher.java +++ b/server/src/main/java/org/opensearch/index/mapper/FieldValueFetcher.java @@ -40,7 +40,7 @@ protected FieldValueFetcher(String simpleName) { * Converts the field value to required representation, should be overridden by field mappers as needed * @param value - value to convert */ - Object convert(Object value) { + public Object convert(Object value) { return value; } From e72d280265b3bcde8c1d0adefdeb318152eac224 Mon Sep 17 00:00:00 2001 From: expani Date: Tue, 31 Mar 2026 20:38:45 -0700 Subject: [PATCH 3/6] Initial commit for Planner with marking on projects, filters and aggregates with unit tests Rebased with changes on main Signed-off-by: expani --- .../analytics/spi/AggregateFunction.java | 59 ++---- .../spi/AnalyticsSearchBackendPlugin.java | 39 +++- .../analytics/spi/FieldTypeFamily.java | 60 ++++++ .../analytics/spi/FilterCapability.java | 35 +--- .../analytics/spi/FilterOperator.java | 82 ++------ .../analytics/spi/FullTextOperator.java | 32 +++ .../analytics/exec/DefaultPlanExecutor.java | 6 +- .../planner/CapabilityResolutionUtils.java | 176 +++++++++++++---- .../analytics/planner/FieldStorageInfo.java | 43 ++-- .../planner/FieldStorageResolver.java | 147 +++----------- .../analytics/planner/PlannerContext.java | 23 ++- .../analytics/planner/RelNodeUtils.java | 10 +- .../rel/AnnotatedProjectExpression.java | 14 +- .../planner/rel/FullTextFunctions.java | 26 ++- .../planner/rel/OpenSearchAggregate.java | 16 +- .../rel/OpenSearchDistributionTraitDef.java | 28 ++- .../rel/OpenSearchExchangeReducer.java | 20 +- .../planner/rel/OpenSearchExchangeWriter.java | 27 ++- .../planner/rel/OpenSearchFilter.java | 13 +- .../planner/rel/OpenSearchProject.java | 44 ++--- .../planner/rel/OpenSearchRelNode.java | 12 +- .../planner/rel/OpenSearchShuffleReader.java | 20 +- .../analytics/planner/rel/OpenSearchSort.java | 13 +- .../planner/rel/OpenSearchTableScan.java | 17 +- .../rules/OpenSearchAggregateRule.java | 106 ++++------ .../rules/OpenSearchAggregateSplitRule.java | 4 + .../planner/rules/OpenSearchFilterRule.java | 184 ++++++++++-------- .../planner/rules/OpenSearchProjectRule.java | 157 ++++----------- .../planner/rules/OpenSearchSortRule.java | 17 +- .../rules/OpenSearchTableScanRule.java | 54 +---- .../analytics/planner/AggregateRuleTests.java | 56 ++---- .../planner/BasePlannerRulesTests.java | 2 +- .../analytics/planner/FilterRuleTests.java | 67 ++----- .../planner/MockDataFusionBackend.java | 104 +++++----- .../analytics/planner/MockLuceneBackend.java | 103 ++++------ .../analytics/planner/ProjectRuleTests.java | 124 +++++------- 36 files changed, 888 insertions(+), 1052 deletions(-) create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java index b72e794e93684..84cce5710b1b2 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java @@ -11,57 +11,40 @@ import org.apache.calcite.sql.SqlKind; /** - * Aggregate functions that a backend may support, categorized by {@link Type}. + * Aggregate functions that a backend may support. + * Used by the aggregate rule to verify the backend can handle + * every {@link org.apache.calcite.rel.core.AggregateCall} in the plan. * *

Note: {@code COUNT} covers both {@code COUNT(*)} and {@code COUNT(DISTINCT x)}. * The distinction is on {@code AggregateCall.isDistinct()}, not on SqlKind. + * Backends that only support non-distinct count should check distinctness + * separately during fragment conversion. * * @opensearch.internal */ public enum AggregateFunction { - // Simple — fixed-size state per key - SUM(Type.SIMPLE, SqlKind.SUM), - SUM0(Type.SIMPLE, SqlKind.SUM0), - MIN(Type.SIMPLE, SqlKind.MIN), - MAX(Type.SIMPLE, SqlKind.MAX), - COUNT(Type.SIMPLE, SqlKind.COUNT), - AVG(Type.SIMPLE, SqlKind.AVG), + SUM(SqlKind.SUM), + SUM0(SqlKind.SUM0), + MIN(SqlKind.MIN), + MAX(SqlKind.MAX), + COUNT(SqlKind.COUNT), + APPROX_COUNT_DISTINCT(SqlKind.OTHER), + AVG(SqlKind.AVG), + STDDEV_POP(SqlKind.STDDEV_POP), + STDDEV_SAMP(SqlKind.STDDEV_SAMP), + VAR_POP(SqlKind.VAR_POP), + VAR_SAMP(SqlKind.VAR_SAMP), + PERCENTILE_CONT(SqlKind.PERCENTILE_CONT), + PERCENTILE_DISC(SqlKind.PERCENTILE_DISC), + COLLECT(SqlKind.COLLECT), + LISTAGG(SqlKind.LISTAGG); - // Statistical — fixed-size state, multi-pass or running stats - STDDEV_POP(Type.STATISTICAL, SqlKind.STDDEV_POP), - STDDEV_SAMP(Type.STATISTICAL, SqlKind.STDDEV_SAMP), - VAR_POP(Type.STATISTICAL, SqlKind.VAR_POP), - VAR_SAMP(Type.STATISTICAL, SqlKind.VAR_SAMP), - - // State-expanding — state grows with input rows per key - PERCENTILE_CONT(Type.STATE_EXPANDING, SqlKind.PERCENTILE_CONT), - PERCENTILE_DISC(Type.STATE_EXPANDING, SqlKind.PERCENTILE_DISC), - COLLECT(Type.STATE_EXPANDING, SqlKind.COLLECT), - LISTAGG(Type.STATE_EXPANDING, SqlKind.LISTAGG), - - // Approximate — probabilistic, fixed-size state - APPROX_COUNT_DISTINCT(Type.APPROXIMATE, SqlKind.OTHER); - - /** Category of aggregate function. Affects execution strategy (shuffle vs map-reduce). */ - public enum Type { - SIMPLE, - STATISTICAL, - STATE_EXPANDING, - APPROXIMATE - } - - private final Type type; private final SqlKind sqlKind; - AggregateFunction(Type type, SqlKind sqlKind) { - this.type = type; + AggregateFunction(SqlKind sqlKind) { this.sqlKind = sqlKind; } - public Type getType() { - return type; - } - public SqlKind getSqlKind() { return sqlKind; } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index c823763e2040d..827ee1e527d96 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java @@ -8,6 +8,10 @@ package org.opensearch.analytics.spi; + +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.analytics.backend.ExecutionContext; +import org.opensearch.analytics.backend.SearchExecEngine; import org.opensearch.index.engine.dataformat.DataFormat; import java.util.Collections; @@ -20,11 +24,25 @@ */ public interface AnalyticsSearchBackendPlugin extends SearchExecEngineProvider { + /** Unique engine name (e.g., "lucene", "datafusion"). */ + String name(); + + /** + * Creates a searcher bound to the given reader snapshot. + * @param ctx the execution context + */ + SearchExecEngine searcher(ExecutionContext ctx); + /** Returns the data formats supported by this backend. */ List getSupportedFormats(); - /** Filter capabilities scoped to operator, field type, and data format. */ - default Set filterCapabilities() { + /** Filter operators this backend can evaluate, scoped by field type family. */ + default Set supportedFilterCapabilities() { + return Collections.emptySet(); + } + + /** Full-text operators this backend can evaluate on indexed fields. */ + default Set supportedFullTextOperators() { return Collections.emptySet(); } @@ -43,18 +61,22 @@ default Set acceptedDelegations() { return Collections.emptySet(); } - /** Aggregate capabilities scoped to function, field type, and data format. */ - default Set aggregateCapabilities() { + /** Aggregate functions this backend can evaluate (SUM, AVG, COUNT, etc.). */ + default Set supportedAggregateFunctions() { return Collections.emptySet(); } - /** Window capabilities scoped to function, field type, and data format. */ - default Set windowCapabilities() { + /** Scalar functions this backend can evaluate in projections (UPPER, CAST, math ops, etc.). */ + default Set supportedScalarFunctions() { return Collections.emptySet(); } - /** Project capabilities: scalar functions and opaque operations, scoped to data format. */ - default Set projectCapabilities() { + /** + * Names of opaque project operations this backend can handle (e.g. "painless", "highlight", "suggest"). + * Used to resolve UnresolvedRexNode from frontend plugins to a backend. + * Analytics Core is agnostic to what these names mean — backends define and consume them. + */ + default Set supportedOpaqueProjectOperations() { return Collections.emptySet(); } @@ -71,5 +93,4 @@ default Set supportedShuffleCapabilities() { default byte[] convertFragment(Object fragment) { throw new UnsupportedOperationException("convertFragment not yet implemented for " + name()); } - } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java new file mode 100644 index 0000000000000..9830f3377419d --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java @@ -0,0 +1,60 @@ +/* + * 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; + +/** + * Logical field type families for capability matching. + * Groups OpenSearch {@code MappedFieldType.typeName()} strings into + * coarse categories that backends declare support for. + * + * @opensearch.internal + */ +public enum FieldTypeFamily { + NUMERIC, + KEYWORD, + TEXT, + DATE, + BOOLEAN, + IP, + GEO_POINT, + GEO_SHAPE, + BINARY, + NESTED, + OBJECT, + RANGE, + COMPLETION; + + /** + * Maps an OpenSearch mapping type string ({@code MappedFieldType.typeName()}) + * to a FieldTypeFamily. Returns null if the type is not recognized. + */ + public static FieldTypeFamily fromMappingType(String mappingType) { + if (mappingType == null) { + return null; + } + return switch (mappingType) { + case "integer", "long", "short", "byte", "float", "double", + "half_float", "scaled_float", "unsigned_long" -> NUMERIC; + case "keyword", "constant_keyword", "wildcard" -> KEYWORD; + case "text", "match_only_text" -> TEXT; + case "date", "date_nanos" -> DATE; + case "boolean" -> BOOLEAN; + case "ip" -> IP; + case "geo_point", "point" -> GEO_POINT; + case "geo_shape", "shape" -> GEO_SHAPE; + case "binary" -> BINARY; + case "nested" -> NESTED; + case "object", "flat_object" -> OBJECT; + case "integer_range", "float_range", "long_range", + "double_range", "date_range", "ip_range" -> RANGE; + case "completion" -> COMPLETION; + default -> null; + }; + } +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java index 3f7e8f05096a8..3fdfcb3f76e11 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java @@ -8,37 +8,18 @@ package org.opensearch.analytics.spi; -import java.util.Set; - /** - * Declares a backend's ability to evaluate filter predicates, scoped to data formats. - * Three variants for the three categories of filter operations. + * Declares that a backend can evaluate a specific {@link FilterOperator} + * on a specific {@link FieldTypeFamily}. + * + *

Example: {@code FilterCapability.of(EQUALS, KEYWORD)} means the backend + * can evaluate equality predicates on keyword fields. * * @opensearch.internal */ -public sealed interface FilterCapability { - - /** Standard comparison filter (EQUALS, GT, IN, LIKE, etc.) on a field type in given formats. */ - record Standard(FilterOperator operator, FieldType fieldType, - Set formats) implements FilterCapability { - public Standard { - formats = Set.copyOf(formats); - } - } - - /** Full-text filter (MATCH, MATCH_PHRASE, FUZZY, etc.) with supported query parameters. */ - record FullText(FilterOperator operator, FieldType fieldType, - Set formats, Set supportedParams) implements FilterCapability { - public FullText { - formats = Set.copyOf(formats); - supportedParams = Set.copyOf(supportedParams); - } - } +public record FilterCapability(FilterOperator operator, FieldTypeFamily fieldTypeFamily) { - /** Expression-based filter on derived columns (e.g., HAVING after aggregate). */ - record Expression(Set formats) implements FilterCapability { - public Expression { - formats = Set.copyOf(formats); - } + public static FilterCapability of(FilterOperator operator, FieldTypeFamily fieldTypeFamily) { + return new FilterCapability(operator, fieldTypeFamily); } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java index 70b3262f9b8b1..d60be4445bb0d 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java @@ -11,83 +11,35 @@ import org.apache.calcite.sql.SqlKind; /** - * All filter operations a backend may support, covering standard comparisons, - * full-text search, and expression-based filtering. - * - *

Each operator carries a {@link Type} indicating its category and whether - * it supports parameters (e.g., full-text operators accept analyzer, slop, etc.). + * Standard comparison/predicate operations that a backend may support. * * @opensearch.internal */ public enum FilterOperator { + EQUALS(SqlKind.EQUALS), + NOT_EQUALS(SqlKind.NOT_EQUALS), + GREATER_THAN(SqlKind.GREATER_THAN), + GREATER_THAN_OR_EQUAL(SqlKind.GREATER_THAN_OR_EQUAL), + LESS_THAN(SqlKind.LESS_THAN), + LESS_THAN_OR_EQUAL(SqlKind.LESS_THAN_OR_EQUAL), + IS_NULL(SqlKind.IS_NULL), + IS_NOT_NULL(SqlKind.IS_NOT_NULL), + IN(SqlKind.IN), + LIKE(SqlKind.LIKE), + PREFIX(SqlKind.OTHER), + REGEXP(SqlKind.OTHER), + WILDCARD(SqlKind.OTHER); - // Standard comparison - EQUALS(Type.STANDARD, SqlKind.EQUALS), - NOT_EQUALS(Type.STANDARD, SqlKind.NOT_EQUALS), - GREATER_THAN(Type.STANDARD, SqlKind.GREATER_THAN), - GREATER_THAN_OR_EQUAL(Type.STANDARD, SqlKind.GREATER_THAN_OR_EQUAL), - LESS_THAN(Type.STANDARD, SqlKind.LESS_THAN), - LESS_THAN_OR_EQUAL(Type.STANDARD, SqlKind.LESS_THAN_OR_EQUAL), - IS_NULL(Type.STANDARD, SqlKind.IS_NULL), - IS_NOT_NULL(Type.STANDARD, SqlKind.IS_NOT_NULL), - IN(Type.STANDARD, SqlKind.IN), - LIKE(Type.STANDARD, SqlKind.LIKE), - PREFIX(Type.STANDARD, SqlKind.OTHER), - - // Full-text search - MATCH(Type.FULL_TEXT, SqlKind.OTHER), - MATCH_PHRASE(Type.FULL_TEXT, SqlKind.OTHER), - MATCH_PHRASE_PREFIX(Type.FULL_TEXT, SqlKind.OTHER), - MATCH_BOOL_PREFIX(Type.FULL_TEXT, SqlKind.OTHER), - MULTI_MATCH(Type.FULL_TEXT, SqlKind.OTHER), - QUERY_STRING(Type.FULL_TEXT, SqlKind.OTHER), - SIMPLE_QUERY_STRING(Type.FULL_TEXT, SqlKind.OTHER), - FUZZY(Type.FULL_TEXT, SqlKind.OTHER), - WILDCARD(Type.FULL_TEXT, SqlKind.OTHER), - REGEXP(Type.FULL_TEXT, SqlKind.OTHER), - - // Expression-based filtering (on derived columns, e.g., HAVING) - EXPRESSION(Type.EXPRESSION, SqlKind.OTHER); - - /** - * Category of filter operator. Declares whether the operator supports parameters. - */ - public enum Type { - STANDARD(false), - FULL_TEXT(true), - EXPRESSION(false); - - private final boolean supportsParams; - - Type(boolean supportsParams) { - this.supportsParams = supportsParams; - } - - public boolean supportsParams() { - return supportsParams; - } - } - - private final Type type; private final SqlKind sqlKind; - FilterOperator(Type type, SqlKind sqlKind) { - this.type = type; + FilterOperator(SqlKind sqlKind) { this.sqlKind = sqlKind; } - public Type getType() { - return type; - } - - public SqlKind getSqlKind() { - return sqlKind; - } - - /** Maps a Calcite SqlKind to a standard FilterOperator, or null if not recognized. */ + /** Maps a Calcite SqlKind to a FilterOperator, or null if not a standard filter op. */ public static FilterOperator fromSqlKind(SqlKind kind) { for (FilterOperator op : values()) { - if (op.type == Type.STANDARD && op.sqlKind == kind && op.sqlKind != SqlKind.OTHER) { + if (op.sqlKind == kind && op.sqlKind != SqlKind.OTHER) { return op; } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java new file mode 100644 index 0000000000000..bb9068134523b --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java @@ -0,0 +1,32 @@ +/* + * 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; + +/** + * Full-text search operations that an inverted index backend may support. + * Intended as a common facade between Lucene, Tantivy, and future + * full-text backends. + * + * @opensearch.internal + */ +public enum FullTextOperator { + MATCH, + MATCH_PHRASE, + MATCH_PHRASE_PREFIX, + MATCH_BOOL_PREFIX, + MULTI_MATCH, + QUERY_STRING, + SIMPLE_QUERY_STRING, + FUZZY, + SPAN_NEAR, + SPAN_OR, + SPAN_NOT, + SPAN_FIRST, + SPAN_TERM +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java index 1875b2fbb49a3..081370f3b2875 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java @@ -17,7 +17,6 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; -import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.PlannerImpl; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; @@ -66,10 +65,7 @@ public DefaultPlanExecutor(List providers, Indices @Override public Iterable execute(RelNode logicalFragment, Object context) { logicalFragment = PlannerImpl.createPlan(logicalFragment, - new PlannerContext( - new CapabilityRegistry(new ArrayList<>(backEnds.values())), - clusterService.state())); - + new PlannerContext(new ArrayList<>(backEnds.values()), clusterService.state())); String tableName = extractTableName(logicalFragment); AnalyticsSearchBackendPlugin provider = selectBackEnd(); if (provider == null) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java index 5b45a75d2fad5..d968cf9ac4fd3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java @@ -8,16 +8,22 @@ package org.opensearch.analytics.planner; +import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.core.Aggregate; import org.opensearch.analytics.planner.rel.ShuffleImpl; +import org.opensearch.analytics.spi.AggregateFunction; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.analytics.spi.ShuffleCapability; import java.util.ArrayList; import java.util.List; +import java.util.Map; /** - * Utility logic that operates on {@link CapabilityRegistry} results. - * Registry handles indexed lookups; this class handles filtering and resolution logic. + * Shared helpers for resolving backend capabilities during planning. + * Used by all marking rules to determine which backend handles each operator. * * @opensearch.internal */ @@ -25,58 +31,148 @@ public final class CapabilityResolutionUtils { private CapabilityResolutionUtils() {} - /** Filters viable backends to those that support COORDINATOR_REDUCE. */ - public static List filterByReduceCapability(CapabilityRegistry registry, - List viableBackends) { - List reduceCapable = registry.operatorBackends(OperatorCapability.COORDINATOR_REDUCE); - List result = new ArrayList<>(); - for (String name : viableBackends) { - if (reduceCapable.contains(name)) { - result.add(name); + /** True if the named backend supports the given operator capability. */ + public static boolean backendSupports(Map backends, + String backendName, OperatorCapability capability) { + AnalyticsSearchBackendPlugin plugin = backends.get(backendName); + return plugin != null && plugin.supportedOperators().contains(capability); + } + + /** Finds the first backend that supports the given capability, or null. */ + public static String findBackendWith(Map backends, + OperatorCapability capability) { + for (AnalyticsSearchBackendPlugin plugin : backends.values()) { + if (plugin.supportedOperators().contains(capability)) { + return plugin.name(); } } - if (result.isEmpty()) { - throw new IllegalStateException( - "No viable backend supports COORDINATOR_REDUCE among " + viableBackends); + return null; + } + + /** + * Resolves the backend for an operator: prefers childBackend if it supports + * the capability, otherwise finds any backend that does. + * Falls back to childBackend if none found. + */ + public static String resolveBackend(Map backends, + String childBackend, OperatorCapability capability) { + if (backendSupports(backends, childBackend, capability)) { + return childBackend; + } + String found = findBackendWith(backends, capability); + return found != null ? found : childBackend; + } + + /** + * True if the named backend supports AGGREGATE and every aggregate function + * required by the plan. + */ + public static boolean backendSupportsAggregate(Map backends, + String backendName, Aggregate aggregate) { + AnalyticsSearchBackendPlugin plugin = backends.get(backendName); + if (plugin == null || !plugin.supportedOperators().contains(OperatorCapability.AGGREGATE)) { + return false; } - return result; + return aggregate.getAggCallList().stream().allMatch(aggCall -> { + AggregateFunction func = AggregateFunction.fromSqlKind(aggCall.getAggregation().getKind()); + if (func == null) { + func = AggregateFunction.fromNameOrError(aggCall.getAggregation().getName()); + } + return plugin.supportedAggregateFunctions().contains(func); + }); } - /** Filters viable backends to those with any shuffle capability. */ - public static List filterByShuffleCapability(CapabilityRegistry registry, - List viableBackends) { - List result = new ArrayList<>(); - for (String name : viableBackends) { - if (!registry.getShuffleCapabilities(name).isEmpty()) { - result.add(name); + /** Finds the first backend that supports AGGREGATE and all required functions. */ + public static String findBackendForAggregate(Map backends, + Aggregate aggregate) { + for (AnalyticsSearchBackendPlugin plugin : backends.values()) { + if (backendSupportsAggregate(backends, plugin.name(), aggregate)) { + return plugin.name(); } } - if (result.isEmpty()) { - throw new IllegalStateException( - "No viable backend supports shuffle among " + viableBackends); + return null; + } + + /** + * Computes all backends viable for an operator capability, considering delegation. + * A backend is viable if it supports the capability natively, OR if it can delegate + * that type of work and at least one other backend can accept the delegation. + */ + public static List computeViableBackends(Map backends, + OperatorCapability capability, + DelegationType delegationType) { + boolean anyAcceptsDelegation = backends.values().stream() + .anyMatch(b -> b.acceptedDelegations().contains(delegationType)); + + List viable = new ArrayList<>(); + for (AnalyticsSearchBackendPlugin plugin : backends.values()) { + if (plugin.supportedOperators().contains(capability)) { + viable.add(plugin.name()); + } else if (anyAcceptsDelegation && plugin.supportedDelegations().contains(delegationType)) { + viable.add(plugin.name()); + } } - return result; + return viable; } - /** Picks the best shuffle impl across all shuffle-viable backends. Prefers STREAM over FILE. */ - public static ShuffleImpl bestShuffleImpl(CapabilityRegistry registry, List shuffleViable) { - boolean hasStream = false; - boolean hasFile = false; - for (String name : shuffleViable) { - var caps = registry.getShuffleCapabilities(name); - if (caps.contains(ShuffleCapability.STREAM_WRITE)) { - hasStream = true; + /** Computes all backends that natively support the given capability. No delegation. */ + public static List computeViableBackends(Map backends, + OperatorCapability capability) { + List viable = new ArrayList<>(); + for (AnalyticsSearchBackendPlugin plugin : backends.values()) { + if (plugin.supportedOperators().contains(capability)) { + viable.add(plugin.name()); } - if (caps.contains(ShuffleCapability.FILE_WRITE)) { - hasFile = true; + } + return viable; + } + + /** + * Computes all backends viable for a capability that also support the given data format. + * Used by scan rule where the backend must match the index's primary data format. + */ + public static List computeViableBackends(Map backends, + OperatorCapability capability, String dataFormat) { + List viable = new ArrayList<>(); + for (AnalyticsSearchBackendPlugin plugin : backends.values()) { + if (plugin.supportedOperators().contains(capability) + && plugin.getSupportedFormats().stream().anyMatch(f -> f.name().equals(dataFormat))) { + viable.add(plugin.name()); } } - if (hasStream) { - return ShuffleImpl.STREAM; + return viable; + } + + /** + * Resolves the shuffle implementation for a HASH/RANGE exchange based on the + * backend's advertised shuffle capabilities. Prefers STREAM over FILE. + * Only called for non-SINGLETON distributions. + */ + public static ShuffleImpl resolveShuffleImpl(Map backends, + String backendName, RelDistribution.Type distributionType) { + AnalyticsSearchBackendPlugin plugin = backends.get(backendName); + if (plugin != null) { + if (plugin.supportedShuffleCapabilities().contains(ShuffleCapability.STREAM_WRITE)) { + return ShuffleImpl.STREAM; + } + if (plugin.supportedShuffleCapabilities().contains(ShuffleCapability.FILE_WRITE)) { + return ShuffleImpl.FILE; + } } - if (hasFile) { - return ShuffleImpl.FILE; + throw new IllegalStateException( + "Backend [" + backendName + "] does not advertise any shuffle capability for distribution [" + + distributionType + "]"); + } + + /** + * Validates that the backend supports coordinator-side reduce. + * Throws if it doesn't. + */ + public static void validateReduceCapability(Map backends, + String backendName) { + if (!backendSupports(backends, backendName, OperatorCapability.COORDINATOR_REDUCE)) { + throw new IllegalStateException( + "Backend [" + backendName + "] does not support COORDINATOR_REDUCE capability"); } - throw new IllegalStateException("No shuffle impl available among " + shuffleViable); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java index 1e8efa39d264a..fc9d6a16323a7 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java @@ -8,57 +8,45 @@ package org.opensearch.analytics.planner; -import org.apache.calcite.sql.type.SqlTypeName; -import org.opensearch.analytics.spi.FieldType; - import java.util.List; /** - * Per-column storage metadata describing where doc values, indices and stored fields live. + * Per-column storage metadata describing where doc values and indices live. * Flows through the plan tree: each OpenSearchRelNode computes this for its * output columns from its input's metadata. * + *

TODO: use {@code DataFormat} instead of String for format identifiers + * once the dependency is wired through. + * * @opensearch.internal */ public class FieldStorageInfo { private final String fieldName; - private final String mappingType; - private final FieldType fieldType; + private final String fieldType; private final List docValueFormats; private final List indexFormats; - private final List storedFieldFormats; private final boolean derived; - public FieldStorageInfo(String fieldName, String mappingType, FieldType fieldType, - List docValueFormats, List indexFormats, - List storedFieldFormats, boolean derived) { + public FieldStorageInfo(String fieldName, String fieldType, List docValueFormats, + List indexFormats, boolean derived) { this.fieldName = fieldName; - this.mappingType = mappingType; this.fieldType = fieldType; this.docValueFormats = docValueFormats; this.indexFormats = indexFormats; - this.storedFieldFormats = storedFieldFormats; this.derived = derived; } - /** Creates a derived column (agg result, expression) with no physical storage. - * FieldType inferred from SqlTypeName. */ - public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTypeName) { - return new FieldStorageInfo(fieldName, sqlTypeName.getName(), - FieldType.fromSqlTypeName(sqlTypeName), - List.of(), List.of(), List.of(), true); + /** Creates a derived column (agg result, expression) with no physical storage. */ + public static FieldStorageInfo derivedColumn(String fieldName, String fieldType) { + return new FieldStorageInfo(fieldName, fieldType, List.of(), List.of(), true); } public String getFieldName() { return fieldName; } - public String getMappingType() { - return mappingType; - } - - public FieldType getFieldType() { + public String getFieldType() { return fieldType; } @@ -72,11 +60,6 @@ public List getIndexFormats() { return indexFormats; } - /** Data formats holding stored fields for this field (e.g. ["parquet"], ["lucene"]). */ - public List getStoredFieldFormats() { - return storedFieldFormats; - } - /** True for computed columns (agg results, expressions) with no physical storage. */ public boolean isDerived() { return derived; @@ -89,8 +72,4 @@ public boolean hasDocValues() { public boolean hasIndex() { return !indexFormats.isEmpty(); } - - public boolean hasStoredFields() { - return !storedFieldFormats.isEmpty(); - } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java index bfef4645ce79b..40b60b6527bcb 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java @@ -8,41 +8,34 @@ package org.opensearch.analytics.planner; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; -import org.opensearch.analytics.spi.FieldType; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.metadata.MappingMetadata; -import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; -import org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability; import java.util.ArrayList; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; /** - * Resolves per-field storage metadata by consulting each backend's {@link DataFormat} - * capabilities. For each field, determines which formats provide doc values - * ({@link Capability#COLUMNAR_STORAGE}), indices ({@link Capability#FULL_TEXT_SEARCH}, - * {@link Capability#POINT_RANGE}), and stored fields ({@link Capability#STORED_FIELDS}). + * Builds {@link FieldStorageInfo} for index fields from {@link IndexMetadata}. * - *

Production constructor queries backends' {@link DataFormat#supportedFields()} to - * build per-field storage info. Test constructor accepts explicit per-field storage. + *

At the coordinator, we only have {@link MappingMetadata} (raw source map). + * {@code MappedFieldType} with {@code hasDocValues()} / {@code isSearchable()} is + * only available at the data node via {@code MapperService}. + * + *

TODO: introduce a coordinator-level typed field metadata API so we don't + * parse raw maps here. Or defer field storage resolution to the data node planner + * where {@code MappedFieldType} is available. * * @opensearch.internal */ public class FieldStorageResolver { - private final Map fieldStorage; - private final List docValueFormats; + private FieldStorageResolver() {} - /** - * Production: resolves per-field storage from IndexMetadata and backend capabilities. - */ @SuppressWarnings("unchecked") - public FieldStorageResolver(IndexMetadata indexMetadata, List backends) { + public static List resolve(IndexMetadata indexMetadata, List fieldNames) { String indexName = indexMetadata.getIndex().getName(); + String primaryFormat = indexMetadata.getSettings().get("index.composite.primary_data_format", "lucene"); MappingMetadata mapping = indexMetadata.mapping(); if (mapping == null) { @@ -54,123 +47,37 @@ public FieldStorageResolver(IndexMetadata indexMetadata, List> formatCapabilities = buildFormatCapabilities(backends); - - this.fieldStorage = new LinkedHashMap<>(); - for (Map.Entry entry : properties.entrySet()) { - String fieldName = entry.getKey(); - Map fieldProps = (Map) entry.getValue(); - String fieldType = (String) fieldProps.get("type"); - if (fieldType == null) { - throw new IllegalStateException("Field [" + fieldName + "] has no type in mapping"); - } - this.fieldStorage.put(fieldName, resolveField(fieldName, fieldType, fieldProps, formatCapabilities)); - } - this.docValueFormats = computeDocValueFormats(this.fieldStorage); - } - - /** - * Test/future: explicit per-field storage info. - * Simulates hybrid indices where doc values exist in multiple formats. - */ - public FieldStorageResolver(Map fieldStorage) { - this.fieldStorage = fieldStorage; - this.docValueFormats = computeDocValueFormats(fieldStorage); - } - - /** Resolves storage info for the requested fields. */ - public List resolve(List fieldNames) { - List result = new ArrayList<>(fieldNames.size()); + List result = new ArrayList<>(); for (String fieldName : fieldNames) { - FieldStorageInfo info = fieldStorage.get(fieldName); - if (info == null) { - throw new IllegalStateException("Field [" + fieldName + "] not found in field storage"); - } - result.add(info); - } - return result; - } - - /** Returns all unique data formats that hold doc values across all fields. Precomputed at creation. */ - public List docValueFormats() { - return docValueFormats; - } - - private static List computeDocValueFormats(Map fieldStorage) { - List formats = new ArrayList<>(); - for (FieldStorageInfo info : fieldStorage.values()) { - for (String format : info.getDocValueFormats()) { - if (!formats.contains(format)) { - formats.add(format); - } - } - } - return formats; - } - - /** - * Builds a lookup: formatName → fieldType → FieldTypeCapabilities - * from all backends' DataFormats. - */ - private static Map> buildFormatCapabilities( - List backends) { - Map> result = new LinkedHashMap<>(); - for (AnalyticsSearchBackendPlugin backend : backends) { - for (DataFormat format : backend.getSupportedFormats()) { - Map byFieldType = result.computeIfAbsent( - format.name(), k -> new LinkedHashMap<>()); - for (FieldTypeCapabilities cap : format.supportedFields()) { - byFieldType.put(cap.fieldType(), cap); - } + Map fieldProps = (Map) properties.get(fieldName); + if (fieldProps == null) { + throw new IllegalStateException("Field [" + fieldName + "] not found in mapping for index [" + indexName + "]"); } + result.add(resolveField(fieldName, fieldProps, primaryFormat)); } return result; } - private static FieldStorageInfo resolveField(String fieldName, String fieldType, - Map fieldProps, - Map> formatCapabilities) { - List docValueFormats = new ArrayList<>(); - List indexFormats = new ArrayList<>(); - List storedFieldFormats = new ArrayList<>(); - - for (Map.Entry> formatEntry : formatCapabilities.entrySet()) { - String formatName = formatEntry.getKey(); - FieldTypeCapabilities caps = formatEntry.getValue().get(fieldType); - if (caps == null) { - continue; - } - if (caps.capabilities().contains(Capability.COLUMNAR_STORAGE)) { - docValueFormats.add(formatName); - } - if (caps.capabilities().contains(Capability.FULL_TEXT_SEARCH) - || caps.capabilities().contains(Capability.POINT_RANGE)) { - indexFormats.add(formatName); - } - if (caps.capabilities().contains(Capability.STORED_FIELDS)) { - storedFieldFormats.add(formatName); - } + private static FieldStorageInfo resolveField(String fieldName, Map fieldProps, String primaryFormat) { + String fieldType = (String) fieldProps.get("type"); + if (fieldType == null) { + throw new IllegalStateException("Field [" + fieldName + "] has no type in mapping"); } - // Respect mapping overrides: doc_values=false or index=false + // TODO: use MappedFieldType.hasDocValues() / isSearchable() when available boolean hasDocValues = Boolean.TRUE.equals(fieldProps.get("doc_values")) || (fieldProps.get("doc_values") == null && !"text".equals(fieldType)); boolean isIndexed = Boolean.TRUE.equals(fieldProps.get("index")) || fieldProps.get("index") == null; - if (!hasDocValues) { - docValueFormats = List.of(); - } - if (!isIndexed) { - indexFormats = List.of(); + if (!hasDocValues && !isIndexed) { + throw new IllegalStateException("Field [" + fieldName + "] has neither doc_values nor index"); } - if (docValueFormats.isEmpty() && indexFormats.isEmpty() && storedFieldFormats.isEmpty()) { - throw new IllegalStateException("Field [" + fieldName + "] has no storage in any format"); - } + // TODO: data format per field should come from MappingMetadata directly + List docValueFormats = hasDocValues ? List.of(primaryFormat) : List.of(); + List indexFormats = isIndexed ? List.of("lucene") : List.of(); - return new FieldStorageInfo(fieldName, fieldType, FieldType.fromMappingType(fieldType), - docValueFormats, indexFormats, storedFieldFormats, false); + return new FieldStorageInfo(fieldName, fieldType, docValueFormats, indexFormats, false); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java index ae69b6fb7bdd4..620e604de7670 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java @@ -9,29 +9,38 @@ package org.opensearch.analytics.planner; import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.cluster.ClusterState; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + /** * Shared context available to all planner rules. - * Holds capability registry (singleton, built at plugin startup) and - * per-query cluster state. + * Holds cluster state and backend plugins. Rules consult this + * to extract index metadata, mappings, and backend capabilities + * for whichever tables they encounter. * * @opensearch.internal */ public class PlannerContext { - private final CapabilityRegistry capabilityRegistry; + private final Map backends; private final ClusterState clusterState; private final OpenSearchDistributionTraitDef distributionTraitDef; - public PlannerContext(CapabilityRegistry capabilityRegistry, ClusterState clusterState) { - this.capabilityRegistry = capabilityRegistry; + public PlannerContext(List backends, ClusterState clusterState) { + this.backends = new LinkedHashMap<>(); + for (AnalyticsSearchBackendPlugin b : backends) { + this.backends.put(b.name(), b); + } this.clusterState = clusterState; this.distributionTraitDef = new OpenSearchDistributionTraitDef(this); } - public CapabilityRegistry getCapabilityRegistry() { - return capabilityRegistry; + public Map getBackends() { + return backends; } public ClusterState getClusterState() { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java index 6418f2c163fd6..098e9dcea8402 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java @@ -57,20 +57,20 @@ public static RelNode copyToCluster(RelNode node, RelOptCluster newCluster, if (node instanceof OpenSearchTableScan scan) { return new OpenSearchTableScan(newCluster, newTraits, scan.getTable(), - scan.getViableBackends(), scan.getOutputFieldStorage()); + scan.getBackend(), scan.getViableBackends(), scan.getOutputFieldStorage()); } else if (node instanceof OpenSearchFilter filter) { return new OpenSearchFilter(newCluster, newTraits, newInputs.getFirst(), - filter.getCondition(), filter.getViableBackends()); + filter.getCondition(), filter.getBackend(), filter.getViableBackends()); } else if (node instanceof OpenSearchAggregate aggregate) { return new OpenSearchAggregate(newCluster, newTraits, newInputs.getFirst(), aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), - aggregate.getMode(), aggregate.getViableBackends()); + aggregate.getMode(), aggregate.getBackend(), aggregate.getViableBackends()); } else if (node instanceof OpenSearchSort sort) { return new OpenSearchSort(newCluster, newTraits, newInputs.getFirst(), - sort.getCollation(), sort.offset, sort.fetch, sort.getViableBackends()); + sort.getCollation(), sort.offset, sort.fetch, sort.getBackend(), sort.getViableBackends()); } else if (node instanceof OpenSearchProject project) { return new OpenSearchProject(newCluster, newTraits, newInputs.getFirst(), - project.getProjects(), project.getRowType(), project.getViableBackends()); + project.getProjects(), project.getRowType(), project.getBackend()); } throw new UnsupportedOperationException("Cannot copy node type: " + node.getClass().getSimpleName()); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java index 03e9851c7a0b5..12b96683264cb 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java @@ -46,25 +46,25 @@ public SqlSyntax getSyntax() { }; private final RexNode original; - private final List viableBackends; + private final String backend; - public AnnotatedProjectExpression(RelDataType type, RexNode original, List viableBackends) { + public AnnotatedProjectExpression(RelDataType type, RexNode original, String backend) { super(type, ANNOTATED_PROJECT_EXPR_OP, List.of(original)); this.original = original; - this.viableBackends = viableBackends; + this.backend = backend; } public RexNode getOriginal() { return original; } - /** Backends that can evaluate this expression. */ - public List getViableBackends() { - return viableBackends; + /** The backend that evaluates this expression. */ + public String getBackend() { + return backend; } @Override protected String computeDigest(boolean withType) { - return "ANNOTATED_PROJECT_EXPR(backends=" + viableBackends + ", " + original + ")"; + return "ANNOTATED_PROJECT_EXPR(backend=" + backend + ", " + original + ")"; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java index 7b2f21a6ac4b4..713830e10878b 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java @@ -13,12 +13,14 @@ import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.type.OperandTypes; import org.apache.calcite.sql.type.ReturnTypes; -import org.opensearch.analytics.spi.FilterOperator; +import org.opensearch.analytics.spi.FullTextOperator; /** * Calcite SqlFunction markers for full-text search operations. - * The filter rule recognizes these and routes to backends that support - * the corresponding {@link FilterOperator} of type {@link FilterOperator.Type#FULL_TEXT}. + * Used by frontend plugins (DSL converter, SQL, PPL) to represent + * full-text queries in the RelNode tree. The filter rule recognizes + * these and routes to backends that support the corresponding + * {@link FullTextOperator}. * * @opensearch.internal */ @@ -36,15 +38,21 @@ public class FullTextFunctions { private FullTextFunctions() {} private static SqlFunction fullTextFunction(String name) { - return new SqlFunction(name, SqlKind.OTHER_FUNCTION, ReturnTypes.BOOLEAN, - null, OperandTypes.ANY, SqlFunctionCategory.USER_DEFINED_FUNCTION); + return new SqlFunction( + name, + SqlKind.OTHER_FUNCTION, + ReturnTypes.BOOLEAN, + null, + OperandTypes.ANY, + SqlFunctionCategory.USER_DEFINED_FUNCTION + ); } - /** Maps a SqlFunction to a FULL_TEXT FilterOperator, or null if not a full-text function. */ - public static FilterOperator toFilterOperator(SqlFunction function) { + /** Maps a SqlFunction to a FullTextOperator, or null if not a full-text function. */ + public static FullTextOperator toFullTextOperator(SqlFunction function) { + String name = function.getName(); try { - FilterOperator op = FilterOperator.valueOf(function.getName()); - return op.getType() == FilterOperator.Type.FULL_TEXT ? op : null; + return FullTextOperator.valueOf(name); } catch (IllegalArgumentException ignored) { return null; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java index df0eb21a54852..78d6eade67fad 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java @@ -26,18 +26,25 @@ */ public class OpenSearchAggregate extends Aggregate implements OpenSearchRelNode { + private final String backend; private final List viableBackends; private final AggregateMode mode; public OpenSearchAggregate(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, ImmutableBitSet groupSet, List groupSets, List aggCalls, AggregateMode mode, - List viableBackends) { + String backend, List viableBackends) { super(cluster, traitSet, List.of(), input, groupSet, groupSets, aggCalls); this.mode = mode; + this.backend = backend; this.viableBackends = viableBackends; } + @Override + public String getBackend() { + return backend; + } + public AggregateMode getMode() { return mode; } @@ -69,8 +76,7 @@ public List getOutputFieldStorage() { // Agg results: derived columns with no physical storage for (AggregateCall aggCall : getAggCallList()) { - outputStorage.add(FieldStorageInfo.derivedColumn(aggCall.getName(), - aggCall.getType().getSqlTypeName())); + outputStorage.add(FieldStorageInfo.derivedColumn(aggCall.getName(), aggCall.getType().toString())); } return outputStorage; @@ -79,7 +85,7 @@ public List getOutputFieldStorage() { @Override public Aggregate copy(RelTraitSet traitSet, RelNode input, ImmutableBitSet groupSet, List groupSets, List aggCalls) { - return new OpenSearchAggregate(getCluster(), traitSet, input, groupSet, groupSets, aggCalls, mode, viableBackends); + return new OpenSearchAggregate(getCluster(), traitSet, input, groupSet, groupSets, aggCalls, mode, backend, viableBackends); } @Override @@ -103,6 +109,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(org.apache.calcite.pla @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("mode", mode).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("mode", mode).item("backend", backend).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java index b381f6c147be3..b36348f0bc274 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java @@ -15,7 +15,6 @@ import org.apache.calcite.rel.RelNode; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.PlannerContext; @@ -94,32 +93,31 @@ public RelNode convert(RelOptPlanner planner, RelNode rel, return rel; } - List viableBackends = resolveViableBackendsFromRel(rel); + String backend = resolveBackendFromRel(rel); LOGGER.info("convert(): rel={}#{}, fromTrait={}, toTrait={}, backend={}", - rel.getClass().getSimpleName(), rel.getId(), fromTrait, toTrait, viableBackends.getFirst()); - - CapabilityRegistry registry = plannerContext.getCapabilityRegistry(); + rel.getClass().getSimpleName(), rel.getId(), fromTrait, toTrait, backend); RelNode result; if (toTrait.getType() == RelDistribution.Type.SINGLETON) { - List reduceViable = CapabilityResolutionUtils.filterByReduceCapability( - registry, viableBackends); + CapabilityResolutionUtils.validateReduceCapability( + plannerContext.getBackends(), backend); result = new OpenSearchExchangeReducer( rel.getCluster(), rel.getTraitSet().replace(toTrait), rel, - reduceViable + backend ); } else { - List shuffleViable = CapabilityResolutionUtils.filterByShuffleCapability( - registry, viableBackends); - ShuffleImpl shuffleImpl = CapabilityResolutionUtils.bestShuffleImpl(registry, shuffleViable); + // HASH/RANGE: Writer at data node partitions and writes shuffle data. + // Reader at target data node reads from source nodes. + ShuffleImpl shuffleImpl = CapabilityResolutionUtils.resolveShuffleImpl( + plannerContext.getBackends(), backend, toTrait.getType()); OpenSearchExchangeWriter writer = new OpenSearchExchangeWriter( rel.getCluster(), rel.getTraitSet(), rel, - shuffleViable, + backend, shuffleImpl, toTrait.getKeys() ); @@ -127,7 +125,7 @@ public RelNode convert(RelOptPlanner planner, RelNode rel, rel.getCluster(), rel.getTraitSet().replace(toTrait), writer, - shuffleViable, + backend, shuffleImpl ); } @@ -141,12 +139,12 @@ public boolean canConvert(RelOptPlanner planner, OpenSearchDistribution fromTrai return true; } - private static List resolveViableBackendsFromRel(RelNode rel) { + private static String resolveBackendFromRel(RelNode rel) { if (rel instanceof RelSubset subset) { rel = subset.getBestOrOriginal(); } if (rel instanceof OpenSearchRelNode openSearchRel) { - return openSearchRel.getViableBackends(); + return openSearchRel.getBackend(); } throw new IllegalStateException( "Expected OpenSearchRelNode but got [" + rel.getClass().getSimpleName() + "#" + rel.getId() + "]"); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java index 00111a8b5ff5c..19f3dedc3a8f5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java @@ -32,17 +32,25 @@ */ public class OpenSearchExchangeReducer extends SingleRel implements OpenSearchRelNode { - private final List viableBackends; + private final String backend; public OpenSearchExchangeReducer(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - List viableBackends) { + String backend) { super(cluster, traitSet, input); - this.viableBackends = viableBackends; + this.backend = backend; + } + + @Override + public String getBackend() { + return backend; } @Override public List getViableBackends() { - return viableBackends; + if (getInput() instanceof OpenSearchRelNode openSearchInput) { + return openSearchInput.getViableBackends(); + } + return List.of(); } @Override @@ -55,7 +63,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), viableBackends); + return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), backend); } @Override @@ -65,6 +73,6 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("backend", backend).item("viableBackends", getViableBackends()); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java index 0f8f6fd7f8ab3..99489200ac1de 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java @@ -32,23 +32,24 @@ */ public class OpenSearchExchangeWriter extends SingleRel implements OpenSearchRelNode { - private final List viableBackends; - private final ShuffleImpl shuffleImpl; - private final List keys; + private final String backend; + private final ShuffleImpl shuffleImpl; // null for SINGLETON (transport, not shuffle) + private final List keys; // empty for SINGLETON public OpenSearchExchangeWriter(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - List viableBackends, ShuffleImpl shuffleImpl, List keys) { + String backend, ShuffleImpl shuffleImpl, List keys) { super(cluster, traitSet, input); - this.viableBackends = viableBackends; + this.backend = backend; this.shuffleImpl = shuffleImpl; this.keys = keys; } @Override - public List getViableBackends() { - return viableBackends; + public String getBackend() { + return backend; } + /** Null for SINGLETON (uses Analytics Core transport). FILE or STREAM for shuffle. */ public ShuffleImpl getShuffleImpl() { return shuffleImpl; } @@ -61,6 +62,14 @@ public boolean isShuffle() { return shuffleImpl != null; } + @Override + public List getViableBackends() { + if (getInput() instanceof OpenSearchRelNode openSearchInput) { + return openSearchInput.getViableBackends(); + } + return List.of(); + } + @Override public List getOutputFieldStorage() { if (getInput() instanceof OpenSearchRelNode openSearchInput) { @@ -72,7 +81,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { return new OpenSearchExchangeWriter(getCluster(), traitSet, sole(inputs), - viableBackends, shuffleImpl, keys); + backend, shuffleImpl, keys); } @Override @@ -82,7 +91,7 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - RelWriter writer = super.explainTerms(pw).item("viableBackends", viableBackends); + RelWriter writer = super.explainTerms(pw).item("backend", backend); if (shuffleImpl != null) { writer.item("shuffleImpl", shuffleImpl).item("keys", keys); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java index 936c443c136e1..45fe5cabc1974 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java @@ -23,14 +23,21 @@ */ public class OpenSearchFilter extends Filter implements OpenSearchRelNode { + private final String backend; private final List viableBackends; public OpenSearchFilter(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - RexNode condition, List viableBackends) { + RexNode condition, String backend, List viableBackends) { super(cluster, traitSet, input, condition); + this.backend = backend; this.viableBackends = viableBackends; } + @Override + public String getBackend() { + return backend; + } + @Override public List getViableBackends() { return viableBackends; @@ -47,7 +54,7 @@ public List getOutputFieldStorage() { @Override public Filter copy(RelTraitSet traitSet, RelNode input, RexNode condition) { - return new OpenSearchFilter(getCluster(), traitSet, input, condition, viableBackends); + return new OpenSearchFilter(getCluster(), traitSet, input, condition, backend, viableBackends); } @Override @@ -58,6 +65,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(org.apache.calcite.pla @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("backend", backend).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java index c31b7d72942ed..413d50de95906 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java @@ -17,13 +17,8 @@ import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.metadata.RelMetadataQuery; import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexNode; import org.opensearch.analytics.planner.FieldStorageInfo; -import org.opensearch.analytics.planner.RelNodeUtils; - -import java.util.ArrayList; -import java.util.List; import java.util.List; @@ -32,46 +27,35 @@ */ public class OpenSearchProject extends Project implements OpenSearchRelNode { - private final List viableBackends; + private final String backend; public OpenSearchProject(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, List projects, RelDataType rowType, - List viableBackends) { + String backend) { super(cluster, traitSet, List.of(), input, projects, rowType); - this.viableBackends = viableBackends; + this.backend = backend; + } + + @Override + public String getBackend() { + return backend; } @Override public List getViableBackends() { - return viableBackends; + return List.of(backend); } @Override public List getOutputFieldStorage() { - RelNode input = RelNodeUtils.unwrapHep(getInput()); - if (!(input instanceof OpenSearchRelNode openSearchChild)) { - throw new IllegalStateException( - "Project child is not OpenSearchRelNode: " + input.getClass().getSimpleName()); - } - List inputStorage = openSearchChild.getOutputFieldStorage(); - - List result = new ArrayList<>(getProjects().size()); - for (int i = 0; i < getProjects().size(); i++) { - RexNode expr = getProjects().get(i); - if (expr instanceof RexInputRef ref && ref.getIndex() < inputStorage.size()) { - result.add(inputStorage.get(ref.getIndex())); - } else { - String fieldName = getRowType().getFieldList().get(i).getName(); - result.add(FieldStorageInfo.derivedColumn(fieldName, - getRowType().getFieldList().get(i).getType().getSqlTypeName())); - } - } - return result; + return getRowType().getFieldList().stream() + .map(field -> FieldStorageInfo.derivedColumn(field.getName(), field.getType().toString())) + .toList(); } @Override public Project copy(RelTraitSet traitSet, RelNode input, List projects, RelDataType rowType) { - return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, viableBackends); + return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, backend); } @Override @@ -81,6 +65,6 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("backend", backend); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java index bacb2b62926f3..36d13b93db68a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java @@ -19,14 +19,20 @@ *

Each node computes {@link #getOutputFieldStorage()} from its input's metadata. * Parent operators read this to make backend routing decisions. * - *

{@link #getViableBackends()} lists all backends that could execute this operator - * (including via delegation). The first entry is the default/preferred backend. - * Consumed during plan forking to generate one complete plan per viable backend. + *

{@link #getBackend()} is the default/preferred backend chosen by the marking rule. + * {@link #getViableBackends()} lists all backends that could execute this operator + * (including via delegation). Consumed by {@code BackendResolver.generateCandidatePlans()} + * during StagePlan alternative generation. * * @opensearch.internal */ public interface OpenSearchRelNode { + String UNRESOLVED = "unresolved"; + + /** Default/preferred backend chosen by the marking rule. */ + String getBackend(); + /** All backends that could execute this operator, including via delegation. */ List getViableBackends(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java index 2359c13a141ad..d41944ccaa531 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java @@ -32,23 +32,31 @@ */ public class OpenSearchShuffleReader extends SingleRel implements OpenSearchRelNode { - private final List viableBackends; + private final String backend; private final ShuffleImpl shuffleImpl; public OpenSearchShuffleReader(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - List viableBackends, ShuffleImpl shuffleImpl) { + String backend, ShuffleImpl shuffleImpl) { super(cluster, traitSet, input); - this.viableBackends = viableBackends; + this.backend = backend; this.shuffleImpl = shuffleImpl; } + @Override + public String getBackend() { + return backend; + } + public ShuffleImpl getShuffleImpl() { return shuffleImpl; } @Override public List getViableBackends() { - return viableBackends; + if (getInput() instanceof OpenSearchRelNode openSearchInput) { + return openSearchInput.getViableBackends(); + } + return List.of(); } @Override @@ -61,7 +69,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchShuffleReader(getCluster(), traitSet, sole(inputs), viableBackends, shuffleImpl); + return new OpenSearchShuffleReader(getCluster(), traitSet, sole(inputs), backend, shuffleImpl); } @Override @@ -71,6 +79,6 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("viableBackends", viableBackends).item("shuffleImpl", shuffleImpl); + return super.explainTerms(pw).item("backend", backend).item("shuffleImpl", shuffleImpl); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java index 3098d37cd971e..58dafcd9f7d23 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java @@ -24,15 +24,22 @@ */ public class OpenSearchSort extends Sort implements OpenSearchRelNode { + private final String backend; private final List viableBackends; public OpenSearchSort(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, RelCollation collation, RexNode offset, RexNode fetch, - List viableBackends) { + String backend, List viableBackends) { super(cluster, traitSet, input, collation, offset, fetch); + this.backend = backend; this.viableBackends = viableBackends; } + @Override + public String getBackend() { + return backend; + } + @Override public List getViableBackends() { return viableBackends; @@ -49,7 +56,7 @@ public List getOutputFieldStorage() { @Override public Sort copy(RelTraitSet traitSet, RelNode input, RelCollation collation, RexNode offset, RexNode fetch) { - return new OpenSearchSort(getCluster(), traitSet, input, collation, offset, fetch, viableBackends); + return new OpenSearchSort(getCluster(), traitSet, input, collation, offset, fetch, backend, viableBackends); } @Override @@ -60,6 +67,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(org.apache.calcite.pla @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("backend", backend).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java index dc22910d9b3fa..93b6cf4b37947 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java @@ -25,13 +25,15 @@ */ public class OpenSearchTableScan extends TableScan implements OpenSearchRelNode { + private final String backend; private final List viableBackends; private final List outputFieldStorage; public OpenSearchTableScan(RelOptCluster cluster, RelTraitSet traitSet, RelOptTable table, - List viableBackends, + String backend, List viableBackends, List outputFieldStorage) { super(cluster, traitSet, List.of(), table); + this.backend = backend; this.viableBackends = viableBackends; this.outputFieldStorage = outputFieldStorage; } @@ -42,7 +44,7 @@ public OpenSearchTableScan(RelOptCluster cluster, RelTraitSet traitSet, RelOptTa * Single shard → SINGLETON (all data on one node). */ public static OpenSearchTableScan create(RelOptCluster cluster, RelOptTable table, - List viableBackends, + String backend, List viableBackends, List outputFieldStorage, int shardCount, OpenSearchDistributionTraitDef distTraitDef) { @@ -52,7 +54,12 @@ public static OpenSearchTableScan create(RelOptCluster cluster, RelOptTable tabl RelTraitSet traitSet = RelTraitSet.createEmpty() .plus(OpenSearchConvention.INSTANCE) .plus(distribution); - return new OpenSearchTableScan(cluster, traitSet, table, viableBackends, outputFieldStorage); + return new OpenSearchTableScan(cluster, traitSet, table, backend, viableBackends, outputFieldStorage); + } + + @Override + public String getBackend() { + return backend; } @Override @@ -67,7 +74,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchTableScan(getCluster(), traitSet, getTable(), viableBackends, outputFieldStorage); + return new OpenSearchTableScan(getCluster(), traitSet, getTable(), backend, viableBackends, outputFieldStorage); } @Override @@ -78,6 +85,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(RelOptPlanner planner, @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("backend", backend).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java index abe1b52acf7bf..7db6276cf0164 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java @@ -16,8 +16,6 @@ import org.apache.calcite.rel.core.AggregateCall; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.analytics.planner.CapabilityRegistry; -import org.opensearch.analytics.planner.FieldStorageInfo; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.AggregateCallAnnotation; @@ -25,8 +23,8 @@ import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.spi.AggregateFunction; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; -import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.OperatorCapability; import java.util.ArrayList; @@ -72,13 +70,12 @@ public void onMatch(RelOptRuleCall call) { "Aggregate rule encountered unmarked child [" + child.getClass().getSimpleName() + "]"); } - List childViableBackends = openSearchChild.getViableBackends(); - List childFieldStorage = openSearchChild.getOutputFieldStorage(); + String childBackend = openSearchChild.getBackend(); // Annotate each AggregateCall with per-call viable backends List annotatedCalls = new ArrayList<>(); for (AggregateCall aggCall : aggregate.getAggCallList()) { - List callViable = resolveViableBackendsForCall(aggCall, childFieldStorage); + List callViable = resolveViableBackendsForCall(aggCall); if (callViable.isEmpty()) { throw new IllegalStateException( "No backend supports aggregate function [" + aggCall.getAggregation().getName() + "]"); @@ -86,8 +83,8 @@ public void onMatch(RelOptRuleCall call) { annotatedCalls.add(AggregateCallAnnotation.annotate(aggCall, callViable)); } - // Compute operator-level viable backends: must be viable for child AND handle agg calls - List viableBackends = computeAggregateViableBackends(annotatedCalls, childViableBackends); + // Compute operator-level viable backends considering delegation + List viableBackends = computeAggregateViableBackends(annotatedCalls, childBackend); if (viableBackends.isEmpty()) { List funcNames = aggregate.getAggCallList().stream() @@ -95,10 +92,12 @@ public void onMatch(RelOptRuleCall call) { .toList(); throw new IllegalStateException( "No backend can execute aggregate: functions " + funcNames - + " not supported by any viable backend among " + childViableBackends); + + " are split across backends and no delegation path exists"); } - logger.debug("Aggregate viable backends: {} (child viable: {})", viableBackends, childViableBackends); + String backend = viableBackends.contains(childBackend) + ? childBackend + : viableBackends.getFirst(); RelTraitSet aggregateTraits = child.getTraitSet(); if (aggregateTraits.size() > 0) { @@ -113,78 +112,41 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), annotatedCalls, AggregateMode.SINGLE, + backend, viableBackends )); } - private List resolveViableBackendsForCall(AggregateCall aggCall, - List childFieldStorage) { + private List resolveViableBackendsForCall(AggregateCall aggCall) { AggregateFunction func = AggregateFunction.fromSqlKind(aggCall.getAggregation().getKind()); if (func == null) { func = AggregateFunction.fromNameOrError(aggCall.getAggregation().getName()); } - CapabilityRegistry registry = context.getCapabilityRegistry(); - - if (aggCall.getArgList().isEmpty()) { - return new ArrayList<>(registry.operatorBackends(OperatorCapability.AGGREGATE)); - } - - List callViable = null; - for (int fieldIndex : aggCall.getArgList()) { - if (fieldIndex >= childFieldStorage.size()) { - continue; - } - FieldStorageInfo storageInfo = childFieldStorage.get(fieldIndex); - FieldType fieldType = storageInfo.getFieldType(); - if (fieldType == null) { - throw new IllegalStateException("Unrecognized field type [" + storageInfo.getMappingType() - + "] for field [" + storageInfo.getFieldName() + "]"); - } - - List perFieldBackends = new ArrayList<>(); - if (storageInfo.isDerived()) { - perFieldBackends.addAll(registry.aggregateBackendsAnyFormat(func, fieldType)); - } else { - // Format-aware: backends that can read the data and compute - for (String format : storageInfo.getDocValueFormats()) { - for (String name : registry.aggregateBackends(func, fieldType, format)) { - if (!perFieldBackends.contains(name)) { - perFieldBackends.add(name); - } - } - } - // Format-agnostic: delegation targets that can compute but don't need data access - for (String name : registry.aggregateBackendsAnyFormat(func, fieldType)) { - if (!perFieldBackends.contains(name)) { - perFieldBackends.add(name); - } - } - } - - if (callViable == null) { - callViable = perFieldBackends; - } else { - callViable.retainAll(perFieldBackends); + List viable = new ArrayList<>(); + for (AnalyticsSearchBackendPlugin plugin : context.getBackends().values()) { + if (plugin.supportedOperators().contains(OperatorCapability.AGGREGATE) + && plugin.supportedAggregateFunctions().contains(func)) { + viable.add(plugin.name()); } } - - return callViable != null ? callViable : new ArrayList<>(registry.operatorBackends(OperatorCapability.AGGREGATE)); + return viable; } + /** + * Computes which backends can execute this aggregate, considering delegation. + * A backend is viable if for every agg call it can handle natively OR delegate + * (supports AGGREGATE delegation AND some other backend accepts it for that call). + */ private List computeAggregateViableBackends(List annotatedCalls, - List childViableBackends) { + String childBackend) { if (annotatedCalls.isEmpty()) { - return new ArrayList<>(childViableBackends); + return new ArrayList<>(context.getBackends().keySet()); } - CapabilityRegistry registry = context.getCapabilityRegistry(); - List delegationSupporters = registry.delegationSupporters(DelegationType.AGGREGATE); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.AGGREGATE); - List viable = new ArrayList<>(); - for (String candidateName : childViableBackends) { - if (!registry.operatorBackends(OperatorCapability.AGGREGATE).contains(candidateName)) { + for (AnalyticsSearchBackendPlugin candidate : context.getBackends().values()) { + if (!candidate.supportedOperators().contains(OperatorCapability.AGGREGATE)) { continue; } @@ -196,18 +158,24 @@ private List computeAggregateViableBackends(List annotate break; } List callViable = annotation.getViableBackends(); - if (callViable.contains(candidateName)) { + if (callViable.contains(candidate.name())) { continue; } - if (delegationSupporters.contains(candidateName) - && callViable.stream().anyMatch(delegationAcceptors::contains)) { - continue; + // Check if candidate can delegate this call + if (candidate.supportedDelegations().contains(DelegationType.AGGREGATE)) { + boolean someoneAccepts = callViable.stream().anyMatch(backendName -> { + AnalyticsSearchBackendPlugin other = context.getBackends().get(backendName); + return other != null && other.acceptedDelegations().contains(DelegationType.AGGREGATE); + }); + if (someoneAccepts) { + continue; + } } canHandleAll = false; break; } if (canHandleAll) { - viable.add(candidateName); + viable.add(candidate.name()); } } return viable; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java index c010728ad5223..d0001238e402b 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java @@ -50,6 +50,8 @@ public void onMatch(RelOptRuleCall call) { OpenSearchAggregate aggregate = call.rel(0); RelNode child = call.rel(1); + String backend = aggregate.getBackend(); + // Partial aggregate: runs on each partition, keeps input's traits RelTraitSet partialTraits = child.getTraitSet() .replace(OpenSearchConvention.INSTANCE); @@ -61,6 +63,7 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), aggregate.getAggCallList(), AggregateMode.PARTIAL, + backend, aggregate.getViableBackends() ); @@ -78,6 +81,7 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), aggregate.getAggCallList(), AggregateMode.FINAL, + backend, aggregate.getViableBackends() ); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java index 21fa5a9d2abd3..773fb0341e051 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java @@ -22,23 +22,24 @@ import org.apache.logging.log4j.Logger; import org.opensearch.analytics.planner.FieldStorageInfo; import org.opensearch.analytics.planner.PlannerContext; -import org.opensearch.analytics.planner.CapabilityRegistry; -import org.opensearch.analytics.planner.FieldStorageInfo; -import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.AnnotatedPredicate; import org.opensearch.analytics.planner.rel.FullTextFunctions; import org.opensearch.analytics.planner.rel.OpenSearchFilter; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; -import org.opensearch.analytics.spi.FieldType; +import org.opensearch.analytics.spi.FieldTypeFamily; +import org.opensearch.analytics.spi.FilterCapability; import org.opensearch.analytics.spi.FilterOperator; +import org.opensearch.analytics.spi.FullTextOperator; import org.opensearch.analytics.spi.OperatorCapability; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.stream.Collectors; /** * Converts {@link Filter} → {@link OpenSearchFilter}. @@ -80,31 +81,34 @@ public void onMatch(RelOptRuleCall call) { + "]. Ensure all child operators are marked before filter."); } - List childViableBackends = openSearchInput.getViableBackends(); + String childBackend = openSearchInput.getBackend(); List childFieldStorage = openSearchInput.getOutputFieldStorage(); RelDataType inputRowType = child.getRowType(); // Annotate every leaf predicate with viable backends - RexNode annotatedCondition = annotateCondition(filter.getCondition(), inputRowType, - childFieldStorage, childViableBackends); + RexNode annotatedCondition = annotateCondition(filter.getCondition(), inputRowType, childFieldStorage); - // Compute operator-level viable backends: must be viable for child AND handle predicates - List viableBackends = computeFilterViableBackends(annotatedCondition, childViableBackends); + // Compute operator-level viable backends from per-predicate annotations + delegation + List viableBackends = computeFilterViableBackends(annotatedCondition, childBackend); if (viableBackends.isEmpty()) { throw new IllegalStateException( - "No backend can execute filter: no viable backend among " + childViableBackends - + " can evaluate all predicates and no delegation path exists"); + "No backend can execute filter: child backend [" + childBackend + + "] cannot evaluate all predicates and no delegation path exists"); } - LOGGER.debug("Filter viable backends: {} (child viable: {})", viableBackends, childViableBackends); + // Preferred backend: child if viable, otherwise first viable + String filterBackend = viableBackends.contains(childBackend) + ? childBackend + : viableBackends.getFirst(); call.transformTo(new OpenSearchFilter( filter.getCluster(), child.getTraitSet(), RelNodeUtils.unwrapHep(filter.getInput()), annotatedCondition, + filterBackend, viableBackends )); } @@ -117,19 +121,18 @@ public void onMatch(RelOptRuleCall call) { * {@link AnnotatedPredicate} with viable backends resolved from child's field storage. */ private RexNode annotateCondition(RexNode condition, RelDataType inputRowType, - List fieldStorage, - List childViableBackends) { + List fieldStorage) { if (!(condition instanceof RexCall rexCall)) { return condition; } if (rexCall.getKind() == SqlKind.AND || rexCall.getKind() == SqlKind.OR || rexCall.getKind() == SqlKind.NOT) { List annotatedOperands = new ArrayList<>(); for (RexNode operand : rexCall.getOperands()) { - annotatedOperands.add(annotateCondition(operand, inputRowType, fieldStorage, childViableBackends)); + annotatedOperands.add(annotateCondition(operand, inputRowType, fieldStorage)); } return rexCall.clone(rexCall.getType(), annotatedOperands); } - List viableBackends = resolveViableBackends(rexCall, inputRowType, fieldStorage, childViableBackends); + List viableBackends = resolveViableBackends(rexCall, inputRowType, fieldStorage); return new AnnotatedPredicate(rexCall.getType(), rexCall, viableBackends); } @@ -140,78 +143,82 @@ private RexNode annotateCondition(RexNode condition, RelDataType inputRowType, * Intersects across all referenced fields. */ private List resolveViableBackends(RexCall predicate, RelDataType inputRowType, - List fieldStorage, - List childViableBackends) { + List fieldStorage) { + // TODO : Try collapsing in one pass Set fieldIndices = new HashSet<>(); collectFieldIndices(predicate, fieldIndices); - - CapabilityRegistry registry = context.getCapabilityRegistry(); - if (fieldIndices.isEmpty()) { - return new ArrayList<>(registry.operatorBackends(OperatorCapability.FILTER)); + // No field refs (e.g. literal expression) — all backends with FILTER capability + return context.getBackends().values().stream() + .filter(backend -> backend.supportedOperators().contains(OperatorCapability.FILTER)) + .map(AnalyticsSearchBackendPlugin::name) + .collect(Collectors.toList()); } - FilterOperator operator = null; + FilterOperator filterOp = FilterOperator.fromSqlKind(predicate.getKind()); + + // Check if this is a full-text function (MATCH, MATCH_PHRASE, etc.) + FullTextOperator fullTextOp = null; if (predicate.getOperator() instanceof SqlFunction sqlFunction) { - operator = FullTextFunctions.toFilterOperator(sqlFunction); - } - if (operator == null) { - operator = FilterOperator.fromSqlKind(predicate.getKind()); - } - if (operator == null) { - throw new IllegalStateException("Unrecognized filter operator [" + predicate.getKind() + "]"); + fullTextOp = FullTextFunctions.toFullTextOperator(sqlFunction); } - Set viableSet = new HashSet<>(registry.operatorBackends(OperatorCapability.FILTER)); + // Start with all backends that have FILTER capability, intersect per field + Set viableSet = new HashSet<>(); + for (AnalyticsSearchBackendPlugin backend : context.getBackends().values()) { + if (backend.supportedOperators().contains(OperatorCapability.FILTER)) { + viableSet.add(backend.name()); + } + } for (int fieldIndex : fieldIndices) { if (fieldIndex >= fieldStorage.size()) { continue; } FieldStorageInfo storageInfo = fieldStorage.get(fieldIndex); - FieldType fieldType = storageInfo.getFieldType(); - if (fieldType == null) { - throw new IllegalStateException("Unrecognized field type [" + storageInfo.getMappingType() - + "] for field [" + storageInfo.getFieldName() + "]"); + + // Derived/expression column — only backends that can filter on expressions + if (storageInfo.isDerived()) { + viableSet.retainAll( + context.getBackends().values().stream() + .filter(b -> b.supportedOperators().contains(OperatorCapability.FILTER_ON_EXPRESSIONS)) + .map(AnalyticsSearchBackendPlugin::name) + .collect(Collectors.toSet()) + ); + continue; } + FieldTypeFamily typeFamily = FieldTypeFamily.fromMappingType(storageInfo.getFieldType()); + final FullTextOperator finalFullTextOp = fullTextOp; - // TODO: for FULL_TEXT operators, extract required params from RexCall - // and use registry.fullTextFilterBackends() instead Set fieldViable = new HashSet<>(); - if (storageInfo.isDerived()) { - // Derived column — only child viable backends + their delegation targets - List anyFormat = registry.filterBackendsAnyFormat(operator, fieldType); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.FILTER); - for (String name : childViableBackends) { - if (anyFormat.contains(name)) { - fieldViable.add(name); - } - } - // Delegation targets reachable from child viable backends - List delegationSupporters = registry.delegationSupporters(DelegationType.FILTER); - if (childViableBackends.stream().anyMatch(delegationSupporters::contains)) { - for (String name : anyFormat) { - if (!fieldViable.contains(name) && delegationAcceptors.contains(name)) { - fieldViable.add(name); - } - } - } - } else { - // Format-aware: backends that can access the field's data - for (String format : storageInfo.getDocValueFormats()) { - fieldViable.addAll(registry.filterBackends(operator, fieldType, format)); + for (AnalyticsSearchBackendPlugin backend : context.getBackends().values()) { + if (!viableSet.contains(backend.name())) { + continue; } - for (String format : storageInfo.getIndexFormats()) { - fieldViable.addAll(registry.filterBackends(operator, fieldType, format)); + + boolean formatMatch = backend.getSupportedFormats().stream().anyMatch(format -> + storageInfo.getDocValueFormats().contains(format.name()) + || storageInfo.getIndexFormats().contains(format.name()) + ); + + // Full-text ops require the backend to support that specific operator + // AND the field must have an index (full-text needs inverted index) + boolean operatorMatch; + if (finalFullTextOp != null) { + operatorMatch = backend.supportedFullTextOperators().contains(finalFullTextOp) + && storageInfo.hasIndex(); + } else if (filterOp != null && typeFamily != null) { + operatorMatch = backend.supportedFilterCapabilities() + .contains(FilterCapability.of(filterOp, typeFamily)); + } else { + // Unknown operator or unrecognized field type — accept if backend has FILTER + operatorMatch = true; } - // Format-agnostic: delegation targets that can evaluate but don't need data access - for (String name : registry.filterBackendsAnyFormat(operator, fieldType)) { - if (!fieldViable.contains(name)) { - fieldViable.add(name); - } + + if (formatMatch && operatorMatch) { + fieldViable.add(backend.name()); } } - viableSet.retainAll(fieldViable); } @@ -219,7 +226,7 @@ private List resolveViableBackends(RexCall predicate, RelDataType inputR throw new IllegalStateException("No backend can evaluate filter predicate [" + predicate.getKind() + "] on fields " + fieldIndices.stream() .filter(i -> i < fieldStorage.size()) - .map(i -> fieldStorage.get(i).getFieldName() + ":" + fieldStorage.get(i).getMappingType()) + .map(i -> fieldStorage.get(i).getFieldName() + ":" + fieldStorage.get(i).getFieldType()) .toList()); } return new ArrayList<>(viableSet); @@ -247,40 +254,55 @@ private void collectFieldIndices(RexNode node, Set result) { * (child supports FILTER delegation, this backend accepts it, and this backend * can handle all predicates natively). */ - private List computeFilterViableBackends(RexNode annotatedCondition, - List childViableBackends) { + private List computeFilterViableBackends(RexNode annotatedCondition, String childBackend) { List predicates = new ArrayList<>(); collectAnnotatedPredicates(annotatedCondition, predicates); if (predicates.isEmpty()) { - return new ArrayList<>(childViableBackends); + return new ArrayList<>(context.getBackends().keySet()); } + AnalyticsSearchBackendPlugin childPlugin = context.getBackends().get(childBackend); List viable = new ArrayList<>(); - CapabilityRegistry registry = context.getCapabilityRegistry(); - List delegationSupporters = registry.delegationSupporters(DelegationType.FILTER); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.FILTER); - for (String candidateName : childViableBackends) { - if (!registry.operatorBackends(OperatorCapability.FILTER).contains(candidateName)) { + for (AnalyticsSearchBackendPlugin candidate : context.getBackends().values()) { + if (!candidate.supportedOperators().contains(OperatorCapability.FILTER)) { continue; } + boolean isChild = candidate.name().equals(childBackend); + + // Non-child backend: only viable if child can delegate entire filter to it + // and it accepts delegation and can handle all predicates natively + if (!isChild) { + if (childPlugin == null + || !childPlugin.supportedDelegations().contains(DelegationType.FILTER) + || !candidate.acceptedDelegations().contains(DelegationType.FILTER)) { + continue; + } + } + boolean canHandleAll = true; for (AnnotatedPredicate predicate : predicates) { List predViable = predicate.getViableBackends(); - if (predViable.contains(candidateName)) { + if (predViable.contains(candidate.name())) { continue; } - if (delegationSupporters.contains(candidateName) - && predViable.stream().anyMatch(delegationAcceptors::contains)) { - continue; + // Child backend can delegate individual predicates + if (isChild && candidate.supportedDelegations().contains(DelegationType.FILTER)) { + boolean someoneAccepts = predViable.stream().anyMatch(backendName -> { + AnalyticsSearchBackendPlugin other = context.getBackends().get(backendName); + return other != null && other.acceptedDelegations().contains(DelegationType.FILTER); + }); + if (someoneAccepts) { + continue; + } } canHandleAll = false; break; } if (canHandleAll) { - viable.add(candidateName); + viable.add(candidate.name()); } } return viable; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java index 161291de1f85a..815ec285916f4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java @@ -15,14 +15,13 @@ import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlFunction; -import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.AnnotatedProjectExpression; import org.opensearch.analytics.planner.rel.OpenSearchProject; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; -import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.ScalarFunction; import java.util.ArrayList; @@ -59,19 +58,11 @@ public void onMatch(RelOptRuleCall call) { "Project rule encountered unmarked child [" + child.getClass().getSimpleName() + "]"); } - List childViableBackends = openSearchChild.getViableBackends(); + String backend = openSearchChild.getBackend(); - // TODO: precompute SqlKind → viable backends map to avoid repeated filtering per node - // TODO: reuse childViableBackends list when all candidates pass instead of allocating List annotatedExprs = new ArrayList<>(project.getProjects().size()); for (RexNode expr : project.getProjects()) { - annotatedExprs.add(annotateExpr(expr, childViableBackends)); - } - - List viableBackends = computeProjectViableBackends(annotatedExprs, childViableBackends); - if (viableBackends.isEmpty()) { - throw new IllegalStateException( - "No backend can execute all project expressions among " + childViableBackends); + annotatedExprs.add(annotateExpr(expr, backend)); } call.transformTo(new OpenSearchProject( @@ -80,41 +71,49 @@ public void onMatch(RelOptRuleCall call) { RelNodeUtils.unwrapHep(project.getInput()), annotatedExprs, project.getRowType(), - viableBackends + backend )); } - private RexNode annotateExpr(RexNode expr, List childViableBackends) { + private RexNode annotateExpr(RexNode expr, String backend) { if (!(expr instanceof RexCall rexCall)) { return expr; } - // Opaque operations — no recursion into operands + AnalyticsSearchBackendPlugin plugin = context.getBackends().get(backend); + if (plugin == null) { + throw new IllegalStateException("Backend [" + backend + "] not found"); + } + + // Opaque operations (painless, highlighting, etc.) — no recursion into operands if (rexCall.getOperator() instanceof SqlFunction sqlFunction) { String funcName = sqlFunction.getName(); if (isOpaqueOperation(funcName)) { - List exprViable = resolveOpaqueViableBackends(funcName, childViableBackends); - if (exprViable.isEmpty()) { - throw new IllegalStateException( - "No backend can evaluate [" + funcName + "] and no delegation path exists"); + if (plugin.supportedOpaqueProjectOperations().contains(funcName)) { + return new AnnotatedProjectExpression(rexCall.getType(), rexCall, backend); + } + String delegationTarget = canDelegateProject(plugin, funcName); + if (delegationTarget != null) { + return new AnnotatedProjectExpression(rexCall.getType(), rexCall, delegationTarget); } - return new AnnotatedProjectExpression(rexCall.getType(), rexCall, exprViable); + throw new IllegalStateException( + "Backend [" + backend + "] cannot evaluate [" + funcName + + "] and no delegation path exists"); } } - // Standard scalar function - List scalarViable = resolveScalarViableBackends(rexCall, childViableBackends); - if (scalarViable.isEmpty()) { + // Standard scalar function — validate support + ScalarFunction scalarFunc = ScalarFunction.fromSqlKind(rexCall.getKind()); + if (scalarFunc != null && !plugin.supportedScalarFunctions().contains(scalarFunc)) { throw new IllegalStateException( - "No backend supports scalar function [" + ScalarFunction.fromSqlKind(rexCall.getKind()) - + "] among " + childViableBackends); + "Backend [" + backend + "] does not support scalar function [" + scalarFunc + "]"); } - // Recurse into operands + // Recurse into operands to validate and annotate nested expressions boolean changed = false; List newOperands = new ArrayList<>(rexCall.getOperands().size()); for (RexNode operand : rexCall.getOperands()) { - RexNode annotated = annotateExpr(operand, childViableBackends); + RexNode annotated = annotateExpr(operand, backend); newOperands.add(annotated); if (annotated != operand) { changed = true; @@ -122,98 +121,24 @@ private RexNode annotateExpr(RexNode expr, List childViableBackends) { } RexCall target = changed ? rexCall.clone(rexCall.getType(), newOperands) : rexCall; - return new AnnotatedProjectExpression(target.getType(), target, scalarViable); - } - - private List resolveOpaqueViableBackends(String funcName, List childViableBackends) { - CapabilityRegistry registry = context.getCapabilityRegistry(); - List viable = registry.opaqueBackendsAnyFormat(funcName); - if (viable.isEmpty()) { - return viable; - } - // At least one child viable backend must be able to reach an evaluator: - // either it's in viable itself (native), or it can delegate to one that accepts - List delegationSupporters = registry.delegationSupporters(DelegationType.PROJECT); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); - boolean reachable = childViableBackends.stream().anyMatch(candidateName -> - viable.contains(candidateName) - || (delegationSupporters.contains(candidateName) - && viable.stream().anyMatch(delegationAcceptors::contains))); - return reachable ? viable : List.of(); + return new AnnotatedProjectExpression(target.getType(), target, backend); } - private List resolveScalarViableBackends(RexCall rexCall, List childViableBackends) { - ScalarFunction scalarFunc = ScalarFunction.fromSqlKind(rexCall.getKind()); - if (scalarFunc == null) { - return List.of(); - } - FieldType fieldType = FieldType.fromSqlTypeName(rexCall.getType().getSqlTypeName()); - if (fieldType == null) { - return List.of(); - } - - CapabilityRegistry registry = context.getCapabilityRegistry(); - List allCapable = registry.scalarBackendsAnyFormat(scalarFunc, fieldType); - - // Prefer child viable backends - List viable = new ArrayList<>(); - for (String candidateName : childViableBackends) { - if (allCapable.contains(candidateName)) { - viable.add(candidateName); - } - } - if (!viable.isEmpty()) { - return viable; - } - // Fallback: other backends if reachable via delegation - List delegationSupporters = registry.delegationSupporters(DelegationType.PROJECT); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); - boolean canDelegate = childViableBackends.stream().anyMatch(delegationSupporters::contains); - if (!canDelegate) { - return viable; - } - for (String backendName : allCapable) { - if (delegationAcceptors.contains(backendName)) { - viable.add(backendName); - } - } - return viable; - } - - private List computeProjectViableBackends(List annotatedExprs, - List childViableBackends) { - // A child viable backend is viable for the project if for every expression it can - // either evaluate natively (present in expression's viableBackends) or delegate to - // a backend that can (supports PROJECT delegation to an acceptor in expression's viableBackends) - CapabilityRegistry registry = context.getCapabilityRegistry(); - List delegationSupporters = registry.delegationSupporters(DelegationType.PROJECT); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); - - List result = new ArrayList<>(); - for (String candidateName : childViableBackends) { - boolean canHandleAll = true; - for (RexNode expr : annotatedExprs) { - if (!(expr instanceof AnnotatedProjectExpression annotation)) { - continue; - } - if (annotation.getViableBackends().contains(candidateName)) { - continue; - } - boolean canDelegate = delegationSupporters.contains(candidateName) - && annotation.getViableBackends().stream().anyMatch(delegationAcceptors::contains); - if (!canDelegate) { - canHandleAll = false; - break; - } - } - if (canHandleAll) { - result.add(candidateName); - } - } - return result; + private boolean isOpaqueOperation(String funcName) { + return context.getBackends().values().stream() + .anyMatch(b -> b.supportedOpaqueProjectOperations().contains(funcName)); } - private boolean isOpaqueOperation(String funcName) { - return context.getCapabilityRegistry().isOpaqueOperation(funcName); + private String canDelegateProject(AnalyticsSearchBackendPlugin plugin, String funcName) { + if (!plugin.supportedDelegations().contains(DelegationType.PROJECT)) { + return null; + } + return context.getBackends().values().stream() + .filter(other -> !other.name().equals(plugin.name()) + && other.acceptedDelegations().contains(DelegationType.PROJECT) + && other.supportedOpaqueProjectOperations().contains(funcName)) + .map(AnalyticsSearchBackendPlugin::name) + .findFirst() + .orElse(null); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java index e2503646c1b70..3f72755c0f02e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java @@ -12,6 +12,7 @@ import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.Sort; +import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; @@ -54,16 +55,13 @@ public void onMatch(RelOptRuleCall call) { "Sort rule encountered unmarked child [" + child.getClass().getSimpleName() + "]"); } - List childViableBackends = openSearchChild.getViableBackends(); - List sortCapable = context.getCapabilityRegistry().operatorBackends(OperatorCapability.SORT); + String childBackend = openSearchChild.getBackend(); - List viableBackends = childViableBackends.stream() - .filter(sortCapable::contains) - .toList(); + String backend = CapabilityResolutionUtils.resolveBackend( + context.getBackends(), childBackend, OperatorCapability.SORT); - if (viableBackends.isEmpty()) { - throw new IllegalStateException( - "No backend supports SORT capability among " + childViableBackends); + if (!CapabilityResolutionUtils.backendSupports(context.getBackends(), backend, OperatorCapability.SORT)) { + throw new IllegalStateException("No backend supports SORT capability"); } call.transformTo(new OpenSearchSort( @@ -73,7 +71,8 @@ public void onMatch(RelOptRuleCall call) { sort.getCollation(), sort.offset, sort.fetch, - viableBackends + backend, + List.of(backend) )); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java index 56bd6e87cabaf..12310c2c83683 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java @@ -12,16 +12,14 @@ import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.rel.core.TableScan; import org.apache.calcite.rel.type.RelDataTypeField; -import org.opensearch.analytics.planner.CapabilityRegistry; +import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.FieldStorageInfo; import org.opensearch.analytics.planner.FieldStorageResolver; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; -import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.cluster.metadata.IndexMetadata; -import java.util.ArrayList; import java.util.List; /** @@ -56,54 +54,20 @@ public void onMatch(RelOptRuleCall call) { } String primaryFormat = indexMetadata.getSettings().get("index.composite.primary_data_format", "lucene"); - CapabilityRegistry registry = context.getCapabilityRegistry(); - FieldStorageResolver fieldStorageResolver = new FieldStorageResolver(indexMetadata, - registry.getBackends()); + List viableBackends = CapabilityResolutionUtils.computeViableBackends( + context.getBackends(), OperatorCapability.SCAN, primaryFormat); + if (viableBackends.isEmpty()) { + throw new IllegalStateException("No backend supports format [" + primaryFormat + + "] with SCAN capability for index [" + indexMetadata.getIndex().getName() + "]"); + } - // TODO : This expects the FrontEnds to attach the row type with all fields. - // TODO : How will they attach if we perform the index resolution List fieldNames = scan.getRowType().getFieldList().stream() .map(RelDataTypeField::getName) .toList(); - List fieldStorage = fieldStorageResolver.resolve(fieldNames); - - // Viable backends: must support SCAN and be able to read ALL requested fields - // (natively or via delegation to another backend that can read the field) - List scanCapable = registry.operatorBackends(OperatorCapability.SCAN); - List delegationSupporters = registry.delegationSupporters(DelegationType.SCAN); - List delegationAcceptors = registry.delegationAcceptors(DelegationType.SCAN); - List viableBackends = new ArrayList<>(scanCapable); - - for (FieldStorageInfo field : fieldStorage) { - if (field.isDerived()) { - continue; - } - // Backends that can natively scan this field's doc values - List fieldBackends = new ArrayList<>(); - for (String format : field.getDocValueFormats()) { - for (String backend : registry.scanBackends(format)) { - if (!fieldBackends.contains(backend)) { - fieldBackends.add(backend); - } - } - } - // Keep candidates that can scan natively or delegate to one that can - viableBackends.removeIf(candidate -> { - if (fieldBackends.contains(candidate)) { - return false; - } - return !delegationSupporters.contains(candidate) - || fieldBackends.stream().noneMatch(delegationAcceptors::contains); - }); - } - - if (viableBackends.isEmpty()) { - throw new IllegalStateException("No backend can scan all requested fields on index [" - + indexMetadata.getIndex().getName() + "]"); - } + List fieldStorage = FieldStorageResolver.resolve(indexMetadata, fieldNames); call.transformTo(OpenSearchTableScan.create( - scan.getCluster(), scan.getTable(), viableBackends, fieldStorage, + scan.getCluster(), scan.getTable(), viableBackends.getFirst(), viableBackends, fieldStorage, indexMetadata.getNumberOfShards(), context.getDistributionTraitDef() )); } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java index 1a3990bdbc1d6..8a7ee88156a1e 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java @@ -22,10 +22,9 @@ import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; -import org.opensearch.analytics.spi.AggregateCapability; import org.opensearch.analytics.spi.AggregateFunction; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; -import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.OperatorCapability; import java.util.EnumSet; @@ -52,7 +51,7 @@ public void testPerCallAnnotation() { public void testViableBackendsPopulated() { OpenSearchAggregate agg = runAggregate(1, sumCall()); - assertTrue(agg.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, agg.getBackend()); assertFalse(agg.getViableBackends().isEmpty()); assertTrue(agg.getViableBackends().contains(MockDataFusionBackend.NAME)); } @@ -100,7 +99,7 @@ public void testAggregateErrorsWhenNoBackendSupportsFunction() { "size", Map.of("type", "integer") ), List.of(new MockDataFusionBackend() { @Override - public Set aggregateCapabilities() { + public Set supportedAggregateFunctions() { return Set.of(); // supports AGGREGATE capability but no functions } })); @@ -120,10 +119,8 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set aggregateCapabilities() { - return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( - AggregateFunction.SUM, Set.of(FieldType.INTEGER), - AggregateFunction.COUNT, Set.of(FieldType.INTEGER))); + public Set supportedAggregateFunctions() { + return EnumSet.of(AggregateFunction.SUM, AggregateFunction.COUNT); } }; @@ -137,13 +134,7 @@ public Set aggregateCapabilities() { OpenSearchAggregate agg = (OpenSearchAggregate) result; assertTrue(agg.getViableBackends().contains(MockDataFusionBackend.NAME)); - // Lucene not viable at operator level (no doc values for scan) - // but per-call annotation shows Lucene as viable for SUM - assertFalse(agg.getViableBackends().contains(MockLuceneBackend.NAME)); - AggregateCallAnnotation annotation = AggregateCallAnnotation.find(agg.getAggCallList().get(0)); - assertNotNull(annotation); - assertTrue(annotation.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(annotation.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(agg.getViableBackends().contains(MockLuceneBackend.NAME)); } // ---- Scan ---- @@ -159,7 +150,7 @@ public void testTableScanResolvesBackendAndFieldStorage() { assertTrue(result instanceof OpenSearchTableScan); OpenSearchTableScan scan = (OpenSearchTableScan) result; - assertTrue(scan.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, scan.getBackend()); assertEquals(2, scan.getOutputFieldStorage().size()); assertEquals("status", scan.getOutputFieldStorage().get(0).getFieldName()); assertFalse(scan.getViableBackends().isEmpty()); @@ -232,9 +223,8 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set aggregateCapabilities() { - return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( - AggregateFunction.SUM, Set.of(FieldType.INTEGER))); + public Set supportedAggregateFunctions() { + return EnumSet.of(AggregateFunction.SUM); // no COUNT } }; @@ -286,9 +276,8 @@ public void testReducerPassthroughViableBackends() { public void testAggregateViableWithDelegation() { MockDataFusionBackend dfWithDelegation = new MockDataFusionBackend() { @Override - public Set aggregateCapabilities() { - return aggCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), Map.of( - AggregateFunction.SUM, Set.of(FieldType.INTEGER))); + public Set supportedAggregateFunctions() { + return EnumSet.of(AggregateFunction.SUM); } @Override public Set supportedDelegations() { @@ -301,9 +290,8 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set aggregateCapabilities() { - return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( - AggregateFunction.STDDEV_POP, Set.of(FieldType.INTEGER))); + public Set supportedAggregateFunctions() { + return EnumSet.of(AggregateFunction.STDDEV_POP); } @Override public Set acceptedDelegations() { @@ -333,9 +321,8 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set aggregateCapabilities() { - return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( - AggregateFunction.STDDEV_POP, Set.of(FieldType.INTEGER))); + public Set supportedAggregateFunctions() { + return EnumSet.of(AggregateFunction.STDDEV_POP); } }; @@ -348,17 +335,6 @@ public Set aggregateCapabilities() { IllegalStateException exception = expectThrows(IllegalStateException.class, () -> runPlanner(aggregate, context)); - assertTrue(exception.getMessage().contains("not supported by any viable backend")); - } - - private static Set aggCaps(Set formats, - Map> funcToTypes) { - Set caps = new java.util.HashSet<>(); - for (var entry : funcToTypes.entrySet()) { - for (FieldType family : entry.getValue()) { - caps.add(new AggregateCapability(entry.getKey(), family, formats)); - } - } - return caps; + assertTrue(exception.getMessage().contains("no delegation path exists")); } } 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 2104bada7ba7b..7b5a8bfb62ce7 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 @@ -114,7 +114,7 @@ protected PlannerContext buildContext(String primaryFormat, int shardCount, ClusterState clusterState = mock(ClusterState.class); when(clusterState.metadata()).thenReturn(metadata); - return new PlannerContext(new CapabilityRegistry(backends), clusterState); + return new PlannerContext(backends, clusterState); } // ---- Table builders ---- diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java index 153221d4cd9fd..3c68bc5d9308b 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java @@ -18,7 +18,6 @@ import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.util.ImmutableBitSet; -import org.apache.calcite.rex.RexCall; import org.opensearch.analytics.planner.rel.AnnotatedPredicate; import org.opensearch.analytics.planner.rel.FullTextFunctions; import org.opensearch.analytics.planner.rel.OpenSearchFilter; @@ -45,7 +44,7 @@ public void testNativePredicateAnnotatedWithBothBackends() { ), new String[]{"status", "size"}, new SqlTypeName[]{SqlTypeName.INTEGER, SqlTypeName.INTEGER}, makeEquals(0, SqlTypeName.INTEGER, 200)); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertTrue(result.getCondition() instanceof AnnotatedPredicate); AnnotatedPredicate annotated = (AnnotatedPredicate) result.getCondition(); assertTrue(annotated.getViableBackends().contains(MockDataFusionBackend.NAME)); @@ -72,26 +71,20 @@ public void testKeywordEqualsAnnotatedWithBothBackends() { /** Full-text with delegation — both backends viable at operator level. */ public void testFullTextViableWithDelegation() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( - "message", Map.of("type", "keyword") + "message", Map.of("type", "text") ), new String[]{"message"}, new SqlTypeName[]{SqlTypeName.VARCHAR}, makeFullTextCall(FullTextFunctions.MATCH_PHRASE, 0, "hello world")); - // DF is viable at operator level (has doc values in parquet) + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - // Lucene not viable at operator level — only delegation target - assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); - // MATCH_PHRASE predicate has Lucene as delegation target - AnnotatedPredicate predicate = (AnnotatedPredicate) result.getCondition(); - assertTrue("MATCH_PHRASE should be evaluable by Lucene", - predicate.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(predicate.getOriginal().toString().contains("MATCH_PHRASE")); + assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); } - /** AND with delegation — DF viable at operator, equals viable for both, MATCH_PHRASE delegated to Lucene. */ + /** AND with delegation — both backends viable. */ public void testAndWithDelegationBothViable() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( "status", Map.of("type", "integer"), - "message", Map.of("type", "keyword") + "message", Map.of("type", "text") ), new String[]{"status", "message"}, new SqlTypeName[]{SqlTypeName.INTEGER, SqlTypeName.VARCHAR}, makeAnd( makeEquals(0, SqlTypeName.INTEGER, 200), @@ -99,21 +92,14 @@ public void testAndWithDelegationBothViable() { )); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); - RexCall andCondition = (RexCall) result.getCondition(); - AnnotatedPredicate equalsPred = (AnnotatedPredicate) andCondition.getOperands().get(0); - AnnotatedPredicate matchPred = (AnnotatedPredicate) andCondition.getOperands().get(1); - assertTrue(equalsPred.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(equalsPred.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(matchPred.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(matchPred.getOriginal().toString().contains("MATCH_PHRASE")); + assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); } - /** OR across backends — DF viable at operator, equals viable for both, MATCH delegated to Lucene. */ + /** OR across backends with delegation — both viable. */ public void testOrAcrossBackendsWithDelegation() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( "status", Map.of("type", "integer"), - "message", Map.of("type", "keyword") + "message", Map.of("type", "text") ), new String[]{"status", "message"}, new SqlTypeName[]{SqlTypeName.INTEGER, SqlTypeName.VARCHAR}, makeCall(SqlStdOperatorTable.OR, makeEquals(0, SqlTypeName.INTEGER, 200), @@ -121,21 +107,14 @@ public void testOrAcrossBackendsWithDelegation() { )); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); - RexCall orCondition = (RexCall) result.getCondition(); - AnnotatedPredicate equalsPred = (AnnotatedPredicate) orCondition.getOperands().get(0); - AnnotatedPredicate matchPred = (AnnotatedPredicate) orCondition.getOperands().get(1); - assertTrue(equalsPred.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(equalsPred.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(matchPred.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(matchPred.getOriginal().toString().contains("MATCH")); + assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); } - /** OR of two full-text predicates — DF viable at operator, both predicates delegated to Lucene. */ + /** OR of two full-text predicates with delegation — both viable. */ public void testMultipleFullTextOrWithDelegation() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( - "title", Map.of("type", "keyword"), - "body", Map.of("type", "keyword") + "title", Map.of("type", "text"), + "body", Map.of("type", "text") ), new String[]{"title", "body"}, new SqlTypeName[]{SqlTypeName.VARCHAR, SqlTypeName.VARCHAR}, makeCall(SqlStdOperatorTable.OR, makeFullTextCall(FullTextFunctions.MATCH, 0, "hello"), @@ -143,14 +122,7 @@ public void testMultipleFullTextOrWithDelegation() { )); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); - RexCall orCondition = (RexCall) result.getCondition(); - AnnotatedPredicate matchPred = (AnnotatedPredicate) orCondition.getOperands().get(0); - AnnotatedPredicate phrasePred = (AnnotatedPredicate) orCondition.getOperands().get(1); - assertTrue(matchPred.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(matchPred.getOriginal().toString().contains("MATCH")); - assertTrue(phrasePred.getViableBackends().contains(MockLuceneBackend.NAME)); - assertTrue(phrasePred.getOriginal().toString().contains("MATCH_PHRASE")); + assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); } // ---- Error cases ---- @@ -162,7 +134,7 @@ public void testFullTextErrorsWithoutDelegation() { LogicalFilter filter = LogicalFilter.create(stubScan(table), condition); PlannerContext context = buildContext("parquet", Map.of( - "message", Map.of("type", "keyword") + "message", Map.of("type", "text") )); IllegalStateException exception = expectThrows(IllegalStateException.class, @@ -182,9 +154,12 @@ public void testErrorForUnsupportedFieldTypeOperatorCombo() { "location", Map.of("type", "geo_point") )); - IllegalStateException exception = expectThrows(IllegalStateException.class, - () -> runPlanner(filter, context)); - assertTrue(exception.getMessage().contains("has no storage")); + try { + runPlanner(filter, context); + fail("Expected IllegalStateException for unsupported field type"); + } catch (IllegalStateException e) { + assertTrue(e.getMessage().contains("No backend can")); + } } // ---- Derived columns ---- 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 8335338c14315..7d93f3a58c5d7 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 @@ -11,48 +11,37 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; -import org.opensearch.analytics.spi.AggregateCapability; import org.opensearch.analytics.spi.AggregateFunction; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; -import org.opensearch.analytics.spi.FieldType; +import org.opensearch.analytics.spi.FieldTypeFamily; import org.opensearch.analytics.spi.FilterCapability; import org.opensearch.analytics.spi.FilterOperator; +import org.opensearch.analytics.spi.FullTextOperator; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; -import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.COLUMNAR_STORAGE; - -import java.util.HashSet; +import java.util.EnumSet; import java.util.List; import java.util.Set; +import java.util.stream.Collectors; /** - * Mock DataFusion backend for tests. Supports parquet format with columnar storage, - * standard filter operators on NUMERIC/KEYWORD/DATE/BOOLEAN, and common aggregates. + * Mock DataFusion backend for tests. Supports parquet format, standard filter + * operators on NUMERIC/KEYWORD/DATE/BOOLEAN families, and common aggregates. * No full-text support. */ public class MockDataFusionBackend implements AnalyticsSearchBackendPlugin { public static final String NAME = "mock-parquet"; - public static final String PARQUET_DATA_FORMAT = "parquet"; - private static final Set DATAFUSION_FORMATS = Set.of(PARQUET_DATA_FORMAT); private static final Set OPERATOR_CAPS = Set.of( - OperatorCapability.SCAN, OperatorCapability.FILTER, + OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.FILTER_ON_EXPRESSIONS, OperatorCapability.AGGREGATE, OperatorCapability.SORT, OperatorCapability.PROJECT, OperatorCapability.COORDINATOR_REDUCE ); - private static final Set SUPPORTED_TYPES = new HashSet<>(); - static { - SUPPORTED_TYPES.addAll(FieldType.numeric()); - SUPPORTED_TYPES.addAll(FieldType.keyword()); - SUPPORTED_TYPES.addAll(FieldType.date()); - SUPPORTED_TYPES.add(FieldType.BOOLEAN); - } - - private static final Set STANDARD_OPS = Set.of( + private static final Set FILTER_OPS = EnumSet.of( FilterOperator.EQUALS, FilterOperator.NOT_EQUALS, FilterOperator.GREATER_THAN, FilterOperator.GREATER_THAN_OR_EQUAL, FilterOperator.LESS_THAN, FilterOperator.LESS_THAN_OR_EQUAL, @@ -60,57 +49,60 @@ public class MockDataFusionBackend implements AnalyticsSearchBackendPlugin { FilterOperator.IN, FilterOperator.LIKE ); - private static final Set AGG_FUNCTIONS = Set.of( - AggregateFunction.SUM, AggregateFunction.SUM0, - AggregateFunction.MIN, AggregateFunction.MAX, - AggregateFunction.COUNT, AggregateFunction.AVG + private static final Set SUPPORTED_FAMILIES = EnumSet.of( + FieldTypeFamily.NUMERIC, FieldTypeFamily.KEYWORD, + FieldTypeFamily.DATE, FieldTypeFamily.BOOLEAN ); private static final Set FILTER_CAPS; static { - Set caps = new HashSet<>(); - for (FilterOperator op : STANDARD_OPS) { - for (FieldType type : SUPPORTED_TYPES) { - caps.add(new FilterCapability.Standard(op, type, DATAFUSION_FORMATS)); - } - } - FILTER_CAPS = caps; + FILTER_CAPS = FILTER_OPS.stream() + .flatMap(op -> SUPPORTED_FAMILIES.stream().map(family -> FilterCapability.of(op, family))) + .collect(Collectors.toUnmodifiableSet()); } - private static final Set AGG_CAPS; - static { - Set caps = new HashSet<>(); - for (AggregateFunction func : AGG_FUNCTIONS) { - for (FieldType type : SUPPORTED_TYPES) { - caps.add(AggregateCapability.simple(func, type, DATAFUSION_FORMATS)); - } - } - AGG_CAPS = caps; - } + private static final Set AGG_FUNCTIONS = EnumSet.of( + AggregateFunction.SUM, AggregateFunction.SUM0, + AggregateFunction.MIN, AggregateFunction.MAX, + AggregateFunction.COUNT, AggregateFunction.AVG + ); - @Override public String name() { return NAME; } + @Override + public String name() { + return NAME; + } - @Override public SearchExecEngine searcher(ExecutionContext ctx) { return null; } + @Override + public SearchExecEngine searcher(ExecutionContext ctx) { + return null; + } @Override public List getSupportedFormats() { return List.of(new DataFormat() { - @Override public String name() { return PARQUET_DATA_FORMAT; } + @Override public String name() { return "parquet"; } @Override public long priority() { return 0; } - @Override public Set supportedFields() { - return Set.of( - new FieldTypeCapabilities("integer", Set.of(COLUMNAR_STORAGE)), - new FieldTypeCapabilities("long", Set.of(COLUMNAR_STORAGE)), - new FieldTypeCapabilities("keyword", Set.of(COLUMNAR_STORAGE)), - new FieldTypeCapabilities("text", Set.of(COLUMNAR_STORAGE)), - new FieldTypeCapabilities("boolean", Set.of(COLUMNAR_STORAGE)), - new FieldTypeCapabilities("date", Set.of(COLUMNAR_STORAGE)) - ); - } + @Override public Set supportedFields() { return Set.of(); } }); } - @Override public Set supportedOperators() { return OPERATOR_CAPS; } - @Override public Set filterCapabilities() { return FILTER_CAPS; } - @Override public Set aggregateCapabilities() { return AGG_CAPS; } + @Override + public Set supportedFilterCapabilities() { + return FILTER_CAPS; + } + + @Override + public Set supportedFullTextOperators() { + return Set.of(); + } + + @Override + public Set supportedOperators() { + return OPERATOR_CAPS; + } + + @Override + public Set supportedAggregateFunctions() { + return AGG_FUNCTIONS; + } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java index 4b34e36e93dec..e31534e768f0e 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java @@ -11,39 +11,35 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; +import org.opensearch.analytics.spi.AggregateFunction; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; -import org.opensearch.analytics.spi.FieldType; +import org.opensearch.analytics.spi.FieldTypeFamily; import org.opensearch.analytics.spi.FilterCapability; import org.opensearch.analytics.spi.FilterOperator; +import org.opensearch.analytics.spi.FullTextOperator; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; -import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.FULL_TEXT_SEARCH; -import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.POINT_RANGE; -import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.STORED_FIELDS; - -import java.util.HashSet; +import java.util.EnumSet; import java.util.List; import java.util.Set; +import java.util.stream.Collectors; /** - * Mock Lucene backend for tests. Supports lucene format with index structures - * (full-text, point range) and stored fields. No columnar storage (doc values - * only in parquet for default tests). Standard + full-text filter capabilities. - * SCAN + FILTER only (no AGGREGATE). + * Mock Lucene backend for tests. Supports lucene format, standard filter + * operators on NUMERIC/KEYWORD/TEXT/DATE/BOOLEAN families, all full-text + * operators, and SCAN + FILTER capabilities (no AGGREGATE). */ public class MockLuceneBackend implements AnalyticsSearchBackendPlugin { public static final String NAME = "mock-lucene"; - public static final String LUCENE_DATA_FORMAT = "lucene"; - private static final Set LUCENE_FORMATS = Set.of(LUCENE_DATA_FORMAT); private static final Set OPERATOR_CAPS = Set.of( OperatorCapability.SCAN, OperatorCapability.FILTER ); - private static final Set STANDARD_OPS = Set.of( + private static final Set FILTER_OPS = EnumSet.of( FilterOperator.EQUALS, FilterOperator.NOT_EQUALS, FilterOperator.GREATER_THAN, FilterOperator.GREATER_THAN_OR_EQUAL, FilterOperator.LESS_THAN, FilterOperator.LESS_THAN_OR_EQUAL, @@ -51,67 +47,54 @@ public class MockLuceneBackend implements AnalyticsSearchBackendPlugin { FilterOperator.IN, FilterOperator.LIKE ); - private static final Set FULL_TEXT_OPS = Set.of( - FilterOperator.MATCH, FilterOperator.MATCH_PHRASE, - FilterOperator.MATCH_PHRASE_PREFIX, FilterOperator.MATCH_BOOL_PREFIX, - FilterOperator.MULTI_MATCH, FilterOperator.QUERY_STRING, - FilterOperator.SIMPLE_QUERY_STRING, FilterOperator.FUZZY, - FilterOperator.WILDCARD, FilterOperator.REGEXP + private static final Set SUPPORTED_FAMILIES = EnumSet.of( + FieldTypeFamily.NUMERIC, FieldTypeFamily.KEYWORD, + FieldTypeFamily.TEXT, FieldTypeFamily.DATE, FieldTypeFamily.BOOLEAN ); - private static final Set STANDARD_TYPES = new HashSet<>(); + private static final Set FILTER_CAPS; static { - STANDARD_TYPES.addAll(FieldType.numeric()); - STANDARD_TYPES.addAll(FieldType.keyword()); - STANDARD_TYPES.addAll(FieldType.text()); - STANDARD_TYPES.addAll(FieldType.date()); - STANDARD_TYPES.add(FieldType.BOOLEAN); + FILTER_CAPS = FILTER_OPS.stream() + .flatMap(op -> SUPPORTED_FAMILIES.stream().map(family -> FilterCapability.of(op, family))) + .collect(Collectors.toUnmodifiableSet()); } - private static final Set FULL_TEXT_TYPES = new HashSet<>(); - static { - FULL_TEXT_TYPES.addAll(FieldType.keyword()); - FULL_TEXT_TYPES.addAll(FieldType.text()); + @Override + public String name() { + return NAME; } - private static final Set FILTER_CAPS; - static { - Set caps = new HashSet<>(); - for (FilterOperator op : STANDARD_OPS) { - for (FieldType type : STANDARD_TYPES) { - caps.add(new FilterCapability.Standard(op, type, LUCENE_FORMATS)); - } - } - for (FilterOperator op : FULL_TEXT_OPS) { - for (FieldType type : FULL_TEXT_TYPES) { - caps.add(new FilterCapability.FullText(op, type, LUCENE_FORMATS, Set.of())); - } - } - FILTER_CAPS = caps; + @Override + public SearchExecEngine searcher(ExecutionContext ctx) { + return null; } - @Override public String name() { return NAME; } - - @Override public SearchExecEngine searcher(ExecutionContext ctx) { return null; } - @Override public List getSupportedFormats() { return List.of(new DataFormat() { - @Override public String name() { return LUCENE_DATA_FORMAT; } + @Override public String name() { return "lucene"; } @Override public long priority() { return 0; } - @Override public Set supportedFields() { - return Set.of( - new FieldTypeCapabilities("integer", Set.of(POINT_RANGE, STORED_FIELDS)), - new FieldTypeCapabilities("long", Set.of(POINT_RANGE, STORED_FIELDS)), - new FieldTypeCapabilities("keyword", Set.of(FULL_TEXT_SEARCH, STORED_FIELDS)), - new FieldTypeCapabilities("text", Set.of(FULL_TEXT_SEARCH, STORED_FIELDS)), - new FieldTypeCapabilities("boolean", Set.of(STORED_FIELDS)), - new FieldTypeCapabilities("date", Set.of(POINT_RANGE, STORED_FIELDS)) - ); - } + @Override public Set supportedFields() { return Set.of(); } }); } - @Override public Set supportedOperators() { return OPERATOR_CAPS; } - @Override public Set filterCapabilities() { return FILTER_CAPS; } + @Override + public Set supportedFilterCapabilities() { + return FILTER_CAPS; + } + + @Override + public Set supportedFullTextOperators() { + return EnumSet.allOf(FullTextOperator.class); + } + + @Override + public Set supportedOperators() { + return OPERATOR_CAPS; + } + + @Override + public Set supportedAggregateFunctions() { + return Set.of(); + } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java index ef1d1b8c37ba5..1b67296191aeb 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java @@ -24,12 +24,9 @@ import org.opensearch.analytics.planner.rel.OpenSearchProject; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; -import org.opensearch.analytics.spi.FieldType; -import org.opensearch.analytics.spi.ProjectCapability; import org.opensearch.analytics.spi.ScalarFunction; import java.util.EnumSet; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; @@ -58,7 +55,7 @@ public void testSimpleFieldProjection() { rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0), rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) ); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); for (RexNode expr : result.getProjects()) { assertFalse("Field ref should not be annotated", expr instanceof AnnotatedProjectExpression); } @@ -73,7 +70,7 @@ public void testSupportedScalarFunction() { rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) ); OpenSearchProject result = runProject(castExpr); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); } @@ -93,8 +90,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } }; @@ -102,14 +99,9 @@ public Set projectCapabilities() { rexBuilder.makeCall(PAINLESS, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - // Operator: only DF viable (Lucene not viable for scan) - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); - // painless expression: only Lucene can natively evaluate it - AnnotatedProjectExpression annotation = (AnnotatedProjectExpression) result.getProjects().get(0); - assertTrue(annotation.getOriginal().toString().contains("painless")); - assertTrue(annotation.getViableBackends().contains(MockLuceneBackend.NAME)); - assertFalse(annotation.getViableBackends().contains(MockDataFusionBackend.NAME)); + // DataFusion is child backend, delegates painless to Lucene + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertAnnotation(result.getProjects().get(0), MockLuceneBackend.NAME); } /** Painless on DataFusion WITHOUT delegation → error. */ @@ -117,8 +109,8 @@ public void testPainlessErrorsWithoutDelegation() { // Lucene supports painless but no delegation configured MockLuceneBackend luceneWithPainless = new MockLuceneBackend() { @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } }; @@ -152,8 +144,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } }; @@ -162,7 +154,7 @@ public Set projectCapabilities() { rexBuilder.makeCall(PAINLESS, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertFalse("Field ref should not be annotated", result.getProjects().get(0) instanceof AnnotatedProjectExpression); assertAnnotation(result.getProjects().get(1), MockLuceneBackend.NAME); @@ -182,8 +174,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless", "highlight"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless", "highlight"); } }; @@ -191,7 +183,7 @@ public Set projectCapabilities() { rexBuilder.makeCall(HIGHLIGHT, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertAnnotation(result.getProjects().get(0), MockLuceneBackend.NAME); } @@ -199,7 +191,7 @@ public Set projectCapabilities() { /** Backend does not support CAST → error. */ public void testUnsupportedScalarFunctionErrors() { - // Default MockDataFusionBackend has empty projectCapabilities + // Default MockDataFusionBackend has empty supportedScalarFunctions RexNode castExpr = rexBuilder.makeCast( typeFactory.createSqlType(SqlTypeName.VARCHAR), rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) @@ -217,7 +209,7 @@ public void testUnsupportedScalarFunctionErrors() { IllegalStateException exception = expectThrows(IllegalStateException.class, () -> runPlanner(project, context)); - assertTrue(exception.getMessage().contains("No backend supports scalar function")); + assertTrue(exception.getMessage().contains("does not support scalar function")); } // ---- Opaque operation natively supported ---- @@ -226,10 +218,12 @@ public void testUnsupportedScalarFunctionErrors() { public void testOpaqueOperationSupportedNatively() { MockDataFusionBackend dfWithPainless = new MockDataFusionBackend() { @Override - public Set projectCapabilities() { - return combine( - scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), EnumSet.allOf(ScalarFunction.class)), - opaqueCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), "painless")); + public Set supportedScalarFunctions() { + return EnumSet.allOf(ScalarFunction.class); + } + @Override + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } }; @@ -237,7 +231,7 @@ public Set projectCapabilities() { rexBuilder.makeCall(PAINLESS, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); } @@ -255,7 +249,7 @@ public void testNestedScalarFunctions() { rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) ); OpenSearchProject result = runProject(plusExpr); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); } @@ -269,9 +263,8 @@ public Set supportedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), - EnumSet.allOf(ScalarFunction.class)); + public Set supportedScalarFunctions() { + return EnumSet.allOf(ScalarFunction.class); } }; MockLuceneBackend luceneAccepting = new MockLuceneBackend() { @@ -280,8 +273,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } }; @@ -295,7 +288,7 @@ public Set projectCapabilities() { OpenSearchProject result = runProject("parquet", List.of(dfWithDelegation, luceneAccepting), fieldRef, painlessExpr, castExpr); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); // field ref — no annotation assertFalse("Field ref should not be annotated", result.getProjects().get(0) instanceof AnnotatedProjectExpression); @@ -313,9 +306,8 @@ public Set supportedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), - EnumSet.allOf(ScalarFunction.class)); + public Set supportedScalarFunctions() { + return EnumSet.allOf(ScalarFunction.class); } }; MockLuceneBackend luceneAccepting = new MockLuceneBackend() { @@ -324,8 +316,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } }; @@ -339,7 +331,7 @@ public Set projectCapabilities() { OpenSearchProject result = runProject("parquet", List.of(dfWithDelegation, luceneAccepting), plusExpr); - assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertEquals(MockDataFusionBackend.NAME, result.getBackend()); // outer PLUS annotated with DF assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); // CAST(painless($0)) — CAST annotated with DF, painless inside annotated with Lucene @@ -372,8 +364,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless", "highlight"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless", "highlight"); } }; @@ -394,8 +386,8 @@ public Set projectCapabilities() { MockLuceneBackend thirdBackend = new MockLuceneBackend() { @Override public String name() { return "mock-third"; } @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "suggest"); + public Set supportedOpaqueProjectOperations() { + return Set.of("suggest"); } }; @@ -418,8 +410,8 @@ public Set supportedDelegations() { }; MockLuceneBackend luceneWithPainlessButNoAccept = new MockLuceneBackend() { @Override - public Set projectCapabilities() { - return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); + public Set supportedOpaqueProjectOperations() { + return Set.of("painless"); } // acceptedDelegations() returns empty by default }; @@ -445,7 +437,7 @@ public Set projectCapabilities() { private void assertAnnotation(RexNode expr, String expectedBackend) { assertTrue("Expected AnnotatedProjectExpression, got " + expr.getClass().getSimpleName(), expr instanceof AnnotatedProjectExpression); - assertTrue(((AnnotatedProjectExpression) expr).getViableBackends().contains(expectedBackend)); + assertEquals(expectedBackend, ((AnnotatedProjectExpression) expr).getBackend()); } private OpenSearchProject runProject(RexNode... exprs) { @@ -480,37 +472,9 @@ private OpenSearchProject runProject(String format, List projectCapabilities() { - return scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), - EnumSet.allOf(ScalarFunction.class)); + public Set supportedScalarFunctions() { + return EnumSet.allOf(ScalarFunction.class); } }; } - - @SafeVarargs - private static Set combine(Set... sets) { - Set result = new HashSet<>(); - for (Set set : sets) { - result.addAll(set); - } - return result; - } - - private static Set scalarCaps(Set formats, Set functions) { - Set caps = new HashSet<>(); - for (ScalarFunction func : functions) { - for (FieldType type : FieldType.values()) { - caps.add(new ProjectCapability.Scalar(func, type, formats)); - } - } - return caps; - } - - private static Set opaqueCaps(Set formats, String... names) { - Set caps = new HashSet<>(); - for (String name : names) { - caps.add(new ProjectCapability.Opaque(name, formats)); - } - return caps; - } } From b671daf291f0d8677793b4b739355c8349a79f77 Mon Sep 17 00:00:00 2001 From: expani Date: Wed, 1 Apr 2026 20:05:48 -0700 Subject: [PATCH 4/6] Second iteration with refactoring and updating tests Signed-off-by: expani --- .../analytics/spi/AggregateFunction.java | 59 ++++-- .../spi/AnalyticsSearchBackendPlugin.java | 25 +-- .../analytics/spi/FieldTypeFamily.java | 60 ------ .../analytics/spi/FilterCapability.java | 35 +++- .../analytics/spi/FilterOperator.java | 82 ++++++-- .../analytics/spi/FullTextOperator.java | 32 --- .../analytics/exec/DefaultPlanExecutor.java | 5 +- .../planner/CapabilityResolutionUtils.java | 176 ++++------------- .../analytics/planner/FieldStorageInfo.java | 43 ++-- .../planner/FieldStorageResolver.java | 147 +++++++++++--- .../analytics/planner/PlannerContext.java | 23 +-- .../analytics/planner/RelNodeUtils.java | 10 +- .../rel/AnnotatedProjectExpression.java | 14 +- .../planner/rel/FullTextFunctions.java | 26 +-- .../planner/rel/OpenSearchAggregate.java | 16 +- .../rel/OpenSearchDistributionTraitDef.java | 28 +-- .../rel/OpenSearchExchangeReducer.java | 20 +- .../planner/rel/OpenSearchExchangeWriter.java | 27 +-- .../planner/rel/OpenSearchFilter.java | 13 +- .../planner/rel/OpenSearchProject.java | 44 +++-- .../planner/rel/OpenSearchRelNode.java | 12 +- .../planner/rel/OpenSearchShuffleReader.java | 20 +- .../analytics/planner/rel/OpenSearchSort.java | 13 +- .../planner/rel/OpenSearchTableScan.java | 17 +- .../rules/OpenSearchAggregateRule.java | 106 ++++++---- .../rules/OpenSearchAggregateSplitRule.java | 4 - .../planner/rules/OpenSearchFilterRule.java | 184 ++++++++---------- .../planner/rules/OpenSearchProjectRule.java | 157 +++++++++++---- .../planner/rules/OpenSearchSortRule.java | 17 +- .../rules/OpenSearchTableScanRule.java | 54 ++++- .../analytics/planner/AggregateRuleTests.java | 56 ++++-- .../planner/BasePlannerRulesTests.java | 2 +- .../analytics/planner/FilterRuleTests.java | 67 +++++-- .../planner/MockDataFusionBackend.java | 104 +++++----- .../analytics/planner/MockLuceneBackend.java | 103 ++++++---- .../analytics/planner/ProjectRuleTests.java | 124 +++++++----- 36 files changed, 1050 insertions(+), 875 deletions(-) delete mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java delete mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java index 84cce5710b1b2..b72e794e93684 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java @@ -11,40 +11,57 @@ import org.apache.calcite.sql.SqlKind; /** - * Aggregate functions that a backend may support. - * Used by the aggregate rule to verify the backend can handle - * every {@link org.apache.calcite.rel.core.AggregateCall} in the plan. + * Aggregate functions that a backend may support, categorized by {@link Type}. * *

Note: {@code COUNT} covers both {@code COUNT(*)} and {@code COUNT(DISTINCT x)}. * The distinction is on {@code AggregateCall.isDistinct()}, not on SqlKind. - * Backends that only support non-distinct count should check distinctness - * separately during fragment conversion. * * @opensearch.internal */ public enum AggregateFunction { - SUM(SqlKind.SUM), - SUM0(SqlKind.SUM0), - MIN(SqlKind.MIN), - MAX(SqlKind.MAX), - COUNT(SqlKind.COUNT), - APPROX_COUNT_DISTINCT(SqlKind.OTHER), - AVG(SqlKind.AVG), - STDDEV_POP(SqlKind.STDDEV_POP), - STDDEV_SAMP(SqlKind.STDDEV_SAMP), - VAR_POP(SqlKind.VAR_POP), - VAR_SAMP(SqlKind.VAR_SAMP), - PERCENTILE_CONT(SqlKind.PERCENTILE_CONT), - PERCENTILE_DISC(SqlKind.PERCENTILE_DISC), - COLLECT(SqlKind.COLLECT), - LISTAGG(SqlKind.LISTAGG); + // Simple — fixed-size state per key + SUM(Type.SIMPLE, SqlKind.SUM), + SUM0(Type.SIMPLE, SqlKind.SUM0), + MIN(Type.SIMPLE, SqlKind.MIN), + MAX(Type.SIMPLE, SqlKind.MAX), + COUNT(Type.SIMPLE, SqlKind.COUNT), + AVG(Type.SIMPLE, SqlKind.AVG), + // Statistical — fixed-size state, multi-pass or running stats + STDDEV_POP(Type.STATISTICAL, SqlKind.STDDEV_POP), + STDDEV_SAMP(Type.STATISTICAL, SqlKind.STDDEV_SAMP), + VAR_POP(Type.STATISTICAL, SqlKind.VAR_POP), + VAR_SAMP(Type.STATISTICAL, SqlKind.VAR_SAMP), + + // State-expanding — state grows with input rows per key + PERCENTILE_CONT(Type.STATE_EXPANDING, SqlKind.PERCENTILE_CONT), + PERCENTILE_DISC(Type.STATE_EXPANDING, SqlKind.PERCENTILE_DISC), + COLLECT(Type.STATE_EXPANDING, SqlKind.COLLECT), + LISTAGG(Type.STATE_EXPANDING, SqlKind.LISTAGG), + + // Approximate — probabilistic, fixed-size state + APPROX_COUNT_DISTINCT(Type.APPROXIMATE, SqlKind.OTHER); + + /** Category of aggregate function. Affects execution strategy (shuffle vs map-reduce). */ + public enum Type { + SIMPLE, + STATISTICAL, + STATE_EXPANDING, + APPROXIMATE + } + + private final Type type; private final SqlKind sqlKind; - AggregateFunction(SqlKind sqlKind) { + AggregateFunction(Type type, SqlKind sqlKind) { + this.type = type; this.sqlKind = sqlKind; } + public Type getType() { + return type; + } + public SqlKind getSqlKind() { return sqlKind; } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index 827ee1e527d96..ef7e5cbf00332 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java @@ -36,13 +36,8 @@ public interface AnalyticsSearchBackendPlugin extends SearchExecEngineProvider { /** Returns the data formats supported by this backend. */ List getSupportedFormats(); - /** Filter operators this backend can evaluate, scoped by field type family. */ - default Set supportedFilterCapabilities() { - return Collections.emptySet(); - } - - /** Full-text operators this backend can evaluate on indexed fields. */ - default Set supportedFullTextOperators() { + /** Filter capabilities scoped to operator, field type, and data format. */ + default Set filterCapabilities() { return Collections.emptySet(); } @@ -61,22 +56,18 @@ default Set acceptedDelegations() { return Collections.emptySet(); } - /** Aggregate functions this backend can evaluate (SUM, AVG, COUNT, etc.). */ - default Set supportedAggregateFunctions() { + /** Aggregate capabilities scoped to function, field type, and data format. */ + default Set aggregateCapabilities() { return Collections.emptySet(); } - /** Scalar functions this backend can evaluate in projections (UPPER, CAST, math ops, etc.). */ - default Set supportedScalarFunctions() { + /** Window capabilities scoped to function, field type, and data format. */ + default Set windowCapabilities() { return Collections.emptySet(); } - /** - * Names of opaque project operations this backend can handle (e.g. "painless", "highlight", "suggest"). - * Used to resolve UnresolvedRexNode from frontend plugins to a backend. - * Analytics Core is agnostic to what these names mean — backends define and consume them. - */ - default Set supportedOpaqueProjectOperations() { + /** Project capabilities: scalar functions and opaque operations, scoped to data format. */ + default Set projectCapabilities() { return Collections.emptySet(); } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java deleted file mode 100644 index 9830f3377419d..0000000000000 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldTypeFamily.java +++ /dev/null @@ -1,60 +0,0 @@ -/* - * 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; - -/** - * Logical field type families for capability matching. - * Groups OpenSearch {@code MappedFieldType.typeName()} strings into - * coarse categories that backends declare support for. - * - * @opensearch.internal - */ -public enum FieldTypeFamily { - NUMERIC, - KEYWORD, - TEXT, - DATE, - BOOLEAN, - IP, - GEO_POINT, - GEO_SHAPE, - BINARY, - NESTED, - OBJECT, - RANGE, - COMPLETION; - - /** - * Maps an OpenSearch mapping type string ({@code MappedFieldType.typeName()}) - * to a FieldTypeFamily. Returns null if the type is not recognized. - */ - public static FieldTypeFamily fromMappingType(String mappingType) { - if (mappingType == null) { - return null; - } - return switch (mappingType) { - case "integer", "long", "short", "byte", "float", "double", - "half_float", "scaled_float", "unsigned_long" -> NUMERIC; - case "keyword", "constant_keyword", "wildcard" -> KEYWORD; - case "text", "match_only_text" -> TEXT; - case "date", "date_nanos" -> DATE; - case "boolean" -> BOOLEAN; - case "ip" -> IP; - case "geo_point", "point" -> GEO_POINT; - case "geo_shape", "shape" -> GEO_SHAPE; - case "binary" -> BINARY; - case "nested" -> NESTED; - case "object", "flat_object" -> OBJECT; - case "integer_range", "float_range", "long_range", - "double_range", "date_range", "ip_range" -> RANGE; - case "completion" -> COMPLETION; - default -> null; - }; - } -} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java index 3fdfcb3f76e11..3f7e8f05096a8 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterCapability.java @@ -8,18 +8,37 @@ package org.opensearch.analytics.spi; +import java.util.Set; + /** - * Declares that a backend can evaluate a specific {@link FilterOperator} - * on a specific {@link FieldTypeFamily}. - * - *

Example: {@code FilterCapability.of(EQUALS, KEYWORD)} means the backend - * can evaluate equality predicates on keyword fields. + * Declares a backend's ability to evaluate filter predicates, scoped to data formats. + * Three variants for the three categories of filter operations. * * @opensearch.internal */ -public record FilterCapability(FilterOperator operator, FieldTypeFamily fieldTypeFamily) { +public sealed interface FilterCapability { + + /** Standard comparison filter (EQUALS, GT, IN, LIKE, etc.) on a field type in given formats. */ + record Standard(FilterOperator operator, FieldType fieldType, + Set formats) implements FilterCapability { + public Standard { + formats = Set.copyOf(formats); + } + } + + /** Full-text filter (MATCH, MATCH_PHRASE, FUZZY, etc.) with supported query parameters. */ + record FullText(FilterOperator operator, FieldType fieldType, + Set formats, Set supportedParams) implements FilterCapability { + public FullText { + formats = Set.copyOf(formats); + supportedParams = Set.copyOf(supportedParams); + } + } - public static FilterCapability of(FilterOperator operator, FieldTypeFamily fieldTypeFamily) { - return new FilterCapability(operator, fieldTypeFamily); + /** Expression-based filter on derived columns (e.g., HAVING after aggregate). */ + record Expression(Set formats) implements FilterCapability { + public Expression { + formats = Set.copyOf(formats); + } } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java index d60be4445bb0d..70b3262f9b8b1 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FilterOperator.java @@ -11,35 +11,83 @@ import org.apache.calcite.sql.SqlKind; /** - * Standard comparison/predicate operations that a backend may support. + * All filter operations a backend may support, covering standard comparisons, + * full-text search, and expression-based filtering. + * + *

Each operator carries a {@link Type} indicating its category and whether + * it supports parameters (e.g., full-text operators accept analyzer, slop, etc.). * * @opensearch.internal */ public enum FilterOperator { - EQUALS(SqlKind.EQUALS), - NOT_EQUALS(SqlKind.NOT_EQUALS), - GREATER_THAN(SqlKind.GREATER_THAN), - GREATER_THAN_OR_EQUAL(SqlKind.GREATER_THAN_OR_EQUAL), - LESS_THAN(SqlKind.LESS_THAN), - LESS_THAN_OR_EQUAL(SqlKind.LESS_THAN_OR_EQUAL), - IS_NULL(SqlKind.IS_NULL), - IS_NOT_NULL(SqlKind.IS_NOT_NULL), - IN(SqlKind.IN), - LIKE(SqlKind.LIKE), - PREFIX(SqlKind.OTHER), - REGEXP(SqlKind.OTHER), - WILDCARD(SqlKind.OTHER); + // Standard comparison + EQUALS(Type.STANDARD, SqlKind.EQUALS), + NOT_EQUALS(Type.STANDARD, SqlKind.NOT_EQUALS), + GREATER_THAN(Type.STANDARD, SqlKind.GREATER_THAN), + GREATER_THAN_OR_EQUAL(Type.STANDARD, SqlKind.GREATER_THAN_OR_EQUAL), + LESS_THAN(Type.STANDARD, SqlKind.LESS_THAN), + LESS_THAN_OR_EQUAL(Type.STANDARD, SqlKind.LESS_THAN_OR_EQUAL), + IS_NULL(Type.STANDARD, SqlKind.IS_NULL), + IS_NOT_NULL(Type.STANDARD, SqlKind.IS_NOT_NULL), + IN(Type.STANDARD, SqlKind.IN), + LIKE(Type.STANDARD, SqlKind.LIKE), + PREFIX(Type.STANDARD, SqlKind.OTHER), + + // Full-text search + MATCH(Type.FULL_TEXT, SqlKind.OTHER), + MATCH_PHRASE(Type.FULL_TEXT, SqlKind.OTHER), + MATCH_PHRASE_PREFIX(Type.FULL_TEXT, SqlKind.OTHER), + MATCH_BOOL_PREFIX(Type.FULL_TEXT, SqlKind.OTHER), + MULTI_MATCH(Type.FULL_TEXT, SqlKind.OTHER), + QUERY_STRING(Type.FULL_TEXT, SqlKind.OTHER), + SIMPLE_QUERY_STRING(Type.FULL_TEXT, SqlKind.OTHER), + FUZZY(Type.FULL_TEXT, SqlKind.OTHER), + WILDCARD(Type.FULL_TEXT, SqlKind.OTHER), + REGEXP(Type.FULL_TEXT, SqlKind.OTHER), + + // Expression-based filtering (on derived columns, e.g., HAVING) + EXPRESSION(Type.EXPRESSION, SqlKind.OTHER); + + /** + * Category of filter operator. Declares whether the operator supports parameters. + */ + public enum Type { + STANDARD(false), + FULL_TEXT(true), + EXPRESSION(false); + + private final boolean supportsParams; + + Type(boolean supportsParams) { + this.supportsParams = supportsParams; + } + + public boolean supportsParams() { + return supportsParams; + } + } + + private final Type type; private final SqlKind sqlKind; - FilterOperator(SqlKind sqlKind) { + FilterOperator(Type type, SqlKind sqlKind) { + this.type = type; this.sqlKind = sqlKind; } - /** Maps a Calcite SqlKind to a FilterOperator, or null if not a standard filter op. */ + public Type getType() { + return type; + } + + public SqlKind getSqlKind() { + return sqlKind; + } + + /** Maps a Calcite SqlKind to a standard FilterOperator, or null if not recognized. */ public static FilterOperator fromSqlKind(SqlKind kind) { for (FilterOperator op : values()) { - if (op.sqlKind == kind && op.sqlKind != SqlKind.OTHER) { + if (op.type == Type.STANDARD && op.sqlKind == kind && op.sqlKind != SqlKind.OTHER) { return op; } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java deleted file mode 100644 index bb9068134523b..0000000000000 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FullTextOperator.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * 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; - -/** - * Full-text search operations that an inverted index backend may support. - * Intended as a common facade between Lucene, Tantivy, and future - * full-text backends. - * - * @opensearch.internal - */ -public enum FullTextOperator { - MATCH, - MATCH_PHRASE, - MATCH_PHRASE_PREFIX, - MATCH_BOOL_PREFIX, - MULTI_MATCH, - QUERY_STRING, - SIMPLE_QUERY_STRING, - FUZZY, - SPAN_NEAR, - SPAN_OR, - SPAN_NOT, - SPAN_FIRST, - SPAN_TERM -} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java index 081370f3b2875..e1d42f9f5a645 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java @@ -17,6 +17,7 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; +import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.PlannerImpl; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; @@ -65,7 +66,9 @@ public DefaultPlanExecutor(List providers, Indices @Override public Iterable execute(RelNode logicalFragment, Object context) { logicalFragment = PlannerImpl.createPlan(logicalFragment, - new PlannerContext(new ArrayList<>(backEnds.values()), clusterService.state())); + new PlannerContext( + new CapabilityRegistry(new ArrayList<>(backEnds.values())), + clusterService.state())); String tableName = extractTableName(logicalFragment); AnalyticsSearchBackendPlugin provider = selectBackEnd(); if (provider == null) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java index d968cf9ac4fd3..5b45a75d2fad5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java @@ -8,22 +8,16 @@ package org.opensearch.analytics.planner; -import org.apache.calcite.rel.RelDistribution; -import org.apache.calcite.rel.core.Aggregate; import org.opensearch.analytics.planner.rel.ShuffleImpl; -import org.opensearch.analytics.spi.AggregateFunction; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; -import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.analytics.spi.ShuffleCapability; import java.util.ArrayList; import java.util.List; -import java.util.Map; /** - * Shared helpers for resolving backend capabilities during planning. - * Used by all marking rules to determine which backend handles each operator. + * Utility logic that operates on {@link CapabilityRegistry} results. + * Registry handles indexed lookups; this class handles filtering and resolution logic. * * @opensearch.internal */ @@ -31,148 +25,58 @@ public final class CapabilityResolutionUtils { private CapabilityResolutionUtils() {} - /** True if the named backend supports the given operator capability. */ - public static boolean backendSupports(Map backends, - String backendName, OperatorCapability capability) { - AnalyticsSearchBackendPlugin plugin = backends.get(backendName); - return plugin != null && plugin.supportedOperators().contains(capability); - } - - /** Finds the first backend that supports the given capability, or null. */ - public static String findBackendWith(Map backends, - OperatorCapability capability) { - for (AnalyticsSearchBackendPlugin plugin : backends.values()) { - if (plugin.supportedOperators().contains(capability)) { - return plugin.name(); + /** Filters viable backends to those that support COORDINATOR_REDUCE. */ + public static List filterByReduceCapability(CapabilityRegistry registry, + List viableBackends) { + List reduceCapable = registry.operatorBackends(OperatorCapability.COORDINATOR_REDUCE); + List result = new ArrayList<>(); + for (String name : viableBackends) { + if (reduceCapable.contains(name)) { + result.add(name); } } - return null; - } - - /** - * Resolves the backend for an operator: prefers childBackend if it supports - * the capability, otherwise finds any backend that does. - * Falls back to childBackend if none found. - */ - public static String resolveBackend(Map backends, - String childBackend, OperatorCapability capability) { - if (backendSupports(backends, childBackend, capability)) { - return childBackend; - } - String found = findBackendWith(backends, capability); - return found != null ? found : childBackend; - } - - /** - * True if the named backend supports AGGREGATE and every aggregate function - * required by the plan. - */ - public static boolean backendSupportsAggregate(Map backends, - String backendName, Aggregate aggregate) { - AnalyticsSearchBackendPlugin plugin = backends.get(backendName); - if (plugin == null || !plugin.supportedOperators().contains(OperatorCapability.AGGREGATE)) { - return false; + if (result.isEmpty()) { + throw new IllegalStateException( + "No viable backend supports COORDINATOR_REDUCE among " + viableBackends); } - return aggregate.getAggCallList().stream().allMatch(aggCall -> { - AggregateFunction func = AggregateFunction.fromSqlKind(aggCall.getAggregation().getKind()); - if (func == null) { - func = AggregateFunction.fromNameOrError(aggCall.getAggregation().getName()); - } - return plugin.supportedAggregateFunctions().contains(func); - }); + return result; } - /** Finds the first backend that supports AGGREGATE and all required functions. */ - public static String findBackendForAggregate(Map backends, - Aggregate aggregate) { - for (AnalyticsSearchBackendPlugin plugin : backends.values()) { - if (backendSupportsAggregate(backends, plugin.name(), aggregate)) { - return plugin.name(); + /** Filters viable backends to those with any shuffle capability. */ + public static List filterByShuffleCapability(CapabilityRegistry registry, + List viableBackends) { + List result = new ArrayList<>(); + for (String name : viableBackends) { + if (!registry.getShuffleCapabilities(name).isEmpty()) { + result.add(name); } } - return null; - } - - /** - * Computes all backends viable for an operator capability, considering delegation. - * A backend is viable if it supports the capability natively, OR if it can delegate - * that type of work and at least one other backend can accept the delegation. - */ - public static List computeViableBackends(Map backends, - OperatorCapability capability, - DelegationType delegationType) { - boolean anyAcceptsDelegation = backends.values().stream() - .anyMatch(b -> b.acceptedDelegations().contains(delegationType)); - - List viable = new ArrayList<>(); - for (AnalyticsSearchBackendPlugin plugin : backends.values()) { - if (plugin.supportedOperators().contains(capability)) { - viable.add(plugin.name()); - } else if (anyAcceptsDelegation && plugin.supportedDelegations().contains(delegationType)) { - viable.add(plugin.name()); - } + if (result.isEmpty()) { + throw new IllegalStateException( + "No viable backend supports shuffle among " + viableBackends); } - return viable; + return result; } - /** Computes all backends that natively support the given capability. No delegation. */ - public static List computeViableBackends(Map backends, - OperatorCapability capability) { - List viable = new ArrayList<>(); - for (AnalyticsSearchBackendPlugin plugin : backends.values()) { - if (plugin.supportedOperators().contains(capability)) { - viable.add(plugin.name()); + /** Picks the best shuffle impl across all shuffle-viable backends. Prefers STREAM over FILE. */ + public static ShuffleImpl bestShuffleImpl(CapabilityRegistry registry, List shuffleViable) { + boolean hasStream = false; + boolean hasFile = false; + for (String name : shuffleViable) { + var caps = registry.getShuffleCapabilities(name); + if (caps.contains(ShuffleCapability.STREAM_WRITE)) { + hasStream = true; } - } - return viable; - } - - /** - * Computes all backends viable for a capability that also support the given data format. - * Used by scan rule where the backend must match the index's primary data format. - */ - public static List computeViableBackends(Map backends, - OperatorCapability capability, String dataFormat) { - List viable = new ArrayList<>(); - for (AnalyticsSearchBackendPlugin plugin : backends.values()) { - if (plugin.supportedOperators().contains(capability) - && plugin.getSupportedFormats().stream().anyMatch(f -> f.name().equals(dataFormat))) { - viable.add(plugin.name()); + if (caps.contains(ShuffleCapability.FILE_WRITE)) { + hasFile = true; } } - return viable; - } - - /** - * Resolves the shuffle implementation for a HASH/RANGE exchange based on the - * backend's advertised shuffle capabilities. Prefers STREAM over FILE. - * Only called for non-SINGLETON distributions. - */ - public static ShuffleImpl resolveShuffleImpl(Map backends, - String backendName, RelDistribution.Type distributionType) { - AnalyticsSearchBackendPlugin plugin = backends.get(backendName); - if (plugin != null) { - if (plugin.supportedShuffleCapabilities().contains(ShuffleCapability.STREAM_WRITE)) { - return ShuffleImpl.STREAM; - } - if (plugin.supportedShuffleCapabilities().contains(ShuffleCapability.FILE_WRITE)) { - return ShuffleImpl.FILE; - } + if (hasStream) { + return ShuffleImpl.STREAM; } - throw new IllegalStateException( - "Backend [" + backendName + "] does not advertise any shuffle capability for distribution [" - + distributionType + "]"); - } - - /** - * Validates that the backend supports coordinator-side reduce. - * Throws if it doesn't. - */ - public static void validateReduceCapability(Map backends, - String backendName) { - if (!backendSupports(backends, backendName, OperatorCapability.COORDINATOR_REDUCE)) { - throw new IllegalStateException( - "Backend [" + backendName + "] does not support COORDINATOR_REDUCE capability"); + if (hasFile) { + return ShuffleImpl.FILE; } + throw new IllegalStateException("No shuffle impl available among " + shuffleViable); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java index fc9d6a16323a7..1e8efa39d264a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageInfo.java @@ -8,45 +8,57 @@ package org.opensearch.analytics.planner; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.analytics.spi.FieldType; + import java.util.List; /** - * Per-column storage metadata describing where doc values and indices live. + * Per-column storage metadata describing where doc values, indices and stored fields live. * Flows through the plan tree: each OpenSearchRelNode computes this for its * output columns from its input's metadata. * - *

TODO: use {@code DataFormat} instead of String for format identifiers - * once the dependency is wired through. - * * @opensearch.internal */ public class FieldStorageInfo { private final String fieldName; - private final String fieldType; + private final String mappingType; + private final FieldType fieldType; private final List docValueFormats; private final List indexFormats; + private final List storedFieldFormats; private final boolean derived; - public FieldStorageInfo(String fieldName, String fieldType, List docValueFormats, - List indexFormats, boolean derived) { + public FieldStorageInfo(String fieldName, String mappingType, FieldType fieldType, + List docValueFormats, List indexFormats, + List storedFieldFormats, boolean derived) { this.fieldName = fieldName; + this.mappingType = mappingType; this.fieldType = fieldType; this.docValueFormats = docValueFormats; this.indexFormats = indexFormats; + this.storedFieldFormats = storedFieldFormats; this.derived = derived; } - /** Creates a derived column (agg result, expression) with no physical storage. */ - public static FieldStorageInfo derivedColumn(String fieldName, String fieldType) { - return new FieldStorageInfo(fieldName, fieldType, List.of(), List.of(), true); + /** Creates a derived column (agg result, expression) with no physical storage. + * FieldType inferred from SqlTypeName. */ + public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTypeName) { + return new FieldStorageInfo(fieldName, sqlTypeName.getName(), + FieldType.fromSqlTypeName(sqlTypeName), + List.of(), List.of(), List.of(), true); } public String getFieldName() { return fieldName; } - public String getFieldType() { + public String getMappingType() { + return mappingType; + } + + public FieldType getFieldType() { return fieldType; } @@ -60,6 +72,11 @@ public List getIndexFormats() { return indexFormats; } + /** Data formats holding stored fields for this field (e.g. ["parquet"], ["lucene"]). */ + public List getStoredFieldFormats() { + return storedFieldFormats; + } + /** True for computed columns (agg results, expressions) with no physical storage. */ public boolean isDerived() { return derived; @@ -72,4 +89,8 @@ public boolean hasDocValues() { public boolean hasIndex() { return !indexFormats.isEmpty(); } + + public boolean hasStoredFields() { + return !storedFieldFormats.isEmpty(); + } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java index 40b60b6527bcb..bfef4645ce79b 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java @@ -8,34 +8,41 @@ package org.opensearch.analytics.planner; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.metadata.MappingMetadata; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; +import org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; /** - * Builds {@link FieldStorageInfo} for index fields from {@link IndexMetadata}. + * Resolves per-field storage metadata by consulting each backend's {@link DataFormat} + * capabilities. For each field, determines which formats provide doc values + * ({@link Capability#COLUMNAR_STORAGE}), indices ({@link Capability#FULL_TEXT_SEARCH}, + * {@link Capability#POINT_RANGE}), and stored fields ({@link Capability#STORED_FIELDS}). * - *

At the coordinator, we only have {@link MappingMetadata} (raw source map). - * {@code MappedFieldType} with {@code hasDocValues()} / {@code isSearchable()} is - * only available at the data node via {@code MapperService}. - * - *

TODO: introduce a coordinator-level typed field metadata API so we don't - * parse raw maps here. Or defer field storage resolution to the data node planner - * where {@code MappedFieldType} is available. + *

Production constructor queries backends' {@link DataFormat#supportedFields()} to + * build per-field storage info. Test constructor accepts explicit per-field storage. * * @opensearch.internal */ public class FieldStorageResolver { - private FieldStorageResolver() {} + private final Map fieldStorage; + private final List docValueFormats; + /** + * Production: resolves per-field storage from IndexMetadata and backend capabilities. + */ @SuppressWarnings("unchecked") - public static List resolve(IndexMetadata indexMetadata, List fieldNames) { + public FieldStorageResolver(IndexMetadata indexMetadata, List backends) { String indexName = indexMetadata.getIndex().getName(); - String primaryFormat = indexMetadata.getSettings().get("index.composite.primary_data_format", "lucene"); MappingMetadata mapping = indexMetadata.mapping(); if (mapping == null) { @@ -47,37 +54,123 @@ public static List resolve(IndexMetadata indexMetadata, List result = new ArrayList<>(); + // Build format → capabilities lookup from all backends + Map> formatCapabilities = buildFormatCapabilities(backends); + + this.fieldStorage = new LinkedHashMap<>(); + for (Map.Entry entry : properties.entrySet()) { + String fieldName = entry.getKey(); + Map fieldProps = (Map) entry.getValue(); + String fieldType = (String) fieldProps.get("type"); + if (fieldType == null) { + throw new IllegalStateException("Field [" + fieldName + "] has no type in mapping"); + } + this.fieldStorage.put(fieldName, resolveField(fieldName, fieldType, fieldProps, formatCapabilities)); + } + this.docValueFormats = computeDocValueFormats(this.fieldStorage); + } + + /** + * Test/future: explicit per-field storage info. + * Simulates hybrid indices where doc values exist in multiple formats. + */ + public FieldStorageResolver(Map fieldStorage) { + this.fieldStorage = fieldStorage; + this.docValueFormats = computeDocValueFormats(fieldStorage); + } + + /** Resolves storage info for the requested fields. */ + public List resolve(List fieldNames) { + List result = new ArrayList<>(fieldNames.size()); for (String fieldName : fieldNames) { - Map fieldProps = (Map) properties.get(fieldName); - if (fieldProps == null) { - throw new IllegalStateException("Field [" + fieldName + "] not found in mapping for index [" + indexName + "]"); + FieldStorageInfo info = fieldStorage.get(fieldName); + if (info == null) { + throw new IllegalStateException("Field [" + fieldName + "] not found in field storage"); + } + result.add(info); + } + return result; + } + + /** Returns all unique data formats that hold doc values across all fields. Precomputed at creation. */ + public List docValueFormats() { + return docValueFormats; + } + + private static List computeDocValueFormats(Map fieldStorage) { + List formats = new ArrayList<>(); + for (FieldStorageInfo info : fieldStorage.values()) { + for (String format : info.getDocValueFormats()) { + if (!formats.contains(format)) { + formats.add(format); + } + } + } + return formats; + } + + /** + * Builds a lookup: formatName → fieldType → FieldTypeCapabilities + * from all backends' DataFormats. + */ + private static Map> buildFormatCapabilities( + List backends) { + Map> result = new LinkedHashMap<>(); + for (AnalyticsSearchBackendPlugin backend : backends) { + for (DataFormat format : backend.getSupportedFormats()) { + Map byFieldType = result.computeIfAbsent( + format.name(), k -> new LinkedHashMap<>()); + for (FieldTypeCapabilities cap : format.supportedFields()) { + byFieldType.put(cap.fieldType(), cap); + } } - result.add(resolveField(fieldName, fieldProps, primaryFormat)); } return result; } - private static FieldStorageInfo resolveField(String fieldName, Map fieldProps, String primaryFormat) { - String fieldType = (String) fieldProps.get("type"); - if (fieldType == null) { - throw new IllegalStateException("Field [" + fieldName + "] has no type in mapping"); + private static FieldStorageInfo resolveField(String fieldName, String fieldType, + Map fieldProps, + Map> formatCapabilities) { + List docValueFormats = new ArrayList<>(); + List indexFormats = new ArrayList<>(); + List storedFieldFormats = new ArrayList<>(); + + for (Map.Entry> formatEntry : formatCapabilities.entrySet()) { + String formatName = formatEntry.getKey(); + FieldTypeCapabilities caps = formatEntry.getValue().get(fieldType); + if (caps == null) { + continue; + } + if (caps.capabilities().contains(Capability.COLUMNAR_STORAGE)) { + docValueFormats.add(formatName); + } + if (caps.capabilities().contains(Capability.FULL_TEXT_SEARCH) + || caps.capabilities().contains(Capability.POINT_RANGE)) { + indexFormats.add(formatName); + } + if (caps.capabilities().contains(Capability.STORED_FIELDS)) { + storedFieldFormats.add(formatName); + } } - // TODO: use MappedFieldType.hasDocValues() / isSearchable() when available + // Respect mapping overrides: doc_values=false or index=false boolean hasDocValues = Boolean.TRUE.equals(fieldProps.get("doc_values")) || (fieldProps.get("doc_values") == null && !"text".equals(fieldType)); boolean isIndexed = Boolean.TRUE.equals(fieldProps.get("index")) || fieldProps.get("index") == null; - if (!hasDocValues && !isIndexed) { - throw new IllegalStateException("Field [" + fieldName + "] has neither doc_values nor index"); + if (!hasDocValues) { + docValueFormats = List.of(); + } + if (!isIndexed) { + indexFormats = List.of(); } - // TODO: data format per field should come from MappingMetadata directly - List docValueFormats = hasDocValues ? List.of(primaryFormat) : List.of(); - List indexFormats = isIndexed ? List.of("lucene") : List.of(); + if (docValueFormats.isEmpty() && indexFormats.isEmpty() && storedFieldFormats.isEmpty()) { + throw new IllegalStateException("Field [" + fieldName + "] has no storage in any format"); + } - return new FieldStorageInfo(fieldName, fieldType, docValueFormats, indexFormats, false); + return new FieldStorageInfo(fieldName, fieldType, FieldType.fromMappingType(fieldType), + docValueFormats, indexFormats, storedFieldFormats, false); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java index 620e604de7670..ae69b6fb7bdd4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java @@ -9,38 +9,29 @@ package org.opensearch.analytics.planner; import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.cluster.ClusterState; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; - /** * Shared context available to all planner rules. - * Holds cluster state and backend plugins. Rules consult this - * to extract index metadata, mappings, and backend capabilities - * for whichever tables they encounter. + * Holds capability registry (singleton, built at plugin startup) and + * per-query cluster state. * * @opensearch.internal */ public class PlannerContext { - private final Map backends; + private final CapabilityRegistry capabilityRegistry; private final ClusterState clusterState; private final OpenSearchDistributionTraitDef distributionTraitDef; - public PlannerContext(List backends, ClusterState clusterState) { - this.backends = new LinkedHashMap<>(); - for (AnalyticsSearchBackendPlugin b : backends) { - this.backends.put(b.name(), b); - } + public PlannerContext(CapabilityRegistry capabilityRegistry, ClusterState clusterState) { + this.capabilityRegistry = capabilityRegistry; this.clusterState = clusterState; this.distributionTraitDef = new OpenSearchDistributionTraitDef(this); } - public Map getBackends() { - return backends; + public CapabilityRegistry getCapabilityRegistry() { + return capabilityRegistry; } public ClusterState getClusterState() { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java index 098e9dcea8402..6418f2c163fd6 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java @@ -57,20 +57,20 @@ public static RelNode copyToCluster(RelNode node, RelOptCluster newCluster, if (node instanceof OpenSearchTableScan scan) { return new OpenSearchTableScan(newCluster, newTraits, scan.getTable(), - scan.getBackend(), scan.getViableBackends(), scan.getOutputFieldStorage()); + scan.getViableBackends(), scan.getOutputFieldStorage()); } else if (node instanceof OpenSearchFilter filter) { return new OpenSearchFilter(newCluster, newTraits, newInputs.getFirst(), - filter.getCondition(), filter.getBackend(), filter.getViableBackends()); + filter.getCondition(), filter.getViableBackends()); } else if (node instanceof OpenSearchAggregate aggregate) { return new OpenSearchAggregate(newCluster, newTraits, newInputs.getFirst(), aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), - aggregate.getMode(), aggregate.getBackend(), aggregate.getViableBackends()); + aggregate.getMode(), aggregate.getViableBackends()); } else if (node instanceof OpenSearchSort sort) { return new OpenSearchSort(newCluster, newTraits, newInputs.getFirst(), - sort.getCollation(), sort.offset, sort.fetch, sort.getBackend(), sort.getViableBackends()); + sort.getCollation(), sort.offset, sort.fetch, sort.getViableBackends()); } else if (node instanceof OpenSearchProject project) { return new OpenSearchProject(newCluster, newTraits, newInputs.getFirst(), - project.getProjects(), project.getRowType(), project.getBackend()); + project.getProjects(), project.getRowType(), project.getViableBackends()); } throw new UnsupportedOperationException("Cannot copy node type: " + node.getClass().getSimpleName()); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java index 12b96683264cb..03e9851c7a0b5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedProjectExpression.java @@ -46,25 +46,25 @@ public SqlSyntax getSyntax() { }; private final RexNode original; - private final String backend; + private final List viableBackends; - public AnnotatedProjectExpression(RelDataType type, RexNode original, String backend) { + public AnnotatedProjectExpression(RelDataType type, RexNode original, List viableBackends) { super(type, ANNOTATED_PROJECT_EXPR_OP, List.of(original)); this.original = original; - this.backend = backend; + this.viableBackends = viableBackends; } public RexNode getOriginal() { return original; } - /** The backend that evaluates this expression. */ - public String getBackend() { - return backend; + /** Backends that can evaluate this expression. */ + public List getViableBackends() { + return viableBackends; } @Override protected String computeDigest(boolean withType) { - return "ANNOTATED_PROJECT_EXPR(backend=" + backend + ", " + original + ")"; + return "ANNOTATED_PROJECT_EXPR(backends=" + viableBackends + ", " + original + ")"; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java index 713830e10878b..7b2f21a6ac4b4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/FullTextFunctions.java @@ -13,14 +13,12 @@ import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.type.OperandTypes; import org.apache.calcite.sql.type.ReturnTypes; -import org.opensearch.analytics.spi.FullTextOperator; +import org.opensearch.analytics.spi.FilterOperator; /** * Calcite SqlFunction markers for full-text search operations. - * Used by frontend plugins (DSL converter, SQL, PPL) to represent - * full-text queries in the RelNode tree. The filter rule recognizes - * these and routes to backends that support the corresponding - * {@link FullTextOperator}. + * The filter rule recognizes these and routes to backends that support + * the corresponding {@link FilterOperator} of type {@link FilterOperator.Type#FULL_TEXT}. * * @opensearch.internal */ @@ -38,21 +36,15 @@ public class FullTextFunctions { private FullTextFunctions() {} private static SqlFunction fullTextFunction(String name) { - return new SqlFunction( - name, - SqlKind.OTHER_FUNCTION, - ReturnTypes.BOOLEAN, - null, - OperandTypes.ANY, - SqlFunctionCategory.USER_DEFINED_FUNCTION - ); + return new SqlFunction(name, SqlKind.OTHER_FUNCTION, ReturnTypes.BOOLEAN, + null, OperandTypes.ANY, SqlFunctionCategory.USER_DEFINED_FUNCTION); } - /** Maps a SqlFunction to a FullTextOperator, or null if not a full-text function. */ - public static FullTextOperator toFullTextOperator(SqlFunction function) { - String name = function.getName(); + /** Maps a SqlFunction to a FULL_TEXT FilterOperator, or null if not a full-text function. */ + public static FilterOperator toFilterOperator(SqlFunction function) { try { - return FullTextOperator.valueOf(name); + FilterOperator op = FilterOperator.valueOf(function.getName()); + return op.getType() == FilterOperator.Type.FULL_TEXT ? op : null; } catch (IllegalArgumentException ignored) { return null; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java index 78d6eade67fad..df0eb21a54852 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java @@ -26,25 +26,18 @@ */ public class OpenSearchAggregate extends Aggregate implements OpenSearchRelNode { - private final String backend; private final List viableBackends; private final AggregateMode mode; public OpenSearchAggregate(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, ImmutableBitSet groupSet, List groupSets, List aggCalls, AggregateMode mode, - String backend, List viableBackends) { + List viableBackends) { super(cluster, traitSet, List.of(), input, groupSet, groupSets, aggCalls); this.mode = mode; - this.backend = backend; this.viableBackends = viableBackends; } - @Override - public String getBackend() { - return backend; - } - public AggregateMode getMode() { return mode; } @@ -76,7 +69,8 @@ public List getOutputFieldStorage() { // Agg results: derived columns with no physical storage for (AggregateCall aggCall : getAggCallList()) { - outputStorage.add(FieldStorageInfo.derivedColumn(aggCall.getName(), aggCall.getType().toString())); + outputStorage.add(FieldStorageInfo.derivedColumn(aggCall.getName(), + aggCall.getType().getSqlTypeName())); } return outputStorage; @@ -85,7 +79,7 @@ public List getOutputFieldStorage() { @Override public Aggregate copy(RelTraitSet traitSet, RelNode input, ImmutableBitSet groupSet, List groupSets, List aggCalls) { - return new OpenSearchAggregate(getCluster(), traitSet, input, groupSet, groupSets, aggCalls, mode, backend, viableBackends); + return new OpenSearchAggregate(getCluster(), traitSet, input, groupSet, groupSets, aggCalls, mode, viableBackends); } @Override @@ -109,6 +103,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(org.apache.calcite.pla @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("mode", mode).item("backend", backend).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("mode", mode).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java index b36348f0bc274..b381f6c147be3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchDistributionTraitDef.java @@ -15,6 +15,7 @@ import org.apache.calcite.rel.RelNode; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.PlannerContext; @@ -93,31 +94,32 @@ public RelNode convert(RelOptPlanner planner, RelNode rel, return rel; } - String backend = resolveBackendFromRel(rel); + List viableBackends = resolveViableBackendsFromRel(rel); LOGGER.info("convert(): rel={}#{}, fromTrait={}, toTrait={}, backend={}", - rel.getClass().getSimpleName(), rel.getId(), fromTrait, toTrait, backend); + rel.getClass().getSimpleName(), rel.getId(), fromTrait, toTrait, viableBackends.getFirst()); + + CapabilityRegistry registry = plannerContext.getCapabilityRegistry(); RelNode result; if (toTrait.getType() == RelDistribution.Type.SINGLETON) { - CapabilityResolutionUtils.validateReduceCapability( - plannerContext.getBackends(), backend); + List reduceViable = CapabilityResolutionUtils.filterByReduceCapability( + registry, viableBackends); result = new OpenSearchExchangeReducer( rel.getCluster(), rel.getTraitSet().replace(toTrait), rel, - backend + reduceViable ); } else { - // HASH/RANGE: Writer at data node partitions and writes shuffle data. - // Reader at target data node reads from source nodes. - ShuffleImpl shuffleImpl = CapabilityResolutionUtils.resolveShuffleImpl( - plannerContext.getBackends(), backend, toTrait.getType()); + List shuffleViable = CapabilityResolutionUtils.filterByShuffleCapability( + registry, viableBackends); + ShuffleImpl shuffleImpl = CapabilityResolutionUtils.bestShuffleImpl(registry, shuffleViable); OpenSearchExchangeWriter writer = new OpenSearchExchangeWriter( rel.getCluster(), rel.getTraitSet(), rel, - backend, + shuffleViable, shuffleImpl, toTrait.getKeys() ); @@ -125,7 +127,7 @@ public RelNode convert(RelOptPlanner planner, RelNode rel, rel.getCluster(), rel.getTraitSet().replace(toTrait), writer, - backend, + shuffleViable, shuffleImpl ); } @@ -139,12 +141,12 @@ public boolean canConvert(RelOptPlanner planner, OpenSearchDistribution fromTrai return true; } - private static String resolveBackendFromRel(RelNode rel) { + private static List resolveViableBackendsFromRel(RelNode rel) { if (rel instanceof RelSubset subset) { rel = subset.getBestOrOriginal(); } if (rel instanceof OpenSearchRelNode openSearchRel) { - return openSearchRel.getBackend(); + return openSearchRel.getViableBackends(); } throw new IllegalStateException( "Expected OpenSearchRelNode but got [" + rel.getClass().getSimpleName() + "#" + rel.getId() + "]"); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java index 19f3dedc3a8f5..00111a8b5ff5c 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java @@ -32,25 +32,17 @@ */ public class OpenSearchExchangeReducer extends SingleRel implements OpenSearchRelNode { - private final String backend; + private final List viableBackends; public OpenSearchExchangeReducer(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - String backend) { + List viableBackends) { super(cluster, traitSet, input); - this.backend = backend; - } - - @Override - public String getBackend() { - return backend; + this.viableBackends = viableBackends; } @Override public List getViableBackends() { - if (getInput() instanceof OpenSearchRelNode openSearchInput) { - return openSearchInput.getViableBackends(); - } - return List.of(); + return viableBackends; } @Override @@ -63,7 +55,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), backend); + return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), viableBackends); } @Override @@ -73,6 +65,6 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("backend", backend).item("viableBackends", getViableBackends()); + return super.explainTerms(pw).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java index 99489200ac1de..0f8f6fd7f8ab3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeWriter.java @@ -32,24 +32,23 @@ */ public class OpenSearchExchangeWriter extends SingleRel implements OpenSearchRelNode { - private final String backend; - private final ShuffleImpl shuffleImpl; // null for SINGLETON (transport, not shuffle) - private final List keys; // empty for SINGLETON + private final List viableBackends; + private final ShuffleImpl shuffleImpl; + private final List keys; public OpenSearchExchangeWriter(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - String backend, ShuffleImpl shuffleImpl, List keys) { + List viableBackends, ShuffleImpl shuffleImpl, List keys) { super(cluster, traitSet, input); - this.backend = backend; + this.viableBackends = viableBackends; this.shuffleImpl = shuffleImpl; this.keys = keys; } @Override - public String getBackend() { - return backend; + public List getViableBackends() { + return viableBackends; } - /** Null for SINGLETON (uses Analytics Core transport). FILE or STREAM for shuffle. */ public ShuffleImpl getShuffleImpl() { return shuffleImpl; } @@ -62,14 +61,6 @@ public boolean isShuffle() { return shuffleImpl != null; } - @Override - public List getViableBackends() { - if (getInput() instanceof OpenSearchRelNode openSearchInput) { - return openSearchInput.getViableBackends(); - } - return List.of(); - } - @Override public List getOutputFieldStorage() { if (getInput() instanceof OpenSearchRelNode openSearchInput) { @@ -81,7 +72,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { return new OpenSearchExchangeWriter(getCluster(), traitSet, sole(inputs), - backend, shuffleImpl, keys); + viableBackends, shuffleImpl, keys); } @Override @@ -91,7 +82,7 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - RelWriter writer = super.explainTerms(pw).item("backend", backend); + RelWriter writer = super.explainTerms(pw).item("viableBackends", viableBackends); if (shuffleImpl != null) { writer.item("shuffleImpl", shuffleImpl).item("keys", keys); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java index 45fe5cabc1974..936c443c136e1 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchFilter.java @@ -23,21 +23,14 @@ */ public class OpenSearchFilter extends Filter implements OpenSearchRelNode { - private final String backend; private final List viableBackends; public OpenSearchFilter(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - RexNode condition, String backend, List viableBackends) { + RexNode condition, List viableBackends) { super(cluster, traitSet, input, condition); - this.backend = backend; this.viableBackends = viableBackends; } - @Override - public String getBackend() { - return backend; - } - @Override public List getViableBackends() { return viableBackends; @@ -54,7 +47,7 @@ public List getOutputFieldStorage() { @Override public Filter copy(RelTraitSet traitSet, RelNode input, RexNode condition) { - return new OpenSearchFilter(getCluster(), traitSet, input, condition, backend, viableBackends); + return new OpenSearchFilter(getCluster(), traitSet, input, condition, viableBackends); } @Override @@ -65,6 +58,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(org.apache.calcite.pla @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("backend", backend).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java index 413d50de95906..c31b7d72942ed 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java @@ -17,8 +17,13 @@ import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.metadata.RelMetadataQuery; import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexNode; import org.opensearch.analytics.planner.FieldStorageInfo; +import org.opensearch.analytics.planner.RelNodeUtils; + +import java.util.ArrayList; +import java.util.List; import java.util.List; @@ -27,35 +32,46 @@ */ public class OpenSearchProject extends Project implements OpenSearchRelNode { - private final String backend; + private final List viableBackends; public OpenSearchProject(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, List projects, RelDataType rowType, - String backend) { + List viableBackends) { super(cluster, traitSet, List.of(), input, projects, rowType); - this.backend = backend; - } - - @Override - public String getBackend() { - return backend; + this.viableBackends = viableBackends; } @Override public List getViableBackends() { - return List.of(backend); + return viableBackends; } @Override public List getOutputFieldStorage() { - return getRowType().getFieldList().stream() - .map(field -> FieldStorageInfo.derivedColumn(field.getName(), field.getType().toString())) - .toList(); + RelNode input = RelNodeUtils.unwrapHep(getInput()); + if (!(input instanceof OpenSearchRelNode openSearchChild)) { + throw new IllegalStateException( + "Project child is not OpenSearchRelNode: " + input.getClass().getSimpleName()); + } + List inputStorage = openSearchChild.getOutputFieldStorage(); + + List result = new ArrayList<>(getProjects().size()); + for (int i = 0; i < getProjects().size(); i++) { + RexNode expr = getProjects().get(i); + if (expr instanceof RexInputRef ref && ref.getIndex() < inputStorage.size()) { + result.add(inputStorage.get(ref.getIndex())); + } else { + String fieldName = getRowType().getFieldList().get(i).getName(); + result.add(FieldStorageInfo.derivedColumn(fieldName, + getRowType().getFieldList().get(i).getType().getSqlTypeName())); + } + } + return result; } @Override public Project copy(RelTraitSet traitSet, RelNode input, List projects, RelDataType rowType) { - return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, backend); + return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, viableBackends); } @Override @@ -65,6 +81,6 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("backend", backend); + return super.explainTerms(pw).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java index 36d13b93db68a..bacb2b62926f3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchRelNode.java @@ -19,20 +19,14 @@ *

Each node computes {@link #getOutputFieldStorage()} from its input's metadata. * Parent operators read this to make backend routing decisions. * - *

{@link #getBackend()} is the default/preferred backend chosen by the marking rule. - * {@link #getViableBackends()} lists all backends that could execute this operator - * (including via delegation). Consumed by {@code BackendResolver.generateCandidatePlans()} - * during StagePlan alternative generation. + *

{@link #getViableBackends()} lists all backends that could execute this operator + * (including via delegation). The first entry is the default/preferred backend. + * Consumed during plan forking to generate one complete plan per viable backend. * * @opensearch.internal */ public interface OpenSearchRelNode { - String UNRESOLVED = "unresolved"; - - /** Default/preferred backend chosen by the marking rule. */ - String getBackend(); - /** All backends that could execute this operator, including via delegation. */ List getViableBackends(); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java index d41944ccaa531..2359c13a141ad 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchShuffleReader.java @@ -32,31 +32,23 @@ */ public class OpenSearchShuffleReader extends SingleRel implements OpenSearchRelNode { - private final String backend; + private final List viableBackends; private final ShuffleImpl shuffleImpl; public OpenSearchShuffleReader(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, - String backend, ShuffleImpl shuffleImpl) { + List viableBackends, ShuffleImpl shuffleImpl) { super(cluster, traitSet, input); - this.backend = backend; + this.viableBackends = viableBackends; this.shuffleImpl = shuffleImpl; } - @Override - public String getBackend() { - return backend; - } - public ShuffleImpl getShuffleImpl() { return shuffleImpl; } @Override public List getViableBackends() { - if (getInput() instanceof OpenSearchRelNode openSearchInput) { - return openSearchInput.getViableBackends(); - } - return List.of(); + return viableBackends; } @Override @@ -69,7 +61,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchShuffleReader(getCluster(), traitSet, sole(inputs), backend, shuffleImpl); + return new OpenSearchShuffleReader(getCluster(), traitSet, sole(inputs), viableBackends, shuffleImpl); } @Override @@ -79,6 +71,6 @@ public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("backend", backend).item("shuffleImpl", shuffleImpl); + return super.explainTerms(pw).item("viableBackends", viableBackends).item("shuffleImpl", shuffleImpl); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java index 58dafcd9f7d23..3098d37cd971e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchSort.java @@ -24,22 +24,15 @@ */ public class OpenSearchSort extends Sort implements OpenSearchRelNode { - private final String backend; private final List viableBackends; public OpenSearchSort(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, RelCollation collation, RexNode offset, RexNode fetch, - String backend, List viableBackends) { + List viableBackends) { super(cluster, traitSet, input, collation, offset, fetch); - this.backend = backend; this.viableBackends = viableBackends; } - @Override - public String getBackend() { - return backend; - } - @Override public List getViableBackends() { return viableBackends; @@ -56,7 +49,7 @@ public List getOutputFieldStorage() { @Override public Sort copy(RelTraitSet traitSet, RelNode input, RelCollation collation, RexNode offset, RexNode fetch) { - return new OpenSearchSort(getCluster(), traitSet, input, collation, offset, fetch, backend, viableBackends); + return new OpenSearchSort(getCluster(), traitSet, input, collation, offset, fetch, viableBackends); } @Override @@ -67,6 +60,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(org.apache.calcite.pla @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("backend", backend).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java index 93b6cf4b37947..dc22910d9b3fa 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java @@ -25,15 +25,13 @@ */ public class OpenSearchTableScan extends TableScan implements OpenSearchRelNode { - private final String backend; private final List viableBackends; private final List outputFieldStorage; public OpenSearchTableScan(RelOptCluster cluster, RelTraitSet traitSet, RelOptTable table, - String backend, List viableBackends, + List viableBackends, List outputFieldStorage) { super(cluster, traitSet, List.of(), table); - this.backend = backend; this.viableBackends = viableBackends; this.outputFieldStorage = outputFieldStorage; } @@ -44,7 +42,7 @@ public OpenSearchTableScan(RelOptCluster cluster, RelTraitSet traitSet, RelOptTa * Single shard → SINGLETON (all data on one node). */ public static OpenSearchTableScan create(RelOptCluster cluster, RelOptTable table, - String backend, List viableBackends, + List viableBackends, List outputFieldStorage, int shardCount, OpenSearchDistributionTraitDef distTraitDef) { @@ -54,12 +52,7 @@ public static OpenSearchTableScan create(RelOptCluster cluster, RelOptTable tabl RelTraitSet traitSet = RelTraitSet.createEmpty() .plus(OpenSearchConvention.INSTANCE) .plus(distribution); - return new OpenSearchTableScan(cluster, traitSet, table, backend, viableBackends, outputFieldStorage); - } - - @Override - public String getBackend() { - return backend; + return new OpenSearchTableScan(cluster, traitSet, table, viableBackends, outputFieldStorage); } @Override @@ -74,7 +67,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchTableScan(getCluster(), traitSet, getTable(), backend, viableBackends, outputFieldStorage); + return new OpenSearchTableScan(getCluster(), traitSet, getTable(), viableBackends, outputFieldStorage); } @Override @@ -85,6 +78,6 @@ public org.apache.calcite.plan.RelOptCost computeSelfCost(RelOptPlanner planner, @Override public RelWriter explainTerms(RelWriter pw) { - return super.explainTerms(pw).item("backend", backend).item("viableBackends", viableBackends); + return super.explainTerms(pw).item("viableBackends", viableBackends); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java index 7db6276cf0164..abe1b52acf7bf 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java @@ -16,6 +16,8 @@ import org.apache.calcite.rel.core.AggregateCall; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.planner.CapabilityRegistry; +import org.opensearch.analytics.planner.FieldStorageInfo; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.AggregateCallAnnotation; @@ -23,8 +25,8 @@ import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.spi.AggregateFunction; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.OperatorCapability; import java.util.ArrayList; @@ -70,12 +72,13 @@ public void onMatch(RelOptRuleCall call) { "Aggregate rule encountered unmarked child [" + child.getClass().getSimpleName() + "]"); } - String childBackend = openSearchChild.getBackend(); + List childViableBackends = openSearchChild.getViableBackends(); + List childFieldStorage = openSearchChild.getOutputFieldStorage(); // Annotate each AggregateCall with per-call viable backends List annotatedCalls = new ArrayList<>(); for (AggregateCall aggCall : aggregate.getAggCallList()) { - List callViable = resolveViableBackendsForCall(aggCall); + List callViable = resolveViableBackendsForCall(aggCall, childFieldStorage); if (callViable.isEmpty()) { throw new IllegalStateException( "No backend supports aggregate function [" + aggCall.getAggregation().getName() + "]"); @@ -83,8 +86,8 @@ public void onMatch(RelOptRuleCall call) { annotatedCalls.add(AggregateCallAnnotation.annotate(aggCall, callViable)); } - // Compute operator-level viable backends considering delegation - List viableBackends = computeAggregateViableBackends(annotatedCalls, childBackend); + // Compute operator-level viable backends: must be viable for child AND handle agg calls + List viableBackends = computeAggregateViableBackends(annotatedCalls, childViableBackends); if (viableBackends.isEmpty()) { List funcNames = aggregate.getAggCallList().stream() @@ -92,12 +95,10 @@ public void onMatch(RelOptRuleCall call) { .toList(); throw new IllegalStateException( "No backend can execute aggregate: functions " + funcNames - + " are split across backends and no delegation path exists"); + + " not supported by any viable backend among " + childViableBackends); } - String backend = viableBackends.contains(childBackend) - ? childBackend - : viableBackends.getFirst(); + logger.debug("Aggregate viable backends: {} (child viable: {})", viableBackends, childViableBackends); RelTraitSet aggregateTraits = child.getTraitSet(); if (aggregateTraits.size() > 0) { @@ -112,41 +113,78 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), annotatedCalls, AggregateMode.SINGLE, - backend, viableBackends )); } - private List resolveViableBackendsForCall(AggregateCall aggCall) { + private List resolveViableBackendsForCall(AggregateCall aggCall, + List childFieldStorage) { AggregateFunction func = AggregateFunction.fromSqlKind(aggCall.getAggregation().getKind()); if (func == null) { func = AggregateFunction.fromNameOrError(aggCall.getAggregation().getName()); } - List viable = new ArrayList<>(); - for (AnalyticsSearchBackendPlugin plugin : context.getBackends().values()) { - if (plugin.supportedOperators().contains(OperatorCapability.AGGREGATE) - && plugin.supportedAggregateFunctions().contains(func)) { - viable.add(plugin.name()); + CapabilityRegistry registry = context.getCapabilityRegistry(); + + if (aggCall.getArgList().isEmpty()) { + return new ArrayList<>(registry.operatorBackends(OperatorCapability.AGGREGATE)); + } + + List callViable = null; + for (int fieldIndex : aggCall.getArgList()) { + if (fieldIndex >= childFieldStorage.size()) { + continue; + } + FieldStorageInfo storageInfo = childFieldStorage.get(fieldIndex); + FieldType fieldType = storageInfo.getFieldType(); + if (fieldType == null) { + throw new IllegalStateException("Unrecognized field type [" + storageInfo.getMappingType() + + "] for field [" + storageInfo.getFieldName() + "]"); + } + + List perFieldBackends = new ArrayList<>(); + if (storageInfo.isDerived()) { + perFieldBackends.addAll(registry.aggregateBackendsAnyFormat(func, fieldType)); + } else { + // Format-aware: backends that can read the data and compute + for (String format : storageInfo.getDocValueFormats()) { + for (String name : registry.aggregateBackends(func, fieldType, format)) { + if (!perFieldBackends.contains(name)) { + perFieldBackends.add(name); + } + } + } + // Format-agnostic: delegation targets that can compute but don't need data access + for (String name : registry.aggregateBackendsAnyFormat(func, fieldType)) { + if (!perFieldBackends.contains(name)) { + perFieldBackends.add(name); + } + } + } + + if (callViable == null) { + callViable = perFieldBackends; + } else { + callViable.retainAll(perFieldBackends); } } - return viable; + + return callViable != null ? callViable : new ArrayList<>(registry.operatorBackends(OperatorCapability.AGGREGATE)); } - /** - * Computes which backends can execute this aggregate, considering delegation. - * A backend is viable if for every agg call it can handle natively OR delegate - * (supports AGGREGATE delegation AND some other backend accepts it for that call). - */ private List computeAggregateViableBackends(List annotatedCalls, - String childBackend) { + List childViableBackends) { if (annotatedCalls.isEmpty()) { - return new ArrayList<>(context.getBackends().keySet()); + return new ArrayList<>(childViableBackends); } + CapabilityRegistry registry = context.getCapabilityRegistry(); + List delegationSupporters = registry.delegationSupporters(DelegationType.AGGREGATE); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.AGGREGATE); + List viable = new ArrayList<>(); - for (AnalyticsSearchBackendPlugin candidate : context.getBackends().values()) { - if (!candidate.supportedOperators().contains(OperatorCapability.AGGREGATE)) { + for (String candidateName : childViableBackends) { + if (!registry.operatorBackends(OperatorCapability.AGGREGATE).contains(candidateName)) { continue; } @@ -158,24 +196,18 @@ private List computeAggregateViableBackends(List annotate break; } List callViable = annotation.getViableBackends(); - if (callViable.contains(candidate.name())) { + if (callViable.contains(candidateName)) { continue; } - // Check if candidate can delegate this call - if (candidate.supportedDelegations().contains(DelegationType.AGGREGATE)) { - boolean someoneAccepts = callViable.stream().anyMatch(backendName -> { - AnalyticsSearchBackendPlugin other = context.getBackends().get(backendName); - return other != null && other.acceptedDelegations().contains(DelegationType.AGGREGATE); - }); - if (someoneAccepts) { - continue; - } + if (delegationSupporters.contains(candidateName) + && callViable.stream().anyMatch(delegationAcceptors::contains)) { + continue; } canHandleAll = false; break; } if (canHandleAll) { - viable.add(candidate.name()); + viable.add(candidateName); } } return viable; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java index d0001238e402b..c010728ad5223 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java @@ -50,8 +50,6 @@ public void onMatch(RelOptRuleCall call) { OpenSearchAggregate aggregate = call.rel(0); RelNode child = call.rel(1); - String backend = aggregate.getBackend(); - // Partial aggregate: runs on each partition, keeps input's traits RelTraitSet partialTraits = child.getTraitSet() .replace(OpenSearchConvention.INSTANCE); @@ -63,7 +61,6 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), aggregate.getAggCallList(), AggregateMode.PARTIAL, - backend, aggregate.getViableBackends() ); @@ -81,7 +78,6 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), aggregate.getAggCallList(), AggregateMode.FINAL, - backend, aggregate.getViableBackends() ); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java index 773fb0341e051..21fa5a9d2abd3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java @@ -22,24 +22,23 @@ import org.apache.logging.log4j.Logger; import org.opensearch.analytics.planner.FieldStorageInfo; import org.opensearch.analytics.planner.PlannerContext; +import org.opensearch.analytics.planner.CapabilityRegistry; +import org.opensearch.analytics.planner.FieldStorageInfo; +import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.AnnotatedPredicate; import org.opensearch.analytics.planner.rel.FullTextFunctions; import org.opensearch.analytics.planner.rel.OpenSearchFilter; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; -import org.opensearch.analytics.spi.FieldTypeFamily; -import org.opensearch.analytics.spi.FilterCapability; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.FilterOperator; -import org.opensearch.analytics.spi.FullTextOperator; import org.opensearch.analytics.spi.OperatorCapability; import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; /** * Converts {@link Filter} → {@link OpenSearchFilter}. @@ -81,34 +80,31 @@ public void onMatch(RelOptRuleCall call) { + "]. Ensure all child operators are marked before filter."); } - String childBackend = openSearchInput.getBackend(); + List childViableBackends = openSearchInput.getViableBackends(); List childFieldStorage = openSearchInput.getOutputFieldStorage(); RelDataType inputRowType = child.getRowType(); // Annotate every leaf predicate with viable backends - RexNode annotatedCondition = annotateCondition(filter.getCondition(), inputRowType, childFieldStorage); + RexNode annotatedCondition = annotateCondition(filter.getCondition(), inputRowType, + childFieldStorage, childViableBackends); - // Compute operator-level viable backends from per-predicate annotations + delegation - List viableBackends = computeFilterViableBackends(annotatedCondition, childBackend); + // Compute operator-level viable backends: must be viable for child AND handle predicates + List viableBackends = computeFilterViableBackends(annotatedCondition, childViableBackends); if (viableBackends.isEmpty()) { throw new IllegalStateException( - "No backend can execute filter: child backend [" + childBackend - + "] cannot evaluate all predicates and no delegation path exists"); + "No backend can execute filter: no viable backend among " + childViableBackends + + " can evaluate all predicates and no delegation path exists"); } - // Preferred backend: child if viable, otherwise first viable - String filterBackend = viableBackends.contains(childBackend) - ? childBackend - : viableBackends.getFirst(); + LOGGER.debug("Filter viable backends: {} (child viable: {})", viableBackends, childViableBackends); call.transformTo(new OpenSearchFilter( filter.getCluster(), child.getTraitSet(), RelNodeUtils.unwrapHep(filter.getInput()), annotatedCondition, - filterBackend, viableBackends )); } @@ -121,18 +117,19 @@ public void onMatch(RelOptRuleCall call) { * {@link AnnotatedPredicate} with viable backends resolved from child's field storage. */ private RexNode annotateCondition(RexNode condition, RelDataType inputRowType, - List fieldStorage) { + List fieldStorage, + List childViableBackends) { if (!(condition instanceof RexCall rexCall)) { return condition; } if (rexCall.getKind() == SqlKind.AND || rexCall.getKind() == SqlKind.OR || rexCall.getKind() == SqlKind.NOT) { List annotatedOperands = new ArrayList<>(); for (RexNode operand : rexCall.getOperands()) { - annotatedOperands.add(annotateCondition(operand, inputRowType, fieldStorage)); + annotatedOperands.add(annotateCondition(operand, inputRowType, fieldStorage, childViableBackends)); } return rexCall.clone(rexCall.getType(), annotatedOperands); } - List viableBackends = resolveViableBackends(rexCall, inputRowType, fieldStorage); + List viableBackends = resolveViableBackends(rexCall, inputRowType, fieldStorage, childViableBackends); return new AnnotatedPredicate(rexCall.getType(), rexCall, viableBackends); } @@ -143,82 +140,78 @@ private RexNode annotateCondition(RexNode condition, RelDataType inputRowType, * Intersects across all referenced fields. */ private List resolveViableBackends(RexCall predicate, RelDataType inputRowType, - List fieldStorage) { - // TODO : Try collapsing in one pass + List fieldStorage, + List childViableBackends) { Set fieldIndices = new HashSet<>(); collectFieldIndices(predicate, fieldIndices); + + CapabilityRegistry registry = context.getCapabilityRegistry(); + if (fieldIndices.isEmpty()) { - // No field refs (e.g. literal expression) — all backends with FILTER capability - return context.getBackends().values().stream() - .filter(backend -> backend.supportedOperators().contains(OperatorCapability.FILTER)) - .map(AnalyticsSearchBackendPlugin::name) - .collect(Collectors.toList()); + return new ArrayList<>(registry.operatorBackends(OperatorCapability.FILTER)); } - FilterOperator filterOp = FilterOperator.fromSqlKind(predicate.getKind()); - - // Check if this is a full-text function (MATCH, MATCH_PHRASE, etc.) - FullTextOperator fullTextOp = null; + FilterOperator operator = null; if (predicate.getOperator() instanceof SqlFunction sqlFunction) { - fullTextOp = FullTextFunctions.toFullTextOperator(sqlFunction); + operator = FullTextFunctions.toFilterOperator(sqlFunction); } - - // Start with all backends that have FILTER capability, intersect per field - Set viableSet = new HashSet<>(); - for (AnalyticsSearchBackendPlugin backend : context.getBackends().values()) { - if (backend.supportedOperators().contains(OperatorCapability.FILTER)) { - viableSet.add(backend.name()); - } + if (operator == null) { + operator = FilterOperator.fromSqlKind(predicate.getKind()); + } + if (operator == null) { + throw new IllegalStateException("Unrecognized filter operator [" + predicate.getKind() + "]"); } + Set viableSet = new HashSet<>(registry.operatorBackends(OperatorCapability.FILTER)); + for (int fieldIndex : fieldIndices) { if (fieldIndex >= fieldStorage.size()) { continue; } FieldStorageInfo storageInfo = fieldStorage.get(fieldIndex); - - // Derived/expression column — only backends that can filter on expressions - if (storageInfo.isDerived()) { - viableSet.retainAll( - context.getBackends().values().stream() - .filter(b -> b.supportedOperators().contains(OperatorCapability.FILTER_ON_EXPRESSIONS)) - .map(AnalyticsSearchBackendPlugin::name) - .collect(Collectors.toSet()) - ); - continue; + FieldType fieldType = storageInfo.getFieldType(); + if (fieldType == null) { + throw new IllegalStateException("Unrecognized field type [" + storageInfo.getMappingType() + + "] for field [" + storageInfo.getFieldName() + "]"); } - FieldTypeFamily typeFamily = FieldTypeFamily.fromMappingType(storageInfo.getFieldType()); - final FullTextOperator finalFullTextOp = fullTextOp; + // TODO: for FULL_TEXT operators, extract required params from RexCall + // and use registry.fullTextFilterBackends() instead Set fieldViable = new HashSet<>(); - for (AnalyticsSearchBackendPlugin backend : context.getBackends().values()) { - if (!viableSet.contains(backend.name())) { - continue; + if (storageInfo.isDerived()) { + // Derived column — only child viable backends + their delegation targets + List anyFormat = registry.filterBackendsAnyFormat(operator, fieldType); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.FILTER); + for (String name : childViableBackends) { + if (anyFormat.contains(name)) { + fieldViable.add(name); + } } - - boolean formatMatch = backend.getSupportedFormats().stream().anyMatch(format -> - storageInfo.getDocValueFormats().contains(format.name()) - || storageInfo.getIndexFormats().contains(format.name()) - ); - - // Full-text ops require the backend to support that specific operator - // AND the field must have an index (full-text needs inverted index) - boolean operatorMatch; - if (finalFullTextOp != null) { - operatorMatch = backend.supportedFullTextOperators().contains(finalFullTextOp) - && storageInfo.hasIndex(); - } else if (filterOp != null && typeFamily != null) { - operatorMatch = backend.supportedFilterCapabilities() - .contains(FilterCapability.of(filterOp, typeFamily)); - } else { - // Unknown operator or unrecognized field type — accept if backend has FILTER - operatorMatch = true; + // Delegation targets reachable from child viable backends + List delegationSupporters = registry.delegationSupporters(DelegationType.FILTER); + if (childViableBackends.stream().anyMatch(delegationSupporters::contains)) { + for (String name : anyFormat) { + if (!fieldViable.contains(name) && delegationAcceptors.contains(name)) { + fieldViable.add(name); + } + } } - - if (formatMatch && operatorMatch) { - fieldViable.add(backend.name()); + } else { + // Format-aware: backends that can access the field's data + for (String format : storageInfo.getDocValueFormats()) { + fieldViable.addAll(registry.filterBackends(operator, fieldType, format)); + } + for (String format : storageInfo.getIndexFormats()) { + fieldViable.addAll(registry.filterBackends(operator, fieldType, format)); + } + // Format-agnostic: delegation targets that can evaluate but don't need data access + for (String name : registry.filterBackendsAnyFormat(operator, fieldType)) { + if (!fieldViable.contains(name)) { + fieldViable.add(name); + } } } + viableSet.retainAll(fieldViable); } @@ -226,7 +219,7 @@ private List resolveViableBackends(RexCall predicate, RelDataType inputR throw new IllegalStateException("No backend can evaluate filter predicate [" + predicate.getKind() + "] on fields " + fieldIndices.stream() .filter(i -> i < fieldStorage.size()) - .map(i -> fieldStorage.get(i).getFieldName() + ":" + fieldStorage.get(i).getFieldType()) + .map(i -> fieldStorage.get(i).getFieldName() + ":" + fieldStorage.get(i).getMappingType()) .toList()); } return new ArrayList<>(viableSet); @@ -254,55 +247,40 @@ private void collectFieldIndices(RexNode node, Set result) { * (child supports FILTER delegation, this backend accepts it, and this backend * can handle all predicates natively). */ - private List computeFilterViableBackends(RexNode annotatedCondition, String childBackend) { + private List computeFilterViableBackends(RexNode annotatedCondition, + List childViableBackends) { List predicates = new ArrayList<>(); collectAnnotatedPredicates(annotatedCondition, predicates); if (predicates.isEmpty()) { - return new ArrayList<>(context.getBackends().keySet()); + return new ArrayList<>(childViableBackends); } - AnalyticsSearchBackendPlugin childPlugin = context.getBackends().get(childBackend); List viable = new ArrayList<>(); + CapabilityRegistry registry = context.getCapabilityRegistry(); + List delegationSupporters = registry.delegationSupporters(DelegationType.FILTER); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.FILTER); - for (AnalyticsSearchBackendPlugin candidate : context.getBackends().values()) { - if (!candidate.supportedOperators().contains(OperatorCapability.FILTER)) { + for (String candidateName : childViableBackends) { + if (!registry.operatorBackends(OperatorCapability.FILTER).contains(candidateName)) { continue; } - boolean isChild = candidate.name().equals(childBackend); - - // Non-child backend: only viable if child can delegate entire filter to it - // and it accepts delegation and can handle all predicates natively - if (!isChild) { - if (childPlugin == null - || !childPlugin.supportedDelegations().contains(DelegationType.FILTER) - || !candidate.acceptedDelegations().contains(DelegationType.FILTER)) { - continue; - } - } - boolean canHandleAll = true; for (AnnotatedPredicate predicate : predicates) { List predViable = predicate.getViableBackends(); - if (predViable.contains(candidate.name())) { + if (predViable.contains(candidateName)) { continue; } - // Child backend can delegate individual predicates - if (isChild && candidate.supportedDelegations().contains(DelegationType.FILTER)) { - boolean someoneAccepts = predViable.stream().anyMatch(backendName -> { - AnalyticsSearchBackendPlugin other = context.getBackends().get(backendName); - return other != null && other.acceptedDelegations().contains(DelegationType.FILTER); - }); - if (someoneAccepts) { - continue; - } + if (delegationSupporters.contains(candidateName) + && predViable.stream().anyMatch(delegationAcceptors::contains)) { + continue; } canHandleAll = false; break; } if (canHandleAll) { - viable.add(candidate.name()); + viable.add(candidateName); } } return viable; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java index 815ec285916f4..161291de1f85a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java @@ -15,13 +15,14 @@ import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlFunction; +import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.AnnotatedProjectExpression; import org.opensearch.analytics.planner.rel.OpenSearchProject; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.ScalarFunction; import java.util.ArrayList; @@ -58,11 +59,19 @@ public void onMatch(RelOptRuleCall call) { "Project rule encountered unmarked child [" + child.getClass().getSimpleName() + "]"); } - String backend = openSearchChild.getBackend(); + List childViableBackends = openSearchChild.getViableBackends(); + // TODO: precompute SqlKind → viable backends map to avoid repeated filtering per node + // TODO: reuse childViableBackends list when all candidates pass instead of allocating List annotatedExprs = new ArrayList<>(project.getProjects().size()); for (RexNode expr : project.getProjects()) { - annotatedExprs.add(annotateExpr(expr, backend)); + annotatedExprs.add(annotateExpr(expr, childViableBackends)); + } + + List viableBackends = computeProjectViableBackends(annotatedExprs, childViableBackends); + if (viableBackends.isEmpty()) { + throw new IllegalStateException( + "No backend can execute all project expressions among " + childViableBackends); } call.transformTo(new OpenSearchProject( @@ -71,49 +80,41 @@ public void onMatch(RelOptRuleCall call) { RelNodeUtils.unwrapHep(project.getInput()), annotatedExprs, project.getRowType(), - backend + viableBackends )); } - private RexNode annotateExpr(RexNode expr, String backend) { + private RexNode annotateExpr(RexNode expr, List childViableBackends) { if (!(expr instanceof RexCall rexCall)) { return expr; } - AnalyticsSearchBackendPlugin plugin = context.getBackends().get(backend); - if (plugin == null) { - throw new IllegalStateException("Backend [" + backend + "] not found"); - } - - // Opaque operations (painless, highlighting, etc.) — no recursion into operands + // Opaque operations — no recursion into operands if (rexCall.getOperator() instanceof SqlFunction sqlFunction) { String funcName = sqlFunction.getName(); if (isOpaqueOperation(funcName)) { - if (plugin.supportedOpaqueProjectOperations().contains(funcName)) { - return new AnnotatedProjectExpression(rexCall.getType(), rexCall, backend); - } - String delegationTarget = canDelegateProject(plugin, funcName); - if (delegationTarget != null) { - return new AnnotatedProjectExpression(rexCall.getType(), rexCall, delegationTarget); + List exprViable = resolveOpaqueViableBackends(funcName, childViableBackends); + if (exprViable.isEmpty()) { + throw new IllegalStateException( + "No backend can evaluate [" + funcName + "] and no delegation path exists"); } - throw new IllegalStateException( - "Backend [" + backend + "] cannot evaluate [" + funcName - + "] and no delegation path exists"); + return new AnnotatedProjectExpression(rexCall.getType(), rexCall, exprViable); } } - // Standard scalar function — validate support - ScalarFunction scalarFunc = ScalarFunction.fromSqlKind(rexCall.getKind()); - if (scalarFunc != null && !plugin.supportedScalarFunctions().contains(scalarFunc)) { + // Standard scalar function + List scalarViable = resolveScalarViableBackends(rexCall, childViableBackends); + if (scalarViable.isEmpty()) { throw new IllegalStateException( - "Backend [" + backend + "] does not support scalar function [" + scalarFunc + "]"); + "No backend supports scalar function [" + ScalarFunction.fromSqlKind(rexCall.getKind()) + + "] among " + childViableBackends); } - // Recurse into operands to validate and annotate nested expressions + // Recurse into operands boolean changed = false; List newOperands = new ArrayList<>(rexCall.getOperands().size()); for (RexNode operand : rexCall.getOperands()) { - RexNode annotated = annotateExpr(operand, backend); + RexNode annotated = annotateExpr(operand, childViableBackends); newOperands.add(annotated); if (annotated != operand) { changed = true; @@ -121,24 +122,98 @@ private RexNode annotateExpr(RexNode expr, String backend) { } RexCall target = changed ? rexCall.clone(rexCall.getType(), newOperands) : rexCall; - return new AnnotatedProjectExpression(target.getType(), target, backend); + return new AnnotatedProjectExpression(target.getType(), target, scalarViable); } - private boolean isOpaqueOperation(String funcName) { - return context.getBackends().values().stream() - .anyMatch(b -> b.supportedOpaqueProjectOperations().contains(funcName)); + private List resolveOpaqueViableBackends(String funcName, List childViableBackends) { + CapabilityRegistry registry = context.getCapabilityRegistry(); + List viable = registry.opaqueBackendsAnyFormat(funcName); + if (viable.isEmpty()) { + return viable; + } + // At least one child viable backend must be able to reach an evaluator: + // either it's in viable itself (native), or it can delegate to one that accepts + List delegationSupporters = registry.delegationSupporters(DelegationType.PROJECT); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); + boolean reachable = childViableBackends.stream().anyMatch(candidateName -> + viable.contains(candidateName) + || (delegationSupporters.contains(candidateName) + && viable.stream().anyMatch(delegationAcceptors::contains))); + return reachable ? viable : List.of(); } - private String canDelegateProject(AnalyticsSearchBackendPlugin plugin, String funcName) { - if (!plugin.supportedDelegations().contains(DelegationType.PROJECT)) { - return null; - } - return context.getBackends().values().stream() - .filter(other -> !other.name().equals(plugin.name()) - && other.acceptedDelegations().contains(DelegationType.PROJECT) - && other.supportedOpaqueProjectOperations().contains(funcName)) - .map(AnalyticsSearchBackendPlugin::name) - .findFirst() - .orElse(null); + private List resolveScalarViableBackends(RexCall rexCall, List childViableBackends) { + ScalarFunction scalarFunc = ScalarFunction.fromSqlKind(rexCall.getKind()); + if (scalarFunc == null) { + return List.of(); + } + FieldType fieldType = FieldType.fromSqlTypeName(rexCall.getType().getSqlTypeName()); + if (fieldType == null) { + return List.of(); + } + + CapabilityRegistry registry = context.getCapabilityRegistry(); + List allCapable = registry.scalarBackendsAnyFormat(scalarFunc, fieldType); + + // Prefer child viable backends + List viable = new ArrayList<>(); + for (String candidateName : childViableBackends) { + if (allCapable.contains(candidateName)) { + viable.add(candidateName); + } + } + if (!viable.isEmpty()) { + return viable; + } + // Fallback: other backends if reachable via delegation + List delegationSupporters = registry.delegationSupporters(DelegationType.PROJECT); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); + boolean canDelegate = childViableBackends.stream().anyMatch(delegationSupporters::contains); + if (!canDelegate) { + return viable; + } + for (String backendName : allCapable) { + if (delegationAcceptors.contains(backendName)) { + viable.add(backendName); + } + } + return viable; + } + + private List computeProjectViableBackends(List annotatedExprs, + List childViableBackends) { + // A child viable backend is viable for the project if for every expression it can + // either evaluate natively (present in expression's viableBackends) or delegate to + // a backend that can (supports PROJECT delegation to an acceptor in expression's viableBackends) + CapabilityRegistry registry = context.getCapabilityRegistry(); + List delegationSupporters = registry.delegationSupporters(DelegationType.PROJECT); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); + + List result = new ArrayList<>(); + for (String candidateName : childViableBackends) { + boolean canHandleAll = true; + for (RexNode expr : annotatedExprs) { + if (!(expr instanceof AnnotatedProjectExpression annotation)) { + continue; + } + if (annotation.getViableBackends().contains(candidateName)) { + continue; + } + boolean canDelegate = delegationSupporters.contains(candidateName) + && annotation.getViableBackends().stream().anyMatch(delegationAcceptors::contains); + if (!canDelegate) { + canHandleAll = false; + break; + } + } + if (canHandleAll) { + result.add(candidateName); + } + } + return result; + } + + private boolean isOpaqueOperation(String funcName) { + return context.getCapabilityRegistry().isOpaqueOperation(funcName); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java index 3f72755c0f02e..e2503646c1b70 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java @@ -12,7 +12,6 @@ import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.Sort; -import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; @@ -55,13 +54,16 @@ public void onMatch(RelOptRuleCall call) { "Sort rule encountered unmarked child [" + child.getClass().getSimpleName() + "]"); } - String childBackend = openSearchChild.getBackend(); + List childViableBackends = openSearchChild.getViableBackends(); + List sortCapable = context.getCapabilityRegistry().operatorBackends(OperatorCapability.SORT); - String backend = CapabilityResolutionUtils.resolveBackend( - context.getBackends(), childBackend, OperatorCapability.SORT); + List viableBackends = childViableBackends.stream() + .filter(sortCapable::contains) + .toList(); - if (!CapabilityResolutionUtils.backendSupports(context.getBackends(), backend, OperatorCapability.SORT)) { - throw new IllegalStateException("No backend supports SORT capability"); + if (viableBackends.isEmpty()) { + throw new IllegalStateException( + "No backend supports SORT capability among " + childViableBackends); } call.transformTo(new OpenSearchSort( @@ -71,8 +73,7 @@ public void onMatch(RelOptRuleCall call) { sort.getCollation(), sort.offset, sort.fetch, - backend, - List.of(backend) + viableBackends )); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java index 12310c2c83683..56bd6e87cabaf 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java @@ -12,14 +12,16 @@ import org.apache.calcite.plan.RelOptRuleCall; import org.apache.calcite.rel.core.TableScan; import org.apache.calcite.rel.type.RelDataTypeField; -import org.opensearch.analytics.planner.CapabilityResolutionUtils; +import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.FieldStorageInfo; import org.opensearch.analytics.planner.FieldStorageResolver; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; +import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.cluster.metadata.IndexMetadata; +import java.util.ArrayList; import java.util.List; /** @@ -54,20 +56,54 @@ public void onMatch(RelOptRuleCall call) { } String primaryFormat = indexMetadata.getSettings().get("index.composite.primary_data_format", "lucene"); - List viableBackends = CapabilityResolutionUtils.computeViableBackends( - context.getBackends(), OperatorCapability.SCAN, primaryFormat); - if (viableBackends.isEmpty()) { - throw new IllegalStateException("No backend supports format [" + primaryFormat - + "] with SCAN capability for index [" + indexMetadata.getIndex().getName() + "]"); - } + CapabilityRegistry registry = context.getCapabilityRegistry(); + FieldStorageResolver fieldStorageResolver = new FieldStorageResolver(indexMetadata, + registry.getBackends()); + // TODO : This expects the FrontEnds to attach the row type with all fields. + // TODO : How will they attach if we perform the index resolution List fieldNames = scan.getRowType().getFieldList().stream() .map(RelDataTypeField::getName) .toList(); - List fieldStorage = FieldStorageResolver.resolve(indexMetadata, fieldNames); + List fieldStorage = fieldStorageResolver.resolve(fieldNames); + + // Viable backends: must support SCAN and be able to read ALL requested fields + // (natively or via delegation to another backend that can read the field) + List scanCapable = registry.operatorBackends(OperatorCapability.SCAN); + List delegationSupporters = registry.delegationSupporters(DelegationType.SCAN); + List delegationAcceptors = registry.delegationAcceptors(DelegationType.SCAN); + List viableBackends = new ArrayList<>(scanCapable); + + for (FieldStorageInfo field : fieldStorage) { + if (field.isDerived()) { + continue; + } + // Backends that can natively scan this field's doc values + List fieldBackends = new ArrayList<>(); + for (String format : field.getDocValueFormats()) { + for (String backend : registry.scanBackends(format)) { + if (!fieldBackends.contains(backend)) { + fieldBackends.add(backend); + } + } + } + // Keep candidates that can scan natively or delegate to one that can + viableBackends.removeIf(candidate -> { + if (fieldBackends.contains(candidate)) { + return false; + } + return !delegationSupporters.contains(candidate) + || fieldBackends.stream().noneMatch(delegationAcceptors::contains); + }); + } + + if (viableBackends.isEmpty()) { + throw new IllegalStateException("No backend can scan all requested fields on index [" + + indexMetadata.getIndex().getName() + "]"); + } call.transformTo(OpenSearchTableScan.create( - scan.getCluster(), scan.getTable(), viableBackends.getFirst(), viableBackends, fieldStorage, + scan.getCluster(), scan.getTable(), viableBackends, fieldStorage, indexMetadata.getNumberOfShards(), context.getDistributionTraitDef() )); } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java index 8a7ee88156a1e..1a3990bdbc1d6 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java @@ -22,9 +22,10 @@ import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; +import org.opensearch.analytics.spi.AggregateCapability; import org.opensearch.analytics.spi.AggregateFunction; -import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.OperatorCapability; import java.util.EnumSet; @@ -51,7 +52,7 @@ public void testPerCallAnnotation() { public void testViableBackendsPopulated() { OpenSearchAggregate agg = runAggregate(1, sumCall()); - assertEquals(MockDataFusionBackend.NAME, agg.getBackend()); + assertTrue(agg.getViableBackends().contains(MockDataFusionBackend.NAME)); assertFalse(agg.getViableBackends().isEmpty()); assertTrue(agg.getViableBackends().contains(MockDataFusionBackend.NAME)); } @@ -99,7 +100,7 @@ public void testAggregateErrorsWhenNoBackendSupportsFunction() { "size", Map.of("type", "integer") ), List.of(new MockDataFusionBackend() { @Override - public Set supportedAggregateFunctions() { + public Set aggregateCapabilities() { return Set.of(); // supports AGGREGATE capability but no functions } })); @@ -119,8 +120,10 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set supportedAggregateFunctions() { - return EnumSet.of(AggregateFunction.SUM, AggregateFunction.COUNT); + public Set aggregateCapabilities() { + return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( + AggregateFunction.SUM, Set.of(FieldType.INTEGER), + AggregateFunction.COUNT, Set.of(FieldType.INTEGER))); } }; @@ -134,7 +137,13 @@ public Set supportedAggregateFunctions() { OpenSearchAggregate agg = (OpenSearchAggregate) result; assertTrue(agg.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(agg.getViableBackends().contains(MockLuceneBackend.NAME)); + // Lucene not viable at operator level (no doc values for scan) + // but per-call annotation shows Lucene as viable for SUM + assertFalse(agg.getViableBackends().contains(MockLuceneBackend.NAME)); + AggregateCallAnnotation annotation = AggregateCallAnnotation.find(agg.getAggCallList().get(0)); + assertNotNull(annotation); + assertTrue(annotation.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertTrue(annotation.getViableBackends().contains(MockLuceneBackend.NAME)); } // ---- Scan ---- @@ -150,7 +159,7 @@ public void testTableScanResolvesBackendAndFieldStorage() { assertTrue(result instanceof OpenSearchTableScan); OpenSearchTableScan scan = (OpenSearchTableScan) result; - assertEquals(MockDataFusionBackend.NAME, scan.getBackend()); + assertTrue(scan.getViableBackends().contains(MockDataFusionBackend.NAME)); assertEquals(2, scan.getOutputFieldStorage().size()); assertEquals("status", scan.getOutputFieldStorage().get(0).getFieldName()); assertFalse(scan.getViableBackends().isEmpty()); @@ -223,8 +232,9 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set supportedAggregateFunctions() { - return EnumSet.of(AggregateFunction.SUM); // no COUNT + public Set aggregateCapabilities() { + return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( + AggregateFunction.SUM, Set.of(FieldType.INTEGER))); } }; @@ -276,8 +286,9 @@ public void testReducerPassthroughViableBackends() { public void testAggregateViableWithDelegation() { MockDataFusionBackend dfWithDelegation = new MockDataFusionBackend() { @Override - public Set supportedAggregateFunctions() { - return EnumSet.of(AggregateFunction.SUM); + public Set aggregateCapabilities() { + return aggCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), Map.of( + AggregateFunction.SUM, Set.of(FieldType.INTEGER))); } @Override public Set supportedDelegations() { @@ -290,8 +301,9 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set supportedAggregateFunctions() { - return EnumSet.of(AggregateFunction.STDDEV_POP); + public Set aggregateCapabilities() { + return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( + AggregateFunction.STDDEV_POP, Set.of(FieldType.INTEGER))); } @Override public Set acceptedDelegations() { @@ -321,8 +333,9 @@ public Set supportedOperators() { return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE); } @Override - public Set supportedAggregateFunctions() { - return EnumSet.of(AggregateFunction.STDDEV_POP); + public Set aggregateCapabilities() { + return aggCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), Map.of( + AggregateFunction.STDDEV_POP, Set.of(FieldType.INTEGER))); } }; @@ -335,6 +348,17 @@ public Set supportedAggregateFunctions() { IllegalStateException exception = expectThrows(IllegalStateException.class, () -> runPlanner(aggregate, context)); - assertTrue(exception.getMessage().contains("no delegation path exists")); + assertTrue(exception.getMessage().contains("not supported by any viable backend")); + } + + private static Set aggCaps(Set formats, + Map> funcToTypes) { + Set caps = new java.util.HashSet<>(); + for (var entry : funcToTypes.entrySet()) { + for (FieldType family : entry.getValue()) { + caps.add(new AggregateCapability(entry.getKey(), family, formats)); + } + } + return caps; } } 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 7b5a8bfb62ce7..2104bada7ba7b 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 @@ -114,7 +114,7 @@ protected PlannerContext buildContext(String primaryFormat, int shardCount, ClusterState clusterState = mock(ClusterState.class); when(clusterState.metadata()).thenReturn(metadata); - return new PlannerContext(backends, clusterState); + return new PlannerContext(new CapabilityRegistry(backends), clusterState); } // ---- Table builders ---- diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java index 3c68bc5d9308b..153221d4cd9fd 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/FilterRuleTests.java @@ -18,6 +18,7 @@ import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.util.ImmutableBitSet; +import org.apache.calcite.rex.RexCall; import org.opensearch.analytics.planner.rel.AnnotatedPredicate; import org.opensearch.analytics.planner.rel.FullTextFunctions; import org.opensearch.analytics.planner.rel.OpenSearchFilter; @@ -44,7 +45,7 @@ public void testNativePredicateAnnotatedWithBothBackends() { ), new String[]{"status", "size"}, new SqlTypeName[]{SqlTypeName.INTEGER, SqlTypeName.INTEGER}, makeEquals(0, SqlTypeName.INTEGER, 200)); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); assertTrue(result.getCondition() instanceof AnnotatedPredicate); AnnotatedPredicate annotated = (AnnotatedPredicate) result.getCondition(); assertTrue(annotated.getViableBackends().contains(MockDataFusionBackend.NAME)); @@ -71,20 +72,26 @@ public void testKeywordEqualsAnnotatedWithBothBackends() { /** Full-text with delegation — both backends viable at operator level. */ public void testFullTextViableWithDelegation() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( - "message", Map.of("type", "text") + "message", Map.of("type", "keyword") ), new String[]{"message"}, new SqlTypeName[]{SqlTypeName.VARCHAR}, makeFullTextCall(FullTextFunctions.MATCH_PHRASE, 0, "hello world")); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + // DF is viable at operator level (has doc values in parquet) assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); + // Lucene not viable at operator level — only delegation target + assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); + // MATCH_PHRASE predicate has Lucene as delegation target + AnnotatedPredicate predicate = (AnnotatedPredicate) result.getCondition(); + assertTrue("MATCH_PHRASE should be evaluable by Lucene", + predicate.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(predicate.getOriginal().toString().contains("MATCH_PHRASE")); } - /** AND with delegation — both backends viable. */ + /** AND with delegation — DF viable at operator, equals viable for both, MATCH_PHRASE delegated to Lucene. */ public void testAndWithDelegationBothViable() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( "status", Map.of("type", "integer"), - "message", Map.of("type", "text") + "message", Map.of("type", "keyword") ), new String[]{"status", "message"}, new SqlTypeName[]{SqlTypeName.INTEGER, SqlTypeName.VARCHAR}, makeAnd( makeEquals(0, SqlTypeName.INTEGER, 200), @@ -92,14 +99,21 @@ public void testAndWithDelegationBothViable() { )); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); + assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); + RexCall andCondition = (RexCall) result.getCondition(); + AnnotatedPredicate equalsPred = (AnnotatedPredicate) andCondition.getOperands().get(0); + AnnotatedPredicate matchPred = (AnnotatedPredicate) andCondition.getOperands().get(1); + assertTrue(equalsPred.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertTrue(equalsPred.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(matchPred.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(matchPred.getOriginal().toString().contains("MATCH_PHRASE")); } - /** OR across backends with delegation — both viable. */ + /** OR across backends — DF viable at operator, equals viable for both, MATCH delegated to Lucene. */ public void testOrAcrossBackendsWithDelegation() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( "status", Map.of("type", "integer"), - "message", Map.of("type", "text") + "message", Map.of("type", "keyword") ), new String[]{"status", "message"}, new SqlTypeName[]{SqlTypeName.INTEGER, SqlTypeName.VARCHAR}, makeCall(SqlStdOperatorTable.OR, makeEquals(0, SqlTypeName.INTEGER, 200), @@ -107,14 +121,21 @@ public void testOrAcrossBackendsWithDelegation() { )); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); + assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); + RexCall orCondition = (RexCall) result.getCondition(); + AnnotatedPredicate equalsPred = (AnnotatedPredicate) orCondition.getOperands().get(0); + AnnotatedPredicate matchPred = (AnnotatedPredicate) orCondition.getOperands().get(1); + assertTrue(equalsPred.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertTrue(equalsPred.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(matchPred.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(matchPred.getOriginal().toString().contains("MATCH")); } - /** OR of two full-text predicates with delegation — both viable. */ + /** OR of two full-text predicates — DF viable at operator, both predicates delegated to Lucene. */ public void testMultipleFullTextOrWithDelegation() { OpenSearchFilter result = runFilterWithDelegation("parquet", Map.of( - "title", Map.of("type", "text"), - "body", Map.of("type", "text") + "title", Map.of("type", "keyword"), + "body", Map.of("type", "keyword") ), new String[]{"title", "body"}, new SqlTypeName[]{SqlTypeName.VARCHAR, SqlTypeName.VARCHAR}, makeCall(SqlStdOperatorTable.OR, makeFullTextCall(FullTextFunctions.MATCH, 0, "hello"), @@ -122,7 +143,14 @@ public void testMultipleFullTextOrWithDelegation() { )); assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); - assertTrue(result.getViableBackends().contains(MockLuceneBackend.NAME)); + assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); + RexCall orCondition = (RexCall) result.getCondition(); + AnnotatedPredicate matchPred = (AnnotatedPredicate) orCondition.getOperands().get(0); + AnnotatedPredicate phrasePred = (AnnotatedPredicate) orCondition.getOperands().get(1); + assertTrue(matchPred.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(matchPred.getOriginal().toString().contains("MATCH")); + assertTrue(phrasePred.getViableBackends().contains(MockLuceneBackend.NAME)); + assertTrue(phrasePred.getOriginal().toString().contains("MATCH_PHRASE")); } // ---- Error cases ---- @@ -134,7 +162,7 @@ public void testFullTextErrorsWithoutDelegation() { LogicalFilter filter = LogicalFilter.create(stubScan(table), condition); PlannerContext context = buildContext("parquet", Map.of( - "message", Map.of("type", "text") + "message", Map.of("type", "keyword") )); IllegalStateException exception = expectThrows(IllegalStateException.class, @@ -154,12 +182,9 @@ public void testErrorForUnsupportedFieldTypeOperatorCombo() { "location", Map.of("type", "geo_point") )); - try { - runPlanner(filter, context); - fail("Expected IllegalStateException for unsupported field type"); - } catch (IllegalStateException e) { - assertTrue(e.getMessage().contains("No backend can")); - } + IllegalStateException exception = expectThrows(IllegalStateException.class, + () -> runPlanner(filter, context)); + assertTrue(exception.getMessage().contains("has no storage")); } // ---- Derived columns ---- 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 7d93f3a58c5d7..8335338c14315 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 @@ -11,37 +11,48 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; +import org.opensearch.analytics.spi.AggregateCapability; import org.opensearch.analytics.spi.AggregateFunction; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; -import org.opensearch.analytics.spi.FieldTypeFamily; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.FilterCapability; import org.opensearch.analytics.spi.FilterOperator; -import org.opensearch.analytics.spi.FullTextOperator; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; -import java.util.EnumSet; +import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.COLUMNAR_STORAGE; + +import java.util.HashSet; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; /** - * Mock DataFusion backend for tests. Supports parquet format, standard filter - * operators on NUMERIC/KEYWORD/DATE/BOOLEAN families, and common aggregates. + * Mock DataFusion backend for tests. Supports parquet format with columnar storage, + * standard filter operators on NUMERIC/KEYWORD/DATE/BOOLEAN, and common aggregates. * No full-text support. */ public class MockDataFusionBackend implements AnalyticsSearchBackendPlugin { public static final String NAME = "mock-parquet"; + public static final String PARQUET_DATA_FORMAT = "parquet"; + private static final Set DATAFUSION_FORMATS = Set.of(PARQUET_DATA_FORMAT); private static final Set OPERATOR_CAPS = Set.of( - OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.FILTER_ON_EXPRESSIONS, + OperatorCapability.SCAN, OperatorCapability.FILTER, OperatorCapability.AGGREGATE, OperatorCapability.SORT, OperatorCapability.PROJECT, OperatorCapability.COORDINATOR_REDUCE ); - private static final Set FILTER_OPS = EnumSet.of( + private static final Set SUPPORTED_TYPES = new HashSet<>(); + static { + SUPPORTED_TYPES.addAll(FieldType.numeric()); + SUPPORTED_TYPES.addAll(FieldType.keyword()); + SUPPORTED_TYPES.addAll(FieldType.date()); + SUPPORTED_TYPES.add(FieldType.BOOLEAN); + } + + private static final Set STANDARD_OPS = Set.of( FilterOperator.EQUALS, FilterOperator.NOT_EQUALS, FilterOperator.GREATER_THAN, FilterOperator.GREATER_THAN_OR_EQUAL, FilterOperator.LESS_THAN, FilterOperator.LESS_THAN_OR_EQUAL, @@ -49,60 +60,57 @@ public class MockDataFusionBackend implements AnalyticsSearchBackendPlugin { FilterOperator.IN, FilterOperator.LIKE ); - private static final Set SUPPORTED_FAMILIES = EnumSet.of( - FieldTypeFamily.NUMERIC, FieldTypeFamily.KEYWORD, - FieldTypeFamily.DATE, FieldTypeFamily.BOOLEAN - ); - - private static final Set FILTER_CAPS; - static { - FILTER_CAPS = FILTER_OPS.stream() - .flatMap(op -> SUPPORTED_FAMILIES.stream().map(family -> FilterCapability.of(op, family))) - .collect(Collectors.toUnmodifiableSet()); - } - - private static final Set AGG_FUNCTIONS = EnumSet.of( + private static final Set AGG_FUNCTIONS = Set.of( AggregateFunction.SUM, AggregateFunction.SUM0, AggregateFunction.MIN, AggregateFunction.MAX, AggregateFunction.COUNT, AggregateFunction.AVG ); - @Override - public String name() { - return NAME; + private static final Set FILTER_CAPS; + static { + Set caps = new HashSet<>(); + for (FilterOperator op : STANDARD_OPS) { + for (FieldType type : SUPPORTED_TYPES) { + caps.add(new FilterCapability.Standard(op, type, DATAFUSION_FORMATS)); + } + } + FILTER_CAPS = caps; } - @Override - public SearchExecEngine searcher(ExecutionContext ctx) { - return null; + private static final Set AGG_CAPS; + static { + Set caps = new HashSet<>(); + for (AggregateFunction func : AGG_FUNCTIONS) { + for (FieldType type : SUPPORTED_TYPES) { + caps.add(AggregateCapability.simple(func, type, DATAFUSION_FORMATS)); + } + } + AGG_CAPS = caps; } + @Override public String name() { return NAME; } + + @Override public SearchExecEngine searcher(ExecutionContext ctx) { return null; } + @Override public List getSupportedFormats() { return List.of(new DataFormat() { - @Override public String name() { return "parquet"; } + @Override public String name() { return PARQUET_DATA_FORMAT; } @Override public long priority() { return 0; } - @Override public Set supportedFields() { return Set.of(); } + @Override public Set supportedFields() { + return Set.of( + new FieldTypeCapabilities("integer", Set.of(COLUMNAR_STORAGE)), + new FieldTypeCapabilities("long", Set.of(COLUMNAR_STORAGE)), + new FieldTypeCapabilities("keyword", Set.of(COLUMNAR_STORAGE)), + new FieldTypeCapabilities("text", Set.of(COLUMNAR_STORAGE)), + new FieldTypeCapabilities("boolean", Set.of(COLUMNAR_STORAGE)), + new FieldTypeCapabilities("date", Set.of(COLUMNAR_STORAGE)) + ); + } }); } - @Override - public Set supportedFilterCapabilities() { - return FILTER_CAPS; - } - - @Override - public Set supportedFullTextOperators() { - return Set.of(); - } - - @Override - public Set supportedOperators() { - return OPERATOR_CAPS; - } - - @Override - public Set supportedAggregateFunctions() { - return AGG_FUNCTIONS; - } + @Override public Set supportedOperators() { return OPERATOR_CAPS; } + @Override public Set filterCapabilities() { return FILTER_CAPS; } + @Override public Set aggregateCapabilities() { return AGG_CAPS; } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java index e31534e768f0e..4b34e36e93dec 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java @@ -11,35 +11,39 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; -import org.opensearch.analytics.spi.AggregateFunction; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; -import org.opensearch.analytics.spi.FieldTypeFamily; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.FilterCapability; import org.opensearch.analytics.spi.FilterOperator; -import org.opensearch.analytics.spi.FullTextOperator; import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; -import java.util.EnumSet; +import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.FULL_TEXT_SEARCH; +import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.POINT_RANGE; +import static org.opensearch.index.engine.dataformat.FieldTypeCapabilities.Capability.STORED_FIELDS; + +import java.util.HashSet; import java.util.List; import java.util.Set; -import java.util.stream.Collectors; /** - * Mock Lucene backend for tests. Supports lucene format, standard filter - * operators on NUMERIC/KEYWORD/TEXT/DATE/BOOLEAN families, all full-text - * operators, and SCAN + FILTER capabilities (no AGGREGATE). + * Mock Lucene backend for tests. Supports lucene format with index structures + * (full-text, point range) and stored fields. No columnar storage (doc values + * only in parquet for default tests). Standard + full-text filter capabilities. + * SCAN + FILTER only (no AGGREGATE). */ public class MockLuceneBackend implements AnalyticsSearchBackendPlugin { public static final String NAME = "mock-lucene"; + public static final String LUCENE_DATA_FORMAT = "lucene"; + private static final Set LUCENE_FORMATS = Set.of(LUCENE_DATA_FORMAT); private static final Set OPERATOR_CAPS = Set.of( OperatorCapability.SCAN, OperatorCapability.FILTER ); - private static final Set FILTER_OPS = EnumSet.of( + private static final Set STANDARD_OPS = Set.of( FilterOperator.EQUALS, FilterOperator.NOT_EQUALS, FilterOperator.GREATER_THAN, FilterOperator.GREATER_THAN_OR_EQUAL, FilterOperator.LESS_THAN, FilterOperator.LESS_THAN_OR_EQUAL, @@ -47,54 +51,67 @@ public class MockLuceneBackend implements AnalyticsSearchBackendPlugin { FilterOperator.IN, FilterOperator.LIKE ); - private static final Set SUPPORTED_FAMILIES = EnumSet.of( - FieldTypeFamily.NUMERIC, FieldTypeFamily.KEYWORD, - FieldTypeFamily.TEXT, FieldTypeFamily.DATE, FieldTypeFamily.BOOLEAN + private static final Set FULL_TEXT_OPS = Set.of( + FilterOperator.MATCH, FilterOperator.MATCH_PHRASE, + FilterOperator.MATCH_PHRASE_PREFIX, FilterOperator.MATCH_BOOL_PREFIX, + FilterOperator.MULTI_MATCH, FilterOperator.QUERY_STRING, + FilterOperator.SIMPLE_QUERY_STRING, FilterOperator.FUZZY, + FilterOperator.WILDCARD, FilterOperator.REGEXP ); - private static final Set FILTER_CAPS; + private static final Set STANDARD_TYPES = new HashSet<>(); static { - FILTER_CAPS = FILTER_OPS.stream() - .flatMap(op -> SUPPORTED_FAMILIES.stream().map(family -> FilterCapability.of(op, family))) - .collect(Collectors.toUnmodifiableSet()); + STANDARD_TYPES.addAll(FieldType.numeric()); + STANDARD_TYPES.addAll(FieldType.keyword()); + STANDARD_TYPES.addAll(FieldType.text()); + STANDARD_TYPES.addAll(FieldType.date()); + STANDARD_TYPES.add(FieldType.BOOLEAN); } - @Override - public String name() { - return NAME; + private static final Set FULL_TEXT_TYPES = new HashSet<>(); + static { + FULL_TEXT_TYPES.addAll(FieldType.keyword()); + FULL_TEXT_TYPES.addAll(FieldType.text()); } - @Override - public SearchExecEngine searcher(ExecutionContext ctx) { - return null; + private static final Set FILTER_CAPS; + static { + Set caps = new HashSet<>(); + for (FilterOperator op : STANDARD_OPS) { + for (FieldType type : STANDARD_TYPES) { + caps.add(new FilterCapability.Standard(op, type, LUCENE_FORMATS)); + } + } + for (FilterOperator op : FULL_TEXT_OPS) { + for (FieldType type : FULL_TEXT_TYPES) { + caps.add(new FilterCapability.FullText(op, type, LUCENE_FORMATS, Set.of())); + } + } + FILTER_CAPS = caps; } + @Override public String name() { return NAME; } + + @Override public SearchExecEngine searcher(ExecutionContext ctx) { return null; } + @Override public List getSupportedFormats() { return List.of(new DataFormat() { - @Override public String name() { return "lucene"; } + @Override public String name() { return LUCENE_DATA_FORMAT; } @Override public long priority() { return 0; } - @Override public Set supportedFields() { return Set.of(); } + @Override public Set supportedFields() { + return Set.of( + new FieldTypeCapabilities("integer", Set.of(POINT_RANGE, STORED_FIELDS)), + new FieldTypeCapabilities("long", Set.of(POINT_RANGE, STORED_FIELDS)), + new FieldTypeCapabilities("keyword", Set.of(FULL_TEXT_SEARCH, STORED_FIELDS)), + new FieldTypeCapabilities("text", Set.of(FULL_TEXT_SEARCH, STORED_FIELDS)), + new FieldTypeCapabilities("boolean", Set.of(STORED_FIELDS)), + new FieldTypeCapabilities("date", Set.of(POINT_RANGE, STORED_FIELDS)) + ); + } }); } - @Override - public Set supportedFilterCapabilities() { - return FILTER_CAPS; - } - - @Override - public Set supportedFullTextOperators() { - return EnumSet.allOf(FullTextOperator.class); - } - - @Override - public Set supportedOperators() { - return OPERATOR_CAPS; - } - - @Override - public Set supportedAggregateFunctions() { - return Set.of(); - } + @Override public Set supportedOperators() { return OPERATOR_CAPS; } + @Override public Set filterCapabilities() { return FILTER_CAPS; } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java index 1b67296191aeb..ef1d1b8c37ba5 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ProjectRuleTests.java @@ -24,9 +24,12 @@ import org.opensearch.analytics.planner.rel.OpenSearchProject; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegationType; +import org.opensearch.analytics.spi.FieldType; +import org.opensearch.analytics.spi.ProjectCapability; import org.opensearch.analytics.spi.ScalarFunction; import java.util.EnumSet; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; @@ -55,7 +58,7 @@ public void testSimpleFieldProjection() { rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0), rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) ); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); for (RexNode expr : result.getProjects()) { assertFalse("Field ref should not be annotated", expr instanceof AnnotatedProjectExpression); } @@ -70,7 +73,7 @@ public void testSupportedScalarFunction() { rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) ); OpenSearchProject result = runProject(castExpr); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); } @@ -90,8 +93,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); } }; @@ -99,9 +102,14 @@ public Set supportedOpaqueProjectOperations() { rexBuilder.makeCall(PAINLESS, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - // DataFusion is child backend, delegates painless to Lucene - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); - assertAnnotation(result.getProjects().get(0), MockLuceneBackend.NAME); + // Operator: only DF viable (Lucene not viable for scan) + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); + assertFalse(result.getViableBackends().contains(MockLuceneBackend.NAME)); + // painless expression: only Lucene can natively evaluate it + AnnotatedProjectExpression annotation = (AnnotatedProjectExpression) result.getProjects().get(0); + assertTrue(annotation.getOriginal().toString().contains("painless")); + assertTrue(annotation.getViableBackends().contains(MockLuceneBackend.NAME)); + assertFalse(annotation.getViableBackends().contains(MockDataFusionBackend.NAME)); } /** Painless on DataFusion WITHOUT delegation → error. */ @@ -109,8 +117,8 @@ public void testPainlessErrorsWithoutDelegation() { // Lucene supports painless but no delegation configured MockLuceneBackend luceneWithPainless = new MockLuceneBackend() { @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); } }; @@ -144,8 +152,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); } }; @@ -154,7 +162,7 @@ public Set supportedOpaqueProjectOperations() { rexBuilder.makeCall(PAINLESS, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); assertFalse("Field ref should not be annotated", result.getProjects().get(0) instanceof AnnotatedProjectExpression); assertAnnotation(result.getProjects().get(1), MockLuceneBackend.NAME); @@ -174,8 +182,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless", "highlight"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless", "highlight"); } }; @@ -183,7 +191,7 @@ public Set supportedOpaqueProjectOperations() { rexBuilder.makeCall(HIGHLIGHT, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); assertAnnotation(result.getProjects().get(0), MockLuceneBackend.NAME); } @@ -191,7 +199,7 @@ public Set supportedOpaqueProjectOperations() { /** Backend does not support CAST → error. */ public void testUnsupportedScalarFunctionErrors() { - // Default MockDataFusionBackend has empty supportedScalarFunctions + // Default MockDataFusionBackend has empty projectCapabilities RexNode castExpr = rexBuilder.makeCast( typeFactory.createSqlType(SqlTypeName.VARCHAR), rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) @@ -209,7 +217,7 @@ public void testUnsupportedScalarFunctionErrors() { IllegalStateException exception = expectThrows(IllegalStateException.class, () -> runPlanner(project, context)); - assertTrue(exception.getMessage().contains("does not support scalar function")); + assertTrue(exception.getMessage().contains("No backend supports scalar function")); } // ---- Opaque operation natively supported ---- @@ -218,12 +226,10 @@ public void testUnsupportedScalarFunctionErrors() { public void testOpaqueOperationSupportedNatively() { MockDataFusionBackend dfWithPainless = new MockDataFusionBackend() { @Override - public Set supportedScalarFunctions() { - return EnumSet.allOf(ScalarFunction.class); - } - @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return combine( + scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), EnumSet.allOf(ScalarFunction.class)), + opaqueCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), "painless")); } }; @@ -231,7 +237,7 @@ public Set supportedOpaqueProjectOperations() { rexBuilder.makeCall(PAINLESS, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 0)) ); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); } @@ -249,7 +255,7 @@ public void testNestedScalarFunctions() { rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 1) ); OpenSearchProject result = runProject(plusExpr); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); } @@ -263,8 +269,9 @@ public Set supportedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedScalarFunctions() { - return EnumSet.allOf(ScalarFunction.class); + public Set projectCapabilities() { + return scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), + EnumSet.allOf(ScalarFunction.class)); } }; MockLuceneBackend luceneAccepting = new MockLuceneBackend() { @@ -273,8 +280,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); } }; @@ -288,7 +295,7 @@ public Set supportedOpaqueProjectOperations() { OpenSearchProject result = runProject("parquet", List.of(dfWithDelegation, luceneAccepting), fieldRef, painlessExpr, castExpr); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); // field ref — no annotation assertFalse("Field ref should not be annotated", result.getProjects().get(0) instanceof AnnotatedProjectExpression); @@ -306,8 +313,9 @@ public Set supportedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedScalarFunctions() { - return EnumSet.allOf(ScalarFunction.class); + public Set projectCapabilities() { + return scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), + EnumSet.allOf(ScalarFunction.class)); } }; MockLuceneBackend luceneAccepting = new MockLuceneBackend() { @@ -316,8 +324,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); } }; @@ -331,7 +339,7 @@ public Set supportedOpaqueProjectOperations() { OpenSearchProject result = runProject("parquet", List.of(dfWithDelegation, luceneAccepting), plusExpr); - assertEquals(MockDataFusionBackend.NAME, result.getBackend()); + assertTrue(result.getViableBackends().contains(MockDataFusionBackend.NAME)); // outer PLUS annotated with DF assertAnnotation(result.getProjects().get(0), MockDataFusionBackend.NAME); // CAST(painless($0)) — CAST annotated with DF, painless inside annotated with Lucene @@ -364,8 +372,8 @@ public Set acceptedDelegations() { return Set.of(DelegationType.PROJECT); } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless", "highlight"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless", "highlight"); } }; @@ -386,8 +394,8 @@ public Set supportedOpaqueProjectOperations() { MockLuceneBackend thirdBackend = new MockLuceneBackend() { @Override public String name() { return "mock-third"; } @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("suggest"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "suggest"); } }; @@ -410,8 +418,8 @@ public Set supportedDelegations() { }; MockLuceneBackend luceneWithPainlessButNoAccept = new MockLuceneBackend() { @Override - public Set supportedOpaqueProjectOperations() { - return Set.of("painless"); + public Set projectCapabilities() { + return opaqueCaps(Set.of(MockLuceneBackend.LUCENE_DATA_FORMAT), "painless"); } // acceptedDelegations() returns empty by default }; @@ -437,7 +445,7 @@ public Set supportedOpaqueProjectOperations() { private void assertAnnotation(RexNode expr, String expectedBackend) { assertTrue("Expected AnnotatedProjectExpression, got " + expr.getClass().getSimpleName(), expr instanceof AnnotatedProjectExpression); - assertEquals(expectedBackend, ((AnnotatedProjectExpression) expr).getBackend()); + assertTrue(((AnnotatedProjectExpression) expr).getViableBackends().contains(expectedBackend)); } private OpenSearchProject runProject(RexNode... exprs) { @@ -472,9 +480,37 @@ private OpenSearchProject runProject(String format, List supportedScalarFunctions() { - return EnumSet.allOf(ScalarFunction.class); + public Set projectCapabilities() { + return scalarCaps(Set.of(MockDataFusionBackend.PARQUET_DATA_FORMAT), + EnumSet.allOf(ScalarFunction.class)); } }; } + + @SafeVarargs + private static Set combine(Set... sets) { + Set result = new HashSet<>(); + for (Set set : sets) { + result.addAll(set); + } + return result; + } + + private static Set scalarCaps(Set formats, Set functions) { + Set caps = new HashSet<>(); + for (ScalarFunction func : functions) { + for (FieldType type : FieldType.values()) { + caps.add(new ProjectCapability.Scalar(func, type, formats)); + } + } + return caps; + } + + private static Set opaqueCaps(Set formats, String... names) { + Set caps = new HashSet<>(); + for (String name : names) { + caps.add(new ProjectCapability.Opaque(name, formats)); + } + return caps; + } } From dccc254e0b0fdf667e01dc7e82676501e190264c Mon Sep 17 00:00:00 2001 From: expani Date: Fri, 3 Apr 2026 12:25:39 -0700 Subject: [PATCH 5/6] Added support for DAG creation from Marked Plan Signed-off-by: expani --- .../planner/FieldStorageResolver.java | 15 ++ .../analytics/planner/PlannerImpl.java | 15 +- .../analytics/planner/dag/DAGBuilder.java | 159 ++++++++++++++++++ .../analytics/planner/dag/ExchangeInfo.java | 37 ++++ .../analytics/planner/dag/QueryDAG.java | 61 +++++++ .../analytics/planner/dag/Stage.java | 57 +++++++ .../planner/rel/OpenSearchStageInputScan.java | 69 ++++++++ .../planner/rules/OpenSearchProjectRule.java | 5 + 8 files changed, 416 insertions(+), 2 deletions(-) create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ExchangeInfo.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/QueryDAG.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java index bfef4645ce79b..4858ce1d788c8 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/FieldStorageResolver.java @@ -39,6 +39,21 @@ public class FieldStorageResolver { /** * Production: resolves per-field storage from IndexMetadata and backend capabilities. + * + *

LIMITATION: This infers storage from what each DataFormat DECLARES it can support + * (via FieldTypeCapabilities), not from what the index ACTUALLY stores. A format declaring + * COLUMNAR_STORAGE for "integer" doesn't mean this index has integer doc values in that format. + * + *

The indexing side has no per-field format metadata at the coordinator level today: + * - DataFormat.supportedFields() = capability, not actual storage + * - Segment.dfGroupedSearchableFiles = per-segment format info, data node only + * - DataFormatRegistry has TODOs to filter by index settings/mapper service + * - primary_data_format index setting is the only coordinator-level hint (index-level, not field-level) + * + *

TODO: Replace inference with actual per-field format metadata once the indexing team adds + * it to MappingMetadata or IndexMetadata. Until then, this over-estimates viable backends — + * the shard-level cost function must handle the mismatch. Consider using primary_data_format + * as a narrowing hint in the interim. */ @SuppressWarnings("unchecked") public FieldStorageResolver(IndexMetadata indexMetadata, List backends) { 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 f1ce9de977e7c..c2c7c902603a6 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 @@ -14,6 +14,8 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef; +import org.opensearch.analytics.planner.dag.DAGBuilder; +import org.opensearch.analytics.planner.dag.QueryDAG; import org.opensearch.analytics.planner.rules.OpenSearchAggregateRule; import org.opensearch.analytics.planner.rules.OpenSearchAggregateSplitRule; import org.opensearch.analytics.planner.rules.OpenSearchFilterRule; @@ -21,22 +23,26 @@ import org.opensearch.analytics.planner.rules.OpenSearchSortRule; import org.opensearch.analytics.planner.rules.OpenSearchTableScanRule; +import org.opensearch.analytics.planner.dag.DAGBuilder; +import org.opensearch.analytics.planner.dag.QueryDAG; + import java.util.List; /** * Central planner for the Analytics Plugin. * - *

Two phases: + *

Three phases: *

    *
  1. HepPlanner (RBO): converts LogicalXxx → OpenSearchXxx with backend * assignment, predicate annotation, and distribution traits.
    2. *
  2. VolcanoPlanner (CBO): requests SINGLETON at root (coordinator must * gather all results). Split rule fires on aggregates, Volcano inserts * exchanges via trait enforcement where distribution mismatches.
    3. + *
  3. DAG construction: cuts at exchange boundaries, builds stage tree.
    4. *
* *

TODO: eliminate copyToCluster — have frontends create RelNodes with Volcano cluster. - *

TODO: DAG construction (cut at exchange boundaries) + *

TODO: Per-stage plan forking (multiple plan generation) *

TODO: Fragment conversion (backend.convertFragment) *

TODO: Join strategy selection, sort removal via CBO * @@ -93,6 +99,11 @@ public static RelNode createPlan(RelNode rawRelNode, PlannerContext context) { RelNode result = volcanoPlanner.findBestExp(); LOGGER.info("After CBO:\n{}", RelOptUtil.toString(result)); + + // Phase 3: DAG construction — cut at exchange boundaries + QueryDAG dag = DAGBuilder.build(result); + LOGGER.info("QueryDAG:\n{}", dag); + return result; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java new file mode 100644 index 0000000000000..861c88a4c7109 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java @@ -0,0 +1,159 @@ +/* + * 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.dag; + +import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchExchangeWriter; +import org.opensearch.analytics.planner.rel.OpenSearchShuffleReader; +import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; + +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; + +/** + * Builds a {@link QueryDAG} from the CBO output by cutting at exchange boundaries. + * + *

SINGLETON: {@link OpenSearchExchangeReducer} is the boundary. Parent + * fragment is everything above the reducer (may be null for a pure gather). + * Child fragment is the reducer's input subtree. + * + *

HASH/RANGE: {@link OpenSearchShuffleReader} → {@link OpenSearchExchangeWriter}. + * ShuffleReader stays in parent as leaf (input severed). Writer + subtree + * below becomes the child fragment. + * + *

Stage IDs assigned bottom-up (leaf stages get lower IDs). + * + * @opensearch.internal + */ +public class DAGBuilder { + + private DAGBuilder() {} + + public static QueryDAG build(RelNode cboOutput) { + int[] counter = { 0 }; + List childStages = new ArrayList<>(); + + RelNode fragment; + if (cboOutput instanceof OpenSearchExchangeReducer reducer) { + // Root is an ExchangeReducer (e.g., shuffle case where coordinator + // is a pure gather). ExchangeReducer becomes the root stage fragment + // with its input severed. Child stage is the subtree below. + fragment = cutSingleton(reducer, counter, childStages); + } else { + fragment = sever(cboOutput, counter, childStages); + } + + Stage rootStage = new Stage(counter[0]++, fragment, childStages, null); + return new QueryDAG(UUID.randomUUID().toString(), rootStage); + } + + /** + * Walks top-down collecting operators into the current stage. When an + * exchange boundary is hit, severs the link: subtree below becomes a + * child stage, current stage continues with the boundary removed or + * replaced by a detached leaf. + * + *

Exchange boundaries are detected at the input level (not the node + * level) so the parent node never receives a null input. For SINGLETON + * cuts the input is simply dropped; for shuffle cuts the ShuffleReader + * replaces the input as a detached leaf. + * + * @param node current node being visited + * @param counter stage ID counter (bottom-up assignment) + * @param childStages accumulator for child stages found below exchanges + * @return the rewritten fragment root for the current stage + */ + private static RelNode sever(RelNode node, int[] counter, List childStages) { + // Check each input for exchange boundaries before recursing + List newInputs = new ArrayList<>(); + for (RelNode input : node.getInputs()) { + if (input instanceof OpenSearchExchangeReducer reducer) { + // SINGLETON cut: ExchangeReducer stays in parent as leaf (input severed). + // Child stage is the reducer's input subtree. Analytics Core streams + // results from data nodes — no exchange operator needed in child fragment. + newInputs.add(cutSingleton(reducer, counter, childStages)); + } else if (input instanceof OpenSearchShuffleReader reader) { + // Shuffle cut: ShuffleReader stays in parent as leaf (input severed). + // Child stage is ExchangeWriter + subtree below. + newInputs.add(cutShuffle(reader, counter, childStages)); + } else { + newInputs.add(sever(input, counter, childStages)); + } + } + + // Leaf node (e.g., TableScan) — no inputs to process + if (node.getInputs().isEmpty()) { + return node; + } + + // Rebuild only if inputs changed + boolean changed = false; + for (int idx = 0; idx < newInputs.size(); idx++) { + if (newInputs.get(idx) != node.getInputs().get(idx)) { + changed = true; + break; + } + } + return changed ? node.copy(node.getTraitSet(), newInputs) : node; + } + + private static RelNode cutSingleton(OpenSearchExchangeReducer reducer, + int[] counter, List parentChildStages) { + List grandchildren = new ArrayList<>(); + RelNode childFragment = sever(reducer.getInput(), counter, grandchildren); + + int childStageId = counter[0]++; + parentChildStages.add(new Stage( + childStageId, childFragment, grandchildren, + new ExchangeInfo(RelDistribution.Type.SINGLETON, null, List.of()) + )); + + // Replace child subtree with StageInputScan, keep ExchangeReducer in parent + RelDataType childRowType = reducer.getInput().getRowType(); + OpenSearchStageInputScan stageInput = new OpenSearchStageInputScan( + reducer.getCluster(), reducer.getTraitSet(), childStageId, childRowType + ); + return new OpenSearchExchangeReducer( + reducer.getCluster(), reducer.getTraitSet(), stageInput, reducer.getViableBackends() + ); + } + + private static RelNode cutShuffle(OpenSearchShuffleReader reader, + int[] counter, List parentChildStages) { + if (!(reader.getInput() instanceof OpenSearchExchangeWriter writer)) { + throw new IllegalStateException( + "ShuffleReader input must be ExchangeWriter, got: " + + reader.getInput().getClass().getSimpleName()); + } + + List grandchildren = new ArrayList<>(); + RelNode belowWriter = sever(writer.getInput(), counter, grandchildren); + RelNode childFragment = writer.copy(writer.getTraitSet(), List.of(belowWriter)); + + int childStageId = counter[0]++; + parentChildStages.add(new Stage( + childStageId, childFragment, grandchildren, + new ExchangeInfo(RelDistribution.Type.HASH_DISTRIBUTED, writer.getShuffleImpl(), writer.getKeys()) + )); + + // Replace child subtree with StageInputScan, keep ShuffleReader in parent + RelDataType childRowType = writer.getInput().getRowType(); + OpenSearchStageInputScan stageInput = new OpenSearchStageInputScan( + reader.getCluster(), reader.getTraitSet(), childStageId, childRowType + ); + return new OpenSearchShuffleReader( + reader.getCluster(), reader.getTraitSet(), stageInput, + reader.getViableBackends(), reader.getShuffleImpl() + ); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ExchangeInfo.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ExchangeInfo.java new file mode 100644 index 0000000000000..d71ad6a017123 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ExchangeInfo.java @@ -0,0 +1,37 @@ +/* + * 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.dag; + +import org.apache.calcite.rel.RelDistribution; +import org.opensearch.analytics.planner.rel.ShuffleImpl; + +import java.util.List; + +/** + * Exchange metadata extracted from exchange RelNodes during DAG construction. + * Describes how a child stage delivers data to its parent stage. + * + * @param distributionType distribution type from the exchange operator's trait + * @param shuffleImpl shuffle implementation (null for SINGLETON) + * @param partitionKeyIndices field indices for hash/range partitioning (empty for SINGLETON) + * @opensearch.internal + */ +public record ExchangeInfo( + RelDistribution.Type distributionType, + ShuffleImpl shuffleImpl, + List partitionKeyIndices +) { + public ExchangeInfo { + partitionKeyIndices = List.copyOf(partitionKeyIndices); + } + + public boolean isShuffle() { + return shuffleImpl != null; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/QueryDAG.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/QueryDAG.java new file mode 100644 index 0000000000000..b2d4f17982d38 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/QueryDAG.java @@ -0,0 +1,61 @@ +/* + * 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.dag; + +import org.apache.calcite.plan.RelOptUtil; + +/** + * Root of the query execution DAG. Recursive tree of {@link Stage}s. + * + * @param queryId unique identifier for this query. Currently a random UUID. + * Future: accept a user-provided ID or generate a cluster-wide + * unique ID (e.g., node-local counter + nodeId prefix). + * @param rootStage the coordinator/root stage + * @opensearch.internal + */ +public record QueryDAG(String queryId, Stage rootStage) { + + /** + * Returns a human-readable representation of the entire DAG showing + * every stage with its fragment indented. + */ + @Override + public String toString() { + StringBuilder builder = new StringBuilder(); + builder.append("QueryDAG(queryId=").append(queryId).append(")\n"); + appendStage(builder, rootStage, 0); + return builder.toString(); + } + + private static void appendStage(StringBuilder builder, Stage stage, int depth) { + String indent = " ".repeat(depth); + builder.append(indent).append("Stage ").append(stage.getStageId()); + if (stage.getExchangeInfo() != null) { + builder.append(" [exchange=").append(stage.getExchangeInfo()).append("]"); + } else { + builder.append(" [root]"); + } + builder.append("\n"); + + if (stage.getFragment() != null) { + String fragmentStr = RelOptUtil.toString(stage.getFragment()); + for (String line : fragmentStr.split("\n")) { + if (!line.isEmpty()) { + builder.append(indent).append(" ").append(line).append("\n"); + } + } + } else { + builder.append(indent).append(" \n"); + } + + for (Stage child : stage.getChildStages()) { + appendStage(builder, child, depth + 1); + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java new file mode 100644 index 0000000000000..255c173f5e416 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java @@ -0,0 +1,57 @@ +/* + * 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.dag; + +import org.apache.calcite.rel.RelNode; + +import java.util.List; + +/** + * A stage in the query DAG. Each stage holds a plan fragment (the marked RelNode + * subtree between exchange boundaries) and references to child stages. + * + *

Fragment may be null for a pure gather stage (coordinator just accumulates + * Arrow batches from child stages). + * + * @opensearch.internal + */ +public class Stage { + + private final int stageId; + private final RelNode fragment; + private final List childStages; + private final ExchangeInfo exchangeInfo; + + // TODO: add List planAlternatives — populated during plan forking phase + + public Stage(int stageId, RelNode fragment, List childStages, ExchangeInfo exchangeInfo) { + this.stageId = stageId; + this.fragment = fragment; + this.childStages = List.copyOf(childStages); + this.exchangeInfo = exchangeInfo; + } + + public int getStageId() { + return stageId; + } + + /** Marked plan fragment with annotations intact. Null for a pure gather stage. */ + public RelNode getFragment() { + return fragment; + } + + public List getChildStages() { + return childStages; + } + + /** How this stage connects to its parent. Null for the root stage. */ + public ExchangeInfo getExchangeInfo() { + return exchangeInfo; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java new file mode 100644 index 0000000000000..7b1eb24056f84 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java @@ -0,0 +1,69 @@ +/* + * 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.rel; + +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptCost; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.rel.AbstractRelNode; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelWriter; +import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; + +import java.util.List; + +/** + * Leaf node representing data arriving from a child stage via exchange. + * Replaces the severed child subtree in the parent stage fragment during + * DAG construction. + * + *

Carries the child stage ID and the row type of the child subtree's + * output. During fragment conversion (RelNode → Substrait/QueryBuilder), + * the backend registers a streaming source or shuffle source for this node. + * + * @opensearch.internal + */ +public class OpenSearchStageInputScan extends AbstractRelNode { + + private final int childStageId; + private final RelDataType rowType; + + public OpenSearchStageInputScan(RelOptCluster cluster, RelTraitSet traitSet, + int childStageId, RelDataType rowType) { + super(cluster, traitSet); + this.childStageId = childStageId; + this.rowType = rowType; + } + + public int getChildStageId() { + return childStageId; + } + + @Override + protected RelDataType deriveRowType() { + return rowType; + } + + @Override + public RelNode copy(RelTraitSet traitSet, List inputs) { + return new OpenSearchStageInputScan(getCluster(), traitSet, childStageId, rowType); + } + + @Override + public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { + return planner.getCostFactory().makeZeroCost(); + } + + @Override + public RelWriter explainTerms(RelWriter pw) { + return super.explainTerms(pw).item("childStageId", childStageId); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java index 161291de1f85a..73317d87bc6a3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java @@ -23,6 +23,7 @@ import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.FieldType; +import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.analytics.spi.ScalarFunction; import java.util.ArrayList; @@ -190,7 +191,11 @@ private List computeProjectViableBackends(List annotatedExprs, List delegationAcceptors = registry.delegationAcceptors(DelegationType.PROJECT); List result = new ArrayList<>(); + List projectCapable = registry.operatorBackends(OperatorCapability.PROJECT); for (String candidateName : childViableBackends) { + if (!projectCapable.contains(candidateName)) { + continue; + } boolean canHandleAll = true; for (RexNode expr : annotatedExprs) { if (!(expr instanceof AnnotatedProjectExpression annotation)) { From 4c1524c678e79b6f38858eea8e667e5dafd68bd2 Mon Sep 17 00:00:00 2001 From: expani Date: Fri, 3 Apr 2026 14:17:45 -0700 Subject: [PATCH 6/6] Rebased with internal working branch for search marc/ani Signed-off-by: expani --- .../spi/AnalyticsSearchBackendPlugin.java | 15 +++++++++---- .../exec/DefaultPlanExecutorTests.java | 22 +++++++++++++++++++ .../planner/MockDataFusionBackend.java | 2 +- .../analytics/planner/MockLuceneBackend.java | 2 +- 4 files changed, 35 insertions(+), 6 deletions(-) diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index ef7e5cbf00332..373fcde77b75f 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java @@ -8,7 +8,6 @@ package org.opensearch.analytics.spi; - import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; @@ -21,6 +20,11 @@ /** * SPI extension point for back-end query engines for query planning and execution capabilities * as needed by the {@link org.opensearch.analytics.exec.QueryPlanExecutor} + * + *

TODO: separate capability declaration (planner, coordinator) from execution engine factory + * (data node) into two interfaces. AnalyticsSearchBackendPlugin should only declare capabilities. + * SearchExecEngineProvider should be discovered separately by the executor. Remove the extends + * relationship and the default createSearchExecEngine() below once that separation is done. */ public interface AnalyticsSearchBackendPlugin extends SearchExecEngineProvider { @@ -28,10 +32,13 @@ public interface AnalyticsSearchBackendPlugin extends SearchExecEngineProvider { String name(); /** - * Creates a searcher bound to the given reader snapshot. - * @param ctx the execution context + * {@inheritDoc} + * Temporary default — remove once SearchExecEngineProvider is separated from this interface. */ - SearchExecEngine searcher(ExecutionContext ctx); + @Override + default SearchExecEngine createSearchExecEngine(ExecutionContext ctx) { + throw new UnsupportedOperationException("createSearchExecEngine not implemented for " + name()); + } /** Returns the data formats supported by this backend. */ List getSupportedFormats(); diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java index 14d9dd2a17ed6..459a7f3786abd 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java @@ -26,11 +26,14 @@ import org.opensearch.analytics.backend.ExecutionContext; import org.opensearch.analytics.backend.SearchExecEngine; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.OperatorCapability; import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.cluster.metadata.MappingMetadata; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.common.settings.Settings; import org.opensearch.core.index.Index; import org.opensearch.index.IndexService; import org.opensearch.index.engine.DataFormatAwareEngine; @@ -123,6 +126,15 @@ public void testEndToEndExecuteWithMockBackend() throws IOException { Index index = new Index("my_index", "uuid"); IndexMetadata indexMetadata = mock(IndexMetadata.class); when(indexMetadata.getIndex()).thenReturn(index); + when(indexMetadata.getSettings()).thenReturn( + Settings.builder().put("index.composite.primary_data_format", "mock-columnar").build() + ); + when(indexMetadata.getNumberOfShards()).thenReturn(1); + MappingMetadata mappingMetadata = mock(MappingMetadata.class); + when(mappingMetadata.sourceAsMap()).thenReturn( + Map.of("properties", Map.of("field_0", Map.of("type", "integer"))) + ); + when(indexMetadata.mapping()).thenReturn(mappingMetadata); Metadata metadata = mock(Metadata.class); when(metadata.index("my_index")).thenReturn(indexMetadata); @@ -398,6 +410,16 @@ public String name() { return "mock-backend"; } + @Override + public List getSupportedFormats() { + return List.of(format); + } + + @Override + public Set supportedOperators() { + return Set.of(OperatorCapability.SCAN, OperatorCapability.FILTER); + } + @Override public SearchExecEngine createSearchExecEngine(ExecutionContext ctx) { Object reader = ctx.getReader().reader(format); 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 8335338c14315..c7e36ef7789ec 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 @@ -90,7 +90,7 @@ public class MockDataFusionBackend implements AnalyticsSearchBackendPlugin { @Override public String name() { return NAME; } - @Override public SearchExecEngine searcher(ExecutionContext ctx) { return null; } + @Override public SearchExecEngine createSearchExecEngine(ExecutionContext ctx) { return null; } @Override public List getSupportedFormats() { diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java index 4b34e36e93dec..7828ae9f46f59 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java @@ -92,7 +92,7 @@ public class MockLuceneBackend implements AnalyticsSearchBackendPlugin { @Override public String name() { return NAME; } - @Override public SearchExecEngine searcher(ExecutionContext ctx) { return null; } + @Override public SearchExecEngine createSearchExecEngine(ExecutionContext ctx) { return null; } @Override public List getSupportedFormats() {