Examples with QueryRewriteContext org.opensearch.index.query.QueryRewriteContext used on opensource projects

Example 1 with QueryRewriteContext

use of org.opensearch.index.query.QueryRewriteContext in project OpenSearch by opensearch-project.

the class DateFieldTypeTests method doTestIsFieldWithinQuery.

private void doTestIsFieldWithinQuery(DateFieldType ft, DirectoryReader reader, DateTimeZone zone, DateMathParser alternateFormat) throws IOException {
    QueryRewriteContext context = new QueryRewriteContext(xContentRegistry(), writableRegistry(), null, () -> nowInMillis);
    assertEquals(Relation.INTERSECTS, ft.isFieldWithinQuery(reader, "2015-10-09", "2016-01-02", randomBoolean(), randomBoolean(), null, null, context));
    assertEquals(Relation.INTERSECTS, ft.isFieldWithinQuery(reader, "2016-01-02", "2016-06-20", randomBoolean(), randomBoolean(), null, null, context));
    assertEquals(Relation.INTERSECTS, ft.isFieldWithinQuery(reader, "2016-01-02", "2016-02-12", randomBoolean(), randomBoolean(), null, null, context));
    assertEquals(Relation.DISJOINT, ft.isFieldWithinQuery(reader, "2014-01-02", "2015-02-12", randomBoolean(), randomBoolean(), null, null, context));
    assertEquals(Relation.DISJOINT, ft.isFieldWithinQuery(reader, "2016-05-11", "2016-08-30", randomBoolean(), randomBoolean(), null, null, context));
    assertEquals(Relation.WITHIN, ft.isFieldWithinQuery(reader, "2015-09-25", "2016-05-29", randomBoolean(), randomBoolean(), null, null, context));
    assertEquals(Relation.WITHIN, ft.isFieldWithinQuery(reader, "2015-10-12", "2016-04-03", true, true, null, null, context));
    assertEquals(Relation.INTERSECTS, ft.isFieldWithinQuery(reader, "2015-10-12", "2016-04-03", false, false, null, null, context));
    assertEquals(Relation.INTERSECTS, ft.isFieldWithinQuery(reader, "2015-10-12", "2016-04-03", false, true, null, null, context));
    assertEquals(Relation.INTERSECTS, ft.isFieldWithinQuery(reader, "2015-10-12", "2016-04-03", true, false, null, null, context));
}

Example 2 with QueryRewriteContext

use of org.opensearch.index.query.QueryRewriteContext in project OpenSearch by opensearch-project.

the class AggregatorFactoriesTests method testRewriteAggregationAtTopLevel.

public void testRewriteAggregationAtTopLevel() throws Exception {
    FilterAggregationBuilder filterAggBuilder = new FilterAggregationBuilder("titles", new MatchAllQueryBuilder());
    AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder).addPipelineAggregator(new RewrittenPipelineAggregationBuilder());
    QueryRewriteContext context = new QueryRewriteContext(xContentRegistry, null, null, () -> 0L);
    AggregatorFactories.Builder rewritten = builder.rewrite(context);
    CountDownLatch latch = new CountDownLatch(1);
    context.executeAsyncActions(new ActionListener<Object>() {

        @Override
        public void onResponse(Object response) {
            assertNotSame(builder, rewritten);
            PipelineAggregationBuilder rewrittenPipeline = rewritten.getPipelineAggregatorFactories().iterator().next();
            assertThat(((RewrittenPipelineAggregationBuilder) rewrittenPipeline).setOnRewrite.get(), equalTo("rewritten"));
            latch.countDown();
        }

        @Override
        public void onFailure(Exception e) {
            throw new AssertionError(e);
        }
    });
    latch.await();
}

Example 3 with QueryRewriteContext

use of org.opensearch.index.query.QueryRewriteContext in project OpenSearch by opensearch-project.

the class AggregatorFactoriesTests method testRewriteAggregation.

public void testRewriteAggregation() throws Exception {
    XContentType xContentType = randomFrom(XContentType.values());
    BytesReference bytesReference;
    try (XContentBuilder builder = XContentFactory.contentBuilder(xContentType)) {
        builder.startObject();
        {
            builder.startObject("terms");
            {
                builder.array("title", "foo");
            }
            builder.endObject();
        }
        builder.endObject();
        bytesReference = BytesReference.bytes(builder);
    }
    FilterAggregationBuilder filterAggBuilder = new FilterAggregationBuilder("titles", new WrapperQueryBuilder(bytesReference));
    BucketScriptPipelineAggregationBuilder pipelineAgg = new BucketScriptPipelineAggregationBuilder("const", new Script("1"));
    AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder).addPipelineAggregator(pipelineAgg);
    AggregatorFactories.Builder rewritten = builder.rewrite(new QueryRewriteContext(xContentRegistry, null, null, () -> 0L));
    assertNotSame(builder, rewritten);
    Collection<AggregationBuilder> aggregatorFactories = rewritten.getAggregatorFactories();
    assertEquals(1, aggregatorFactories.size());
    assertThat(aggregatorFactories.iterator().next(), instanceOf(FilterAggregationBuilder.class));
    FilterAggregationBuilder rewrittenFilterAggBuilder = (FilterAggregationBuilder) aggregatorFactories.iterator().next();
    assertNotSame(filterAggBuilder, rewrittenFilterAggBuilder);
    assertNotEquals(filterAggBuilder, rewrittenFilterAggBuilder);
    // Check the filter was rewritten from a wrapper query to a terms query
    QueryBuilder rewrittenFilter = rewrittenFilterAggBuilder.getFilter();
    assertThat(rewrittenFilter, instanceOf(TermsQueryBuilder.class));
    // Check that a further rewrite returns the same aggregation factories builder
    AggregatorFactories.Builder secondRewritten = rewritten.rewrite(new QueryRewriteContext(xContentRegistry, null, null, () -> 0L));
    assertSame(rewritten, secondRewritten);
}

