Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.opensearch.sql.exception.CalciteUnsupportedException;
import org.opensearch.sql.exception.QueryEngineException;
import org.opensearch.sql.exception.SemanticCheckException;
import org.opensearch.sql.ppl.parser.PPLSearchPredicateCompiler;

/**
* {@code UnifiedQueryPlanner} provides a high-level API for parsing and analyzing queries using the
Expand Down Expand Up @@ -137,13 +138,15 @@ public RelNode plan(String query) throws Exception {
private static class CustomVisitorStrategy implements PlanningStrategy {
private final UnifiedQueryContext context;
private final UnifiedQueryParser<UnresolvedPlan> parser;
private final CalciteRelNodeVisitor relNodeVisitor =
new CalciteRelNodeVisitor(new EmptyDataSourceService());
private final CalciteRelNodeVisitor relNodeVisitor;

@SuppressWarnings("unchecked")
CustomVisitorStrategy(UnifiedQueryContext context) {
this.context = context;
this.parser = (UnifiedQueryParser<UnresolvedPlan>) context.getParser();
this.relNodeVisitor =
new CalciteRelNodeVisitor(
new EmptyDataSourceService(), PPLSearchPredicateCompiler.INSTANCE);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@
import org.opensearch.sql.ast.tree.Filter;
import org.opensearch.sql.ast.tree.Flatten;
import org.opensearch.sql.ast.tree.Foreach;
import org.opensearch.sql.ast.tree.Format;
import org.opensearch.sql.ast.tree.GraphLookup;
import org.opensearch.sql.ast.tree.Head;
import org.opensearch.sql.ast.tree.Join;
Expand Down Expand Up @@ -299,6 +300,9 @@ public LogicalPlan visitLimit(Limit node, AnalysisContext context) {

@Override
public LogicalPlan visitSearch(Search node, AnalysisContext context) {
if (node.hasImplicitSubquery()) {
throw getOnlyForCalciteException("Implicit format subsearch");
}
LogicalPlan child = node.getChild().get(0).accept(this, context);
Function queryStringFunc =
AstDSL.function(
Expand Down Expand Up @@ -833,6 +837,11 @@ public LogicalPlan visitReverse(Reverse node, AnalysisContext context) {
throw getOnlyForCalciteException("Reverse");
}

@Override
public LogicalPlan visitFormat(Format node, AnalysisContext context) {
throw getOnlyForCalciteException("Format");
}

@Override
public LogicalPlan visitSpath(SPath node, AnalysisContext context) {
throw getOnlyForCalciteException("Spath");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@
import org.opensearch.sql.ast.tree.Filter;
import org.opensearch.sql.ast.tree.Flatten;
import org.opensearch.sql.ast.tree.Foreach;
import org.opensearch.sql.ast.tree.Format;
import org.opensearch.sql.ast.tree.GraphLookup;
import org.opensearch.sql.ast.tree.Head;
import org.opensearch.sql.ast.tree.Join;
Expand Down Expand Up @@ -166,6 +167,10 @@ public T visitFlatten(Flatten node, C context) {
return visitChildren(node, context);
}

public T visitFormat(Format node, C context) {
return visitChildren(node, context);
}

public T visitTrendline(Trendline node, C context) {
return visitChildren(node, context);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.sql.ast.expression;

import java.util.List;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.ToString;
import org.opensearch.sql.ast.tree.UnresolvedPlan;

/** A subsearch used as a predicate term by the parent {@code search} command. */
@Getter
@RequiredArgsConstructor
@EqualsAndHashCode(callSuper = false)
@ToString
public class SearchSubquery extends SearchExpression {

private final UnresolvedPlan query;

@Override
public String toQueryString() {
throw new IllegalStateException(
"An implicit search subquery must be planned as a correlated runtime input");
}

@Override
public String toAnonymizedString() {
return "[ subsearch ]";
}

@Override
public List<? extends UnresolvedExpression> getChild() {
return List.of();
}
}
65 changes: 65 additions & 0 deletions core/src/main/java/org/opensearch/sql/ast/tree/Format.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.sql.ast.tree;

import com.google.common.collect.ImmutableList;
import java.util.List;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.Setter;
import lombok.ToString;
import org.opensearch.sql.ast.AbstractNodeVisitor;

/** AST node that collapses input rows into a formatted search expression. */
@Getter
@Setter
@ToString
@EqualsAndHashCode(callSuper = false)
@RequiredArgsConstructor
public class Format extends UnresolvedPlan {

public static final String DEFAULT_MV_SEPARATOR = "OR";
public static final int DEFAULT_MAX_RESULTS = 0;
public static final String DEFAULT_ROW_PREFIX = "(";
public static final String DEFAULT_COLUMN_PREFIX = "(";
public static final String DEFAULT_COLUMN_SEPARATOR = "AND";
public static final String DEFAULT_COLUMN_END = ")";
public static final String DEFAULT_ROW_SEPARATOR = "OR";
public static final String DEFAULT_ROW_END = ")";
public static final String DEFAULT_EMPTY_STRING = "NOT ()";

private final String mvSeparator;
private final int maxResults;
private final String rowPrefix;
private final String columnPrefix;
private final String columnSeparator;
private final String columnEnd;
private final String rowSeparator;
private final String rowEnd;
private final String emptyString;

/** Whether the result is consumed by an enclosing search rather than returned to the user. */
private boolean implicit;

private UnresolvedPlan child;

@Override
public Format attach(UnresolvedPlan child) {
this.child = child;
return this;
}

@Override
public List<UnresolvedPlan> getChild() {
return child == null ? ImmutableList.of() : ImmutableList.of(child);
}

@Override
public <T, C> T accept(AbstractNodeVisitor<T, C> visitor, C context) {
return visitor.visitFormat(this, context);
}
}
25 changes: 24 additions & 1 deletion core/src/main/java/org/opensearch/sql/ast/tree/Search.java
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import lombok.ToString;
import org.opensearch.sql.ast.AbstractNodeVisitor;
import org.opensearch.sql.ast.expression.SearchExpression;
import org.opensearch.sql.ast.expression.SearchSubquery;

/**
* Logical plan node for Search operation. Represents search expressions that get converted to
Expand All @@ -26,7 +27,7 @@
public class Search extends UnresolvedPlan {

@EqualsAndHashCode.Include private final UnresolvedPlan child;
@EqualsAndHashCode.Include private final String queryString;
@EqualsAndHashCode.Include private final @Nullable String queryString;

// Currently it's only for anonymizer
private final @Nullable SearchExpression originalExpression;
Expand All @@ -35,6 +36,28 @@ public Search(UnresolvedPlan child, String queryString) {
this(child, queryString, null);
}

/** Creates a search that may contain an implicit subquery. */
public static Search fromExpression(UnresolvedPlan child, SearchExpression originalExpression) {
return new Search(
child,
containsSubquery(originalExpression) ? null : originalExpression.toQueryString(),
originalExpression);
}

public boolean hasImplicitSubquery() {
return queryString == null;
}

private static boolean containsSubquery(SearchExpression expression) {
if (expression instanceof SearchSubquery) {
return true;
}
return expression.getChild().stream()
.filter(SearchExpression.class::isInstance)
.map(SearchExpression.class::cast)
.anyMatch(Search::containsSubquery);
}

@Override
public List<UnresolvedPlan> getChild() {
return ImmutableList.of(child);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,13 @@

import java.sql.Connection;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.Stack;
import java.util.function.BiFunction;
import lombok.Getter;
Expand All @@ -21,6 +24,7 @@
import org.apache.calcite.rex.RexCorrelVariable;
import org.apache.calcite.rex.RexLambdaRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexSubQuery;
import org.apache.calcite.tools.FrameworkConfig;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.opensearch.sql.ast.expression.AggregateFunction;
Expand Down Expand Up @@ -83,6 +87,10 @@ public class CalcitePlanContext {
private final Stack<RexCorrelVariable> correlVar = new Stack<>();
private final Stack<List<RexNode>> windowPartitions = new Stack<>();

/** Identity marker for scalar subqueries produced by an implicit format inside search. */
private final Set<RexSubQuery> implicitFormatSubqueries =
Collections.newSetFromMap(new IdentityHashMap<>());

@Getter public Map<String, RexLambdaRef> rexLambdaRefMap;

/**
Expand Down Expand Up @@ -205,6 +213,14 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
}
}

void registerImplicitFormatSubquery(RexSubQuery subquery) {
implicitFormatSubqueries.add(subquery);
}

boolean isImplicitFormatSubquery(RexSubQuery subquery) {
return implicitFormatSubqueries.contains(subquery);
}

/**
* Creates a clone of this context that shares the relBuilder with the parent. This allows lambda
* expressions to reference fields from the current row while having their own lambda variable
Expand Down
Loading
Loading