|
44 | 44 | import org.opensearch.search.aggregations.AggregationBuilder; |
45 | 45 | import org.opensearch.search.aggregations.AggregationBuilders; |
46 | 46 | import org.opensearch.search.aggregations.AggregatorFactories.Builder; |
| 47 | +import org.opensearch.search.aggregations.BucketOrder; |
47 | 48 | import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder; |
48 | 49 | import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceBuilder; |
49 | 50 | import org.opensearch.search.aggregations.bucket.missing.MissingOrder; |
| 51 | +import org.opensearch.search.aggregations.bucket.terms.MultiTermsAggregationBuilder; |
| 52 | +import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder; |
50 | 53 | import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder; |
51 | 54 | import org.opensearch.search.sort.ScoreSortBuilder; |
52 | 55 | import org.opensearch.search.sort.SortBuilder; |
@@ -397,53 +400,88 @@ public void apply(OpenSearchRequestBuilder requestBuilder) { |
397 | 400 | } |
398 | 401 |
|
399 | 402 | public void pushDownSortIntoAggBucket(List<RelFieldCollation> collations) { |
400 | | - // It will always use a single CompositeAggregationBuilder for the aggregation with GroupBy |
401 | | - // See {@link AggregateAnalyzer} |
402 | | - CompositeAggregationBuilder compositeAggregationBuilder = |
403 | | - (CompositeAggregationBuilder) aggregationBuilder.getLeft().getFirst(); |
404 | | - List<CompositeValuesSourceBuilder<?>> buckets = |
405 | | - ((CompositeAggregationBuilder) aggregationBuilder.getLeft().getFirst()).sources(); |
406 | | - List<CompositeValuesSourceBuilder<?>> newBuckets = new ArrayList<>(buckets.size()); |
| 403 | + AggregationBuilder builder = aggregationBuilder.getLeft().getFirst(); |
407 | 404 | List<Integer> selected = new ArrayList<>(collations.size()); |
408 | | - // Have to put the collation required buckets first, then the rest of buckets. |
409 | | - collations.forEach( |
410 | | - collation -> { |
411 | | - CompositeValuesSourceBuilder<?> bucket = buckets.get(collation.getFieldIndex()); |
412 | | - Direction direction = collation.getDirection(); |
413 | | - NullDirection nullDirection = collation.nullDirection; |
414 | | - SortOrder order = |
415 | | - Direction.DESCENDING.equals(direction) ? SortOrder.DESC : SortOrder.ASC; |
416 | | - MissingOrder missingOrder = |
417 | | - switch (nullDirection) { |
418 | | - case FIRST -> MissingOrder.FIRST; |
419 | | - case LAST -> MissingOrder.LAST; |
420 | | - default -> MissingOrder.DEFAULT; |
421 | | - }; |
422 | | - newBuckets.add(bucket.order(order).missingOrder(missingOrder)); |
423 | | - selected.add(collation.getFieldIndex()); |
424 | | - }); |
425 | | - IntStream.range(0, buckets.size()) |
426 | | - .filter(i -> !selected.contains(i)) |
427 | | - .forEach(i -> newBuckets.add(buckets.get(i))); |
428 | | - Builder newAggBuilder = new Builder(); |
429 | | - compositeAggregationBuilder.getSubAggregations().forEach(newAggBuilder::addAggregator); |
430 | | - aggregationBuilder = |
431 | | - Pair.of( |
432 | | - Collections.singletonList( |
433 | | - AggregationBuilders.composite("composite_buckets", newBuckets) |
434 | | - .subAggregations(newAggBuilder) |
435 | | - .size(AGGREGATION_BUCKET_SIZE)), |
436 | | - aggregationBuilder.getRight()); |
| 405 | + if (builder instanceof CompositeAggregationBuilder compositeAggBuilder) { |
| 406 | + // It will always use a single CompositeAggregationBuilder for the aggregation with GroupBy |
| 407 | + // See {@link AggregateAnalyzer} |
| 408 | + List<CompositeValuesSourceBuilder<?>> buckets = compositeAggBuilder.sources(); |
| 409 | + List<CompositeValuesSourceBuilder<?>> newBuckets = new ArrayList<>(buckets.size()); |
| 410 | + // Have to put the collation required buckets first, then the rest of buckets. |
| 411 | + collations.forEach( |
| 412 | + collation -> { |
| 413 | + CompositeValuesSourceBuilder<?> bucket = buckets.get(collation.getFieldIndex()); |
| 414 | + Direction direction = collation.getDirection(); |
| 415 | + NullDirection nullDirection = collation.nullDirection; |
| 416 | + SortOrder order = |
| 417 | + Direction.DESCENDING.equals(direction) ? SortOrder.DESC : SortOrder.ASC; |
| 418 | + MissingOrder missingOrder = |
| 419 | + switch (nullDirection) { |
| 420 | + case FIRST -> MissingOrder.FIRST; |
| 421 | + case LAST -> MissingOrder.LAST; |
| 422 | + default -> MissingOrder.DEFAULT; |
| 423 | + }; |
| 424 | + newBuckets.add(bucket.order(order).missingOrder(missingOrder)); |
| 425 | + selected.add(collation.getFieldIndex()); |
| 426 | + }); |
| 427 | + IntStream.range(0, buckets.size()) |
| 428 | + .filter(i -> !selected.contains(i)) |
| 429 | + .forEach(i -> newBuckets.add(buckets.get(i))); |
| 430 | + Builder newAggBuilder = new Builder(); |
| 431 | + compositeAggBuilder.getSubAggregations().forEach(newAggBuilder::addAggregator); |
| 432 | + aggregationBuilder = |
| 433 | + Pair.of( |
| 434 | + Collections.singletonList( |
| 435 | + AggregationBuilders.composite("composite_buckets", newBuckets) |
| 436 | + .subAggregations(newAggBuilder) |
| 437 | + .size(AGGREGATION_BUCKET_SIZE)), |
| 438 | + aggregationBuilder.getRight()); |
| 439 | + } |
| 440 | + if (builder instanceof TermsAggregationBuilder termsAggBuilder) { |
| 441 | + termsAggBuilder.order( |
| 442 | + BucketOrder.key(!collations.getFirst().getDirection().isDescending())); |
| 443 | + } |
| 444 | + // TODO for MultiTermsAggregationBuilder |
437 | 445 | } |
438 | 446 |
|
| 447 | + /** |
| 448 | + * Check if the limit can be pushed down into aggregation bucket when the limit size is less |
| 449 | + * than bucket number. |
| 450 | + */ |
439 | 451 | public boolean pushDownLimitIntoBucketSize(Integer size) { |
440 | | - CompositeAggregationBuilder compositeAggregationBuilder = |
441 | | - (CompositeAggregationBuilder) aggregationBuilder.getLeft().getFirst(); |
442 | | - if (size < compositeAggregationBuilder.size()) { |
443 | | - compositeAggregationBuilder.size(size); |
| 452 | + AggregationBuilder builder = aggregationBuilder.getLeft().getFirst(); |
| 453 | + if (builder instanceof CompositeAggregationBuilder compositeAggBuilder) { |
| 454 | + if (size < compositeAggBuilder.size()) { |
| 455 | + compositeAggBuilder.size(size); |
| 456 | + return true; |
| 457 | + } else { |
| 458 | + return false; |
| 459 | + } |
| 460 | + } |
| 461 | + if (builder instanceof TermsAggregationBuilder termsAggBuilder) { |
| 462 | + if (size < termsAggBuilder.size()) { |
| 463 | + termsAggBuilder.size(size); |
| 464 | + return true; |
| 465 | + } else { |
| 466 | + return false; |
| 467 | + } |
| 468 | + } |
| 469 | + if (builder instanceof MultiTermsAggregationBuilder multiTermsAggBuilder) { |
| 470 | + if (size < multiTermsAggBuilder.size()) { |
| 471 | + multiTermsAggBuilder.size(size); |
| 472 | + return true; |
| 473 | + } else { |
| 474 | + return false; |
| 475 | + } |
| 476 | + } |
| 477 | + // now we only have Composite, Terms and MultiTerms bucket aggregations, |
| 478 | + // add code here when we could support more in the future. |
| 479 | + if (builder instanceof ValuesSourceAggregationBuilder.LeafOnly<?, ?>) { |
| 480 | + // Note: all metric aggregations will be treated as pushed since it generates only one row. |
444 | 481 | return true; |
445 | 482 | } |
446 | | - return false; |
| 483 | + throw new OpenSearchRequestBuilder.PushDownUnSupportedException( |
| 484 | + "Unknown aggregation builder " + builder.getClass().getSimpleName()); |
447 | 485 | } |
448 | 486 | } |
449 | 487 | } |
0 commit comments