Example 4 with QueryRewriteContext

use of org.opensearch.index.query.QueryRewriteContext in project OpenSearch by opensearch-project.

the class AggregatorFactoriesTests method testRewritePipelineAggregationUnderAggregation.

public void testRewritePipelineAggregationUnderAggregation() throws Exception {
    FilterAggregationBuilder filterAggBuilder = new FilterAggregationBuilder("titles", new MatchAllQueryBuilder()).subAggregation(new RewrittenPipelineAggregationBuilder());
    AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder);
    QueryRewriteContext context = new QueryRewriteContext(xContentRegistry, null, null, () -> 0L);
    AggregatorFactories.Builder rewritten = builder.rewrite(context);
    CountDownLatch latch = new CountDownLatch(1);
    context.executeAsyncActions(new ActionListener<Object>() {

        @Override
        public void onResponse(Object response) {
            assertNotSame(builder, rewritten);
            Collection<AggregationBuilder> aggregatorFactories = rewritten.getAggregatorFactories();
            assertEquals(1, aggregatorFactories.size());
            FilterAggregationBuilder rewrittenFilterAggBuilder = (FilterAggregationBuilder) aggregatorFactories.iterator().next();
            PipelineAggregationBuilder rewrittenPipeline = rewrittenFilterAggBuilder.getPipelineAggregations().iterator().next();
            assertThat(((RewrittenPipelineAggregationBuilder) rewrittenPipeline).setOnRewrite.get(), equalTo("rewritten"));
            latch.countDown();
        }

        @Override
        public void onFailure(Exception e) {
            throw new AssertionError(e);
        }
    });
    latch.await();
}

Example 5 with QueryRewriteContext

use of org.opensearch.index.query.QueryRewriteContext in project OpenSearch by opensearch-project.

the class FiltersTests method testRewrite.

public void testRewrite() throws IOException {
    // test non-keyed filter that doesn't rewrite
    AggregationBuilder original = new FiltersAggregationBuilder("my-agg", new MatchAllQueryBuilder());
    original.setMetadata(Collections.singletonMap(randomAlphaOfLengthBetween(1, 20), randomAlphaOfLengthBetween(1, 20)));
    AggregationBuilder rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L));
    assertSame(original, rewritten);
    // test non-keyed filter that does rewrite
    original = new FiltersAggregationBuilder("my-agg", new BoolQueryBuilder());
    rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L));
    assertNotSame(original, rewritten);
    assertThat(rewritten, instanceOf(FiltersAggregationBuilder.class));
    assertEquals("my-agg", ((FiltersAggregationBuilder) rewritten).getName());
    assertEquals(1, ((FiltersAggregationBuilder) rewritten).filters().size());
    assertEquals("0", ((FiltersAggregationBuilder) rewritten).filters().get(0).key());
    assertThat(((FiltersAggregationBuilder) rewritten).filters().get(0).filter(), instanceOf(MatchAllQueryBuilder.class));
    assertFalse(((FiltersAggregationBuilder) rewritten).isKeyed());
    // test keyed filter that doesn't rewrite
    original = new FiltersAggregationBuilder("my-agg", new KeyedFilter("my-filter", new MatchAllQueryBuilder()));
    rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L));
    assertSame(original, rewritten);
    // test non-keyed filter that does rewrite
    original = new FiltersAggregationBuilder("my-agg", new KeyedFilter("my-filter", new BoolQueryBuilder()));
    rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L));
    assertNotSame(original, rewritten);
    assertThat(rewritten, instanceOf(FiltersAggregationBuilder.class));
    assertEquals("my-agg", ((FiltersAggregationBuilder) rewritten).getName());
    assertEquals(1, ((FiltersAggregationBuilder) rewritten).filters().size());
    assertEquals("my-filter", ((FiltersAggregationBuilder) rewritten).filters().get(0).key());
    assertThat(((FiltersAggregationBuilder) rewritten).filters().get(0).filter(), instanceOf(MatchAllQueryBuilder.class));
    assertTrue(((FiltersAggregationBuilder) rewritten).isKeyed());
    // test sub-agg filter that does rewrite
    original = new TermsAggregationBuilder("terms").userValueTypeHint(ValueType.BOOLEAN).subAggregation(new FiltersAggregationBuilder("my-agg", new KeyedFilter("my-filter", new BoolQueryBuilder())));
    rewritten = original.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L));
    assertNotSame(original, rewritten);
    assertNotEquals(original, rewritten);
    assertThat(rewritten, instanceOf(TermsAggregationBuilder.class));
    assertThat(rewritten.getSubAggregations().size(), equalTo(1));
    AggregationBuilder subAgg = rewritten.getSubAggregations().iterator().next();
    assertThat(subAgg, instanceOf(FiltersAggregationBuilder.class));
    assertNotSame(original.getSubAggregations().iterator().next(), subAgg);
    assertEquals("my-agg", subAgg.getName());
    assertSame(rewritten, rewritten.rewrite(new QueryRewriteContext(xContentRegistry(), null, null, () -> 0L)));
}