Support max=n option

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

common/src/main/java/org/opensearch/sql/common/setting/Settings.java

@@ -36,6 +36,7 @@ public enum Key {
     CALCITE_PUSHDOWN_ENABLED("plugins.calcite.pushdown.enabled"),
     CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR(
         "plugins.calcite.pushdown.rowcount.estimation.factor"),
+    CALCITE_SUPPORT_ALL_JOIN_TYPES("plugins.calcite.all_join_types.allowed"),
 
     /** Query Settings. */
     FIELD_TYPE_TOLERANCE("plugins.query.field_type_tolerance"),

core/src/main/java/org/opensearch/sql/ast/expression/Literal.java

@@ -49,4 +49,5 @@ public String toString() {
 
   public static Literal TRUE = new Literal(true, DataType.BOOLEAN);
   public static Literal FALSE = new Literal(false, DataType.BOOLEAN);
+  public static Literal ZERO = new Literal(Integer.valueOf("0"), DataType.INTEGER);
 }

core/src/main/java/org/opensearch/sql/ast/tree/Join.java

@@ -97,6 +97,11 @@ public enum JoinType {
     FULL
   }
 
+  /** RIGHT, CROSS, FULL are performance sensitive join types */
+  public static List<JoinType> highCostJoinTypes() {
+    return List.of(JoinType.RIGHT, JoinType.CROSS, JoinType.FULL);
+  }
+
   @Getter
   @RequiredArgsConstructor
   public static class JoinHint {

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -677,7 +677,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
     List<UnresolvedPlan> children = node.getChildren();
     children.forEach(c -> analyze(c, context));
     if (node.getJoinCondition().isEmpty()) {
-      // For spl compatible grammar
+      // join-with-field-list grammar
       List<String> leftColumns = context.relBuilder.peek(1).getRowType().getFieldNames();
       List<String> rightColumns = context.relBuilder.peek().getRowType().getFieldNames();
       List<String> duplicatedFieldNames =
@@ -715,14 +715,31 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
                 .map(field -> JoinAndLookupUtils.analyzeFieldsForLookUp(field, false, context))
                 .toList();
       }
+      Literal max = node.getArgumentMap().get("max");
+      if (max != null && !max.equals(Literal.ZERO)) {
+        // max != 0 means the right-side should be dedup
+        Integer allowedDuplication = (Integer) max.getValue();
+        if (allowedDuplication < 0) {
+          throw new SemanticCheckException("max option must be a positive integer");
+        }
+        List<RexNode> dedupeFields =
+            node.getJoinFields().isPresent()
+                ? node.getJoinFields().get().stream()
+                    .map(a -> (RexNode) context.relBuilder.field(a.getField().toString()))
+                    .toList()
+                : duplicatedFieldNames.stream()
+                    .map(a -> (RexNode) context.relBuilder.field(a))
+                    .toList();
+        buildDedup(context, dedupeFields, allowedDuplication);
+      }
       context.relBuilder.join(
           JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
       if (!toBeRemovedFields.isEmpty()) {
         context.relBuilder.projectExcept(toBeRemovedFields);
       }
       return context.relBuilder.peek();
     }
-    // The join old grammar doesn't allow empty join condition
+    // The join-with-criteria grammar doesn't allow empty join condition
     RexNode joinCondition =
         node.getJoinCondition()
             .map(c -> rexVisitor.analyzeJoinCondition(c, context))
@@ -925,42 +942,46 @@ public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
       // DropColumns('_row_number_)
       context.relBuilder.projectExcept(_row_number_);
     } else {
-      /*
-       * | dedup 2 a, b keepempty=false
-       * DropColumns('_row_number_)
-       * +- Filter ('_row_number_ <= n)
-       *    +- Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST, specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_], ['a, 'b], ['a ASC NULLS FIRST, 'b ASC NULLS FIRST]
-       *       +- Filter (isnotnull('a) AND isnotnull('b))
-       *          +- ...
-       */
-      // Filter (isnotnull('a) AND isnotnull('b))
-      context.relBuilder.filter(
-          context.relBuilder.and(
-              dedupeFields.stream().map(context.relBuilder::isNotNull).toList()));
-      // Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
-      // specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_], ['a, 'b], ['a ASC
-      // NULLS FIRST, 'b ASC NULLS FIRST]
-      RexNode rowNumber =
-          context
-              .relBuilder
-              .aggregateCall(SqlStdOperatorTable.ROW_NUMBER)
-              .over()
-              .partitionBy(dedupeFields)
-              .orderBy(dedupeFields)
-              .rowsTo(RexWindowBounds.CURRENT_ROW)
-              .as("_row_number_");
-      context.relBuilder.projectPlus(rowNumber);
-      RexNode _row_number_ = context.relBuilder.field("_row_number_");
-      // Filter ('_row_number_ <= n)
-      context.relBuilder.filter(
-          context.relBuilder.lessThanOrEqual(
-              _row_number_, context.relBuilder.literal(allowedDuplication)));
-      // DropColumns('_row_number_)
-      context.relBuilder.projectExcept(_row_number_);
+      buildDedup(context, dedupeFields, allowedDuplication);
     }
     return context.relBuilder.peek();
   }
 
+  private static void buildDedup(
+      CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
+    /*
+     * | dedup 2 a, b keepempty=false
+     * DropColumns('_row_number_)
+     * +- Filter ('_row_number_ <= n)
+     *    +- Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST, specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_], ['a, 'b], ['a ASC NULLS FIRST, 'b ASC NULLS FIRST]
+     *       +- Filter (isnotnull('a) AND isnotnull('b))
+     *          +- ...
+     */
+    // Filter (isnotnull('a) AND isnotnull('b))
+    context.relBuilder.filter(
+        context.relBuilder.and(dedupeFields.stream().map(context.relBuilder::isNotNull).toList()));
+    // Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
+    // specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_], ['a, 'b], ['a ASC
+    // NULLS FIRST, 'b ASC NULLS FIRST]
+    RexNode rowNumber =
+        context
+            .relBuilder
+            .aggregateCall(SqlStdOperatorTable.ROW_NUMBER)
+            .over()
+            .partitionBy(dedupeFields)
+            .orderBy(dedupeFields)
+            .rowsTo(RexWindowBounds.CURRENT_ROW)
+            .as("_row_number_");
+    context.relBuilder.projectPlus(rowNumber);
+    RexNode _row_number_ = context.relBuilder.field("_row_number_");
+    // Filter ('_row_number_ <= n)
+    context.relBuilder.filter(
+        context.relBuilder.lessThanOrEqual(
+            _row_number_, context.relBuilder.literal(allowedDuplication)));
+    // DropColumns('_row_number_)
+    context.relBuilder.projectExcept(_row_number_);
+  }
+
   @Override
   public RelNode visitWindow(Window node, CalcitePlanContext context) {
     visitChildren(node, context);

docs/user/ppl/cmd/join.rst

@@ -19,22 +19,24 @@ Version
 
 Syntax-1
 ========
-| [joinType] join [leftAlias] [rightAlias] on <joinCriteria> <right-dataset>
+| [joinType] join [max=n] [leftAlias] [rightAlias] (on | where) <joinCriteria> <right-dataset>
 
-* joinType: optional. The type of join to perform. The default is ``INNER`` if not specified. Other option is ``LEFT [OUTER]``, ``RIGHT [OUTER]``, ``FULL [OUTER]``, ``CROSS``, ``[LEFT] SEMI``, ``[LEFT] ANTI``.
+* joinType: optional. The type of join to perform. The default is ``inner`` if not specified. Other option is ``left``, ``outer``(alias of ``left``), ``semi``, ``anti`` and performance sensitive types ``right``, ``full`` and ``cross``.
+* max=n: optional. Controls how many subsearch results could be joined against to each row in main search. The default value is 0, means unlimited.
 * leftAlias: optional. The subsearch alias to use with the left join side, to avoid ambiguous naming. Fixed pattern: ``left = <leftAlias>``
 * rightAlias: optional. The subsearch alias to use with the right join side, to avoid ambiguous naming. Fixed pattern: ``right = <rightAlias>``
-* joinCriteria: mandatory. It could be any comparison expression.
+* joinCriteria: mandatory. It could be any comparison expression. Must follow with ``on`` (since 3.0.0) or ``where`` (since 3.3.0) keyword.
 * right-dataset: mandatory. Right dataset could be either an ``index`` or a ``subsearch`` with/without alias.
 
 Syntax-2
 ========
 | (Since 3.3.0)
-| join [type=<joinType>] [overwrite=<bool>] <join-field-list> <right-dataset>
+| join [type=<joinType>] [overwrite=<bool>] [max=n] <join-field-list> <right-dataset>
 
-* type=<joinType>: optional. The type of join to perform. The default is ``INNER`` if not specified. Other option is ``LEFT``, ``RIGHT``, ``FULL``, ``CROSS``, ``SEMI``, ``ANTI``.
+* type=<joinType>: optional. The type of join to perform. The default is ``INNER`` if not specified. Other option is ``left``, ``outer``(alias of ``left``), ``semi``, ``anti`` and performance sensitive types ``right``, ``full`` and ``cross``.
 * overwrite=<bool>: optional. Specifies whether duplicate-named fields from <right-dataset> (subsearch results) should replace corresponding fields in the main search results. The default value is ``true``.
-* join-field-list: optional. The fields to use to build join criteria. The ``join-field-list`` must be present in both sides. If no <join-field-list> is present, all fields that are common to both sides are used.
+* max=n: optional. Controls how many subsearch results could be joined against to each row in main search. The default value is 0, means unlimited.
+* join-field-list: optional. The fields to use to build join criteria. The ``join-field-list`` must be present in both sides. If no <join-field-list> is present, all fields that are common to both sides are used. The comma is optional.
 * right-dataset: mandatory. Right dataset could be either an ``index`` or a ``subsearch`` with/without alias.
 
 Configuration
@@ -69,6 +71,7 @@ Usage
 Join (syntax-1)::
 
     source = table1 | inner join left = l right = r on l.a = r.a table2 | fields l.a, r.a, b, c
+    source = table1 | inner join left = l right = r where l.a = r.a table2 | fields l.a, r.a, b, c
     source = table1 | left join left = l right = r on l.a = r.a table2 | fields l.a, r.a, b, c
     source = table1 | right join left = l right = r on l.a = r.a table2 | fields l.a, r.a, b, c
     source = table1 | full left = l right = r on l.a = r.a table2 | fields l.a, r.a, b, c
@@ -86,6 +89,7 @@ Join (syntax-2)::
 
     source = table1 | join a table2 | fields a, b, c
     source = table1 | join a, b table2 | fields a, b, c
+    source = table1 | join type=outer a b table2 | fields a, b, c
     source = table1 | join type=left overwrite=false a, b [source=table2 | rename d as b] | fields a, b, c
 
 Example 1: Two indices join
@@ -179,3 +183,5 @@ Assume table1 and table2 only contain field ``id``, following PPL queries and th
      - table1.id, tt.id, tt.b, a
 
 But for the syntax-2 (since 3.2.0), duplicate-named fields in output results are deduplicated, with field retention determined by the value of 'overwrite' option.
+
+Join types ``inner``, ``left``, ``outer`` (alias of ``left``), ``semi`` and ``anti`` are supported by default. ``right``, ``full``, ``cross`` are performance sensitive join types which are disabled by default. Set config ``plugins.calcite.all_join_types.allowed = true`` to enable.

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLJoinIT.java

@@ -29,6 +29,7 @@ public class CalcitePPLJoinIT extends PPLIntegTestCase {
   public void init() throws Exception {
     super.init();
     enableCalcite();
+    supportAllJoinTypes();
 
     loadIndex(Index.STATE_COUNTRY);
     loadIndex(Index.OCCUPATION);

integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java

@@ -212,6 +212,12 @@ public static void withFallbackEnabled(Runnable f, String msg) throws IOExceptio
     }
   }
 
+  public static void supportAllJoinTypes() throws IOException {
+    updateClusterSettings(
+        new SQLIntegTestCase.ClusterSetting(
+            "persistent", Key.CALCITE_SUPPORT_ALL_JOIN_TYPES.getKeyValue(), "true"));
+  }
+
   public static void withSettings(Key setting, String value, Runnable f) throws IOException {
     String originalValue = getClusterSetting(setting.getKeyValue(), "transient");
     if (originalValue.equals(value)) f.run();

opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java

@@ -127,6 +127,13 @@ public class OpenSearchSettings extends Settings {
           Setting.Property.NodeScope,
           Setting.Property.Dynamic);
 
+  public static final Setting<?> CALCITE_SUPPORT_ALL_JOIN_TYPES_SETTING =
+      Setting.boolSetting(
+          Key.CALCITE_SUPPORT_ALL_JOIN_TYPES.getKeyValue(),
+          false,
+          Setting.Property.NodeScope,
+          Setting.Property.Dynamic);
+
   public static final Setting<?> QUERY_MEMORY_LIMIT_SETTING =
       Setting.memorySizeSetting(
           Key.QUERY_MEMORY_LIMIT.getKeyValue(),
@@ -351,6 +358,12 @@ public OpenSearchSettings(ClusterSettings clusterSettings) {
         Key.CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR,
         CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR_SETTING,
         new Updater(Key.CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR));
+    register(
+        settingBuilder,
+        clusterSettings,
+        Key.CALCITE_SUPPORT_ALL_JOIN_TYPES,
+        CALCITE_SUPPORT_ALL_JOIN_TYPES_SETTING,
+        new Updater(Key.CALCITE_SUPPORT_ALL_JOIN_TYPES));
     register(
         settingBuilder,
         clusterSettings,
@@ -527,6 +540,7 @@ public static List<Setting<?>> pluginSettings() {
         .add(CALCITE_FALLBACK_ALLOWED_SETTING)
         .add(CALCITE_PUSHDOWN_ENABLED_SETTING)
         .add(CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR_SETTING)
+        .add(CALCITE_SUPPORT_ALL_JOIN_TYPES_SETTING)
         .add(DEFAULT_PATTERN_METHOD_SETTING)
         .add(DEFAULT_PATTERN_MODE_SETTING)
         .add(DEFAULT_PATTERN_MAX_SAMPLE_COUNT_SETTING)

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -316,26 +316,37 @@ tableSourceClause
 
 // join
 joinCommand
-   : joinType JOIN sideAlias joinHintList? joinCriteria right = tableOrSubqueryClause
-   | JOIN (joinOption)* (fieldList)? right = tableOrSubqueryClause
+   : JOIN (joinOption)* (fieldList)? right = tableOrSubqueryClause
+   | sqlLikeJoinType? JOIN (joinOption)* sideAlias joinHintList? joinCriteria right = tableOrSubqueryClause
    ;
 
-joinType
-   : INNER?
+sqlLikeJoinType
+   : INNER
    | CROSS
-   | LEFT OUTER?
+   | (LEFT OUTER? | OUTER)
    | RIGHT OUTER?
    | FULL OUTER?
    | LEFT? SEMI
    | LEFT? ANTI
    ;
 
+joinType
+   : INNER
+   | CROSS
+   | OUTER
+   | LEFT
+   | RIGHT
+   | FULL
+   | SEMI
+   | ANTI
+   ;
+
 sideAlias
    : (LEFT EQUAL leftAlias = qualifiedName)? COMMA? (RIGHT EQUAL rightAlias = qualifiedName)?
    ;
 
 joinCriteria
-   : ON logicalExpression
+   : (ON | WHERE) logicalExpression
    ;
 
 joinHintList
@@ -350,6 +361,7 @@ hintPair
 joinOption
    : OVERWRITE EQUAL booleanLiteral                     # overwriteOption
    | TYPE EQUAL joinType                                # typeOption
+   | MAX EQUAL integerLiteral                           # maxOption
    ;
 
 renameClasue
@@ -530,7 +542,7 @@ tableFunction
 
 // fields
 fieldList
-   : fieldExpression (COMMA fieldExpression)*
+   : fieldExpression ((COMMA)? fieldExpression)*
    ;
 
 wcFieldList
@@ -1118,8 +1130,8 @@ keywordsCanBeId
    | multiFieldRelevanceFunctionName
    | commandName
    | collectionFunctionName
-   | comparisonOperator
    | explainMode
+   | REGEXP
    // commands assist keywords
    | CASE
    | ELSE