Pushdown case function in aggregations as range queries

[yuancu]

Description

This PR push down CASE functions used in aggregations as range queries.

For example, the query source=bank | eval age_range = case (age < 30, 'u30', age < 40, 'u40' else 'u100') | stats avg(balance) by age_range will be pushed down as the following OpenSearch DSL:

{
  "aggregations": {
    "age_range": {
      "range": {
        "field": "age",
        "ranges": [
          {
            "key": "u30",
            "to": 30
          },
          {
            "key": "u40",
            "from": 30,
            "to": 40
          },
          {
            "key": "u100",
            "from": 40
          }
        ],
        "keyed": true
      },
      "aggregations": {
        "avg(balance)": {
          "avg": {
            "field": "balance"
          }
        }
      }
    }
  }
}

A CASE function used in aggregation can be pushed down only if it satisfied the following criteria:

• Result expressions must be string literals
• The field referenced in the condition must be the same
• Field references must be numeric
• Ranges must be closed-open intervals: $[a, b)$, $[a, +inf)$, or $(-inf, b)$.
• No further operations can be performed on the result of the case function

Limitations:

• It only handle cases where the result expression is a string literal. E.g. case(balance<10, 'poor' else 'rich') will be pushed down, while case(balance<10, 10 else 100) won't.
• Red flag: range query will ignore null values. E.g. eval b = case(balance<10, 'poor' else 'rich') | stats avg(balance) by b will not properly handle cases when there are balance with null values. For case function, null values be categorized into the else group; while with pushed-down aggregation, rows with null balance will be ignored.
• With case function, the default else group is null. However, since null can not be a key for a range query, we substitute it with "null". This can be fixed later by assigning a secret key to the else group, and substituting it later when parsing the response.

Examples of generated DSL

Case 1: Group by the case field only, with sub-aggregations

source=bank | eval age_range = case (age < 30, 'u30', age < 40, 'u40' else 'u100') | stats avg(balance) by age_range

RangeAgg
    MetricAgg

{
  "aggregations": {
    "age_range": {
      "range": {
        "field": "age",
        "ranges": [
          {
            "key": "u30",
            "to": 30
          },
          {
            "key": "u40",
            "from": 30,
            "to": 40
          },
          {
            "key": "u100",
            "from": 40
          }
        ],
        "keyed": true
      },
      "aggregations": {
        "avg(balance)": {
          "avg": {
            "field": "balance"
          }
        }
      }
    }
  }
}

Case 2: Group by multiple ranges with sub-aggregations

source=bank | eval age_range = case (age < 30, 'u30', age < 35, 'u35', age < 40, 'u40', age >= 40, 'u100'), balance_range = case(balance < 20000, 'medium' else 'high') | stats avg(balance) by age_range, balance_range

RangeAgg
    RangeAgg
        MetricAgg

{
  "aggregations": {
    "age_range": {
      "range": {
        "field": "age",
        "ranges": [
          {
            "key": "u30",
            "to": 30
          },
          {
            "key": "u35",
            "from": 30,
            "to": 35
          },
          {
            "key": "u40",
            "from": 35,
            "to": 40
          },
          {
            "key": "u100",
            "from": 40
          }
        ],
        "keyed": true
      },
      "aggregations": {
        "balance_range": {
          "range": {
            "field": "balance",
            "ranges": [
              {
                "key": "medium",
                "to": 20000
              },
              {
                "key": "high",
                "from": 20000
              }
            ],
            "keyed": true
          },
          "aggregations": {
            "avg(balance)": {
              "avg": {
                "field": "balance"
              }
            }
          }
        }
      }
    }
  }
}

Case 3: Group by case field and keyword field

source=bank | eval age_range = case (age < 30, 'u30', age < 35, 'u35', age < 40, 'u40', age >= 40, 'u100') | stats avg(balance), count() by firstname, lastname, age_range

CompositeAgg
    RangeAgg
        MetricAgg

{
  "aggregations": {
    "composite_buckets": {
      "composite": {
        "size": 1000,
        "sources": [
          {
            "firstname": {
              "terms": {
                "field": "firstname",
                "missing_bucket": true,
                "missing_order": "first",
                "order": "asc"
              }
            }
          },
          {
            "lastname": {
              "terms": {
                "field": "lastname",
                "missing_bucket": true,
                "missing_order": "first",
                "order": "asc"
              }
            }
          }
        ]
      },
      "aggregations": {
        "age_range": {
          "range": {
            "field": "age",
            "ranges": [
              {
                "key": "u30",
                "to": 30
              },
              {
                "key": "u35",
                "from": 30,
                "to": 35
              },
              {
                "key": "u40",
                "from": 35,
                "to": 40
              },
              {
                "key": "u100",
                "from": 40
              }
            ],
            "keyed": true
          },
          "aggregations": {
            "avg(balance)": {
              "avg": {
                "field": "balance"
              }
            }
          }
        }
      }
    }
  }
}

TODOs:

• Fix the discrepancy of null as expression results (in the pushed down version, it is "null" instead of null) left as a limitation.
• Test cases where the result expressions are numbers
• Test cases where there contains null values
• Javadoc for newly added interfaces
• Refactor bucket parsers, etc once Fallback to sub-aggregation if composite aggregation doesn't support #4413 is merged
• Unify how sub-aggregations are created across AutoDataHistogramAggregation, RangeAggregation and CompositeAggregation

Related Issues

Resolves #4201 , partially resolves #4338

Check List

• New functionality includes testing.
• New functionality has been documented.
• New functionality has javadoc added.
• New functionality has a user manual doc added.
• New PPL command checklist all confirmed.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff or -s.
• Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[qianheng-aws]

[non-blocking] Should this part of code be put in the second branch of the below if for better readability? It should have structure like:

if (aggregate.getGroupSet().isEmpty()) {
  // no group by
} else {
    // Push auto date span & case in group-by list into nested aggregations
   ...
   ...
   if (groupList.isEmpty()) {
        // No composite aggregation at top-level
        ...
   } else {
        // Composite aggregation at top level
        ...
    }
}

It works well with current code but performing useless operations on some empty collections for no-group-by.

[yuancu]

Thanks for the suggestion! optimized the logic here

[qianheng-aws]

Put this at the beginning of this method to skip meaningless operations in advance?

[qianheng-aws]

Please include InternalAutoDateHistogram.Bucket here as well to skip empty bucket for auto_date_span.

And also remove the code

sql/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java

Lines 312 to 329 in 1e62fba

	if (aggregationBuilder.getLeft().size() == 1
	&& aggregationBuilder.getLeft().getFirst()
	instanceof AutoDateHistogramAggregationBuilder autoDateHistogram) {
	// If it's auto_date_histogram, filter the empty bucket by using the first aggregate metrics
	RexBuilder rexBuilder = getCluster().getRexBuilder();
	Optional<AggregationBuilder> aggBuilderOpt =
	autoDateHistogram.getSubAggregations().stream().toList().stream().findFirst();
	RexNode condition =
	aggBuilderOpt.isEmpty() || aggBuilderOpt.get() instanceof ValueCountAggregationBuilder
	? rexBuilder.makeCall(
	SqlStdOperatorTable.GREATER_THAN,
	rexBuilder.makeInputRef(newScan, 1),
	rexBuilder.makeLiteral(
	0, rexBuilder.getTypeFactory().createSqlType(SqlTypeName.INTEGER)))
	: rexBuilder.makeCall(
	SqlStdOperatorTable.IS_NOT_NULL, rexBuilder.makeInputRef(newScan, 1));
	return LogicalFilter.create(newScan, condition);
	}

here. Previously we add another filter to filter out the empty bucket.

[yuancu]

Refactored as suggested

[qianheng-aws]

Should there be another else here? Otherwise it will put the bucket.getKey into the results for composite agg as well.

[yuancu]

Yep. I used to put all contents in bucket.getKey to pass unit tests in a wrong way; corrected the behavior now.

[qianheng-aws]

Please add test case for composite - auto_date_span - range - metric. Set bins to a big enough value like 100 to verify whether the empty buckets are filtered out properly.

[yuancu]

Test added

[LantaoJin]

should change to explainQueryYaml

[yuancu]

Fixed. Thanks for reminding!

[opensearch-trigger-bot]

The backport to 2.19-dev failed:

The process '/usr/bin/git' failed with exit code 128

To backport manually, run these commands in your terminal:

# Navigate to the root of your repository
cd $(git rev-parse --show-toplevel)
# Fetch latest updates from GitHub
git fetch
# Create a new working tree
git worktree add ../.worktrees/sql/backport-2.19-dev 2.19-dev
# Navigate to the new working tree
pushd ../.worktrees/sql/backport-2.19-dev
# Create a new branch
git switch --create backport/backport-4400-to-2.19-dev
# Cherry-pick the merged commit of this pull request and resolve the conflicts
git cherry-pick -x --mainline 1 18ab4dc9a02589e4348d67fa4b38a66d27ad7f59
# Push it to GitHub
git push --set-upstream origin backport/backport-4400-to-2.19-dev
# Go back to the original working tree
popd
# Delete the working tree
git worktree remove ../.worktrees/sql/backport-2.19-dev

Then, create a pull request where the base branch is 2.19-dev and the compare/head branch is backport/backport-4400-to-2.19-dev.