diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index 81b8abb27b0..ea9c74bbb7f 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -661,52 +661,34 @@ private boolean isMetadataField(String fieldName) { return OpenSearchConstants.METADATAFIELD_TYPE_MAP.containsKey(fieldName); } - /** See logic in {@link org.opensearch.sql.analysis.symbol.SymbolTable#lookupAllFields} */ + /** + * See logic in {@link org.opensearch.sql.analysis.symbol.SymbolTable#lookupAllFields}. + * + *

Single pass over field ordinals: index patterns with dynamic sub-fields can yield 100k+ + * columns, so resolving each field by name ({@code RelBuilder#field(String)} is a linear scan) or + * removing from a list per field makes this quadratic. + */ private static void tryToRemoveNestedFields(CalcitePlanContext context) { - Set allFields = new HashSet<>(context.relBuilder.peek().getRowType().getFieldNames()); - List duplicatedNestedFields = - allFields.stream() - .filter( - field -> { - int lastDot = field.lastIndexOf("."); - return -1 != lastDot && allFields.contains(field.substring(0, lastDot)); - }) - .map(field -> (RexNode) context.relBuilder.field(field)) - .toList(); - if (!duplicatedNestedFields.isEmpty()) { - // This is a workaround to avoid the bug in Calcite: + List fieldNames = context.relBuilder.peek().getRowType().getFieldNames(); + Set allFields = new HashSet<>(fieldNames); + List keptFields = new ArrayList<>(fieldNames.size()); + for (int i = 0; i < fieldNames.size(); i++) { + String field = fieldNames.get(i); + int lastDot = field.lastIndexOf('.'); + if (lastDot == -1 || !allFields.contains(field.substring(0, lastDot))) { + keptFields.add(context.relBuilder.field(i)); + } + } + if (keptFields.size() < fieldNames.size()) { + // Project with force=true to work around a Calcite bug: // In {@link RelBuilder#project_(Iterable, Iterable, Iterable, boolean, Iterable)}, // the check `RexUtil.isIdentity(nodeList, inputRowType)` will pass when the input // and the output nodeList refer to the same fields, even if the field name list // is different. As a result, renaming operation will not be applied. This makes // the logical plan for the flatten command incorrect, where the operation is // equivalent to renaming the flattened sub-fields. E.g. emp.name -> name. - forceProjectExcept(context.relBuilder, duplicatedNestedFields); - } - } - - /** - * Project except with force. - * - *

This method is copied from {@link RelBuilder#projectExcept(Iterable)} and modified with the - * force flag in project set to true. It is subject to future changes in Calcite. - * - * @param relBuilder RelBuilder - * @param expressions Expressions to exclude from the project - */ - private static void forceProjectExcept(RelBuilder relBuilder, Iterable expressions) { - List allExpressions = new ArrayList<>(relBuilder.fields()); - Set excludeExpressions = new HashSet<>(); - for (RexNode excludeExp : expressions) { - if (!excludeExpressions.add(excludeExp)) { - throw new IllegalArgumentException( - "Input list contains duplicates. Expression " + excludeExp + " exists multiple times."); - } - if (!allExpressions.remove(excludeExp)) { - throw new IllegalArgumentException("Expression " + excludeExp.toString() + " not found."); - } + context.relBuilder.project(keptFields, ImmutableList.of(), true); } - relBuilder.project(allExpressions, ImmutableList.of(), true); } /** diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLWideNestedSchemaTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLWideNestedSchemaTest.java new file mode 100644 index 00000000000..aa3ad0990c2 --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLWideNestedSchemaTest.java @@ -0,0 +1,143 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.calcite; + +import com.google.common.collect.ImmutableList; +import java.util.List; +import org.apache.calcite.DataContext; +import org.apache.calcite.config.CalciteConnectionConfig; +import org.apache.calcite.linq4j.Enumerable; +import org.apache.calcite.linq4j.Linq4j; +import org.apache.calcite.plan.RelTraitDef; +import org.apache.calcite.rel.RelCollations; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.schema.ScannableTable; +import org.apache.calcite.schema.Schema; +import org.apache.calcite.schema.SchemaPlus; +import org.apache.calcite.schema.Statistic; +import org.apache.calcite.schema.Statistics; +import org.apache.calcite.sql.SqlCall; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql.parser.SqlParser; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.test.CalciteAssert; +import org.apache.calcite.tools.Frameworks; +import org.apache.calcite.tools.Programs; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; + +/** + * The implicit {@code fields *} projection drops flattened sub-fields whose parent struct is also a + * column. Index patterns over many indices with dynamic sub-fields (e.g. {@code attributes.*}) can + * produce 100k+ columns, so this pass must be linear in the column count. + */ +public class CalcitePPLWideNestedSchemaTest extends CalcitePPLAbstractTest { + + private static final int WIDE_SUBFIELDS = 100_000; + + public CalcitePPLWideNestedSchemaTest() { + super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL); + } + + @Override + protected Frameworks.ConfigBuilder config(CalciteAssert.SchemaSpec... schemaSpecs) { + final SchemaPlus rootSchema = Frameworks.createRootSchema(true); + final SchemaPlus schema = CalciteAssert.addSchema(rootSchema, schemaSpecs); + schema.add("SMALL", new WideNestedTable(3)); + schema.add("WIDE", new WideNestedTable(WIDE_SUBFIELDS)); + return Frameworks.newConfigBuilder() + .parserConfig(SqlParser.Config.DEFAULT) + .defaultSchema(schema) + .traitDefs((List) null) + .programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2)); + } + + /** + * Sub-fields with an existing parent column are dropped (including multi-level ones), a dotted + * field without a parent column is kept, and the remaining columns keep their original order. + */ + @Test + public void testNestedSubFieldsRemovedAndOrderPreserved() { + RelNode root = getRelNode("source=SMALL"); + String expectedLogical = + "LogicalProject(ID=[$0], ATTRS=[$1], LONE.X=[$3], NAME=[$8])\n" + + " LogicalTableScan(table=[[scott, SMALL]])\n"; + verifyLogical(root, expectedLogical); + } + + /** Quadratic name lookups made this take ~36s at ~112k columns (V2429433093). */ + @Test(timeout = 10_000) + public void testWideNestedSchemaIsLinear() { + RelNode root = getRelNode("source=WIDE"); + String expectedLogical = + "LogicalProject(ID=[$0], ATTRS=[$1], LONE.X=[$3], NAME=[$" + + (WIDE_SUBFIELDS + 5) + + "])\n" + + " LogicalTableScan(table=[[scott, WIDE]])\n"; + verifyLogical(root, expectedLogical); + } + + /** + * Columns: ID, ATTRS (struct), ATTRS.SUB (struct), LONE.X (no LONE parent), ATTRS.SUB.X, then + * {@code subfields} flattened ATTRS.F* leaves, then NAME. + */ + static class WideNestedTable implements ScannableTable { + private final int subfields; + + WideNestedTable(int subfields) { + this.subfields = subfields; + } + + @Override + public RelDataType getRowType(RelDataTypeFactory factory) { + RelDataType varchar = factory.createSqlType(SqlTypeName.VARCHAR); + RelDataType sub = factory.createStructType(List.of(varchar), List.of("X")); + RelDataTypeFactory.Builder builder = + factory + .builder() + .add("ID", SqlTypeName.INTEGER) + .add("ATTRS", factory.createStructType(List.of(sub), List.of("SUB"))) + .add("ATTRS.SUB", sub) + .add("LONE.X", SqlTypeName.VARCHAR) + .add("ATTRS.SUB.X", SqlTypeName.VARCHAR); + for (int i = 0; i < subfields; i++) { + builder.add("ATTRS.F" + i, SqlTypeName.VARCHAR); + } + return builder.add("NAME", SqlTypeName.VARCHAR).build(); + } + + @Override + public Enumerable<@Nullable Object[]> scan(DataContext root) { + return Linq4j.asEnumerable(ImmutableList.of()); + } + + @Override + public Statistic getStatistic() { + return Statistics.of(0d, ImmutableList.of(), RelCollations.createSingleton(0)); + } + + @Override + public Schema.TableType getJdbcTableType() { + return Schema.TableType.TABLE; + } + + @Override + public boolean isRolledUp(String column) { + return false; + } + + @Override + public boolean rolledUpColumnValidInsideAgg( + String column, + SqlCall call, + @Nullable SqlNode parent, + @Nullable CalciteConnectionConfig config) { + return false; + } + } +}