Skip to content

Commit

Permalink
HSEARCH-5133 Support temporal metric aggregations
Browse files Browse the repository at this point in the history
  • Loading branch information
fax4ever committed Jun 28, 2024
1 parent b94609b commit 18982c6
Show file tree
Hide file tree
Showing 2 changed files with 128 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
import java.time.temporal.TemporalAccessor;

import org.hibernate.search.backend.elasticsearch.lowlevel.index.mapping.impl.DataTypes;
import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchMetricFieldAggregation;
import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchMetricLongAggregation;
import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchRangeAggregation;
import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchTermsAggregation;
import org.hibernate.search.backend.elasticsearch.search.predicate.impl.ElasticsearchExistsPredicate;
Expand Down Expand Up @@ -77,6 +79,18 @@ protected final void complete() {
builder.aggregable( true );
builder.queryElementFactory( AggregationTypeKeys.TERMS, new ElasticsearchTermsAggregation.Factory<>( codec ) );
builder.queryElementFactory( AggregationTypeKeys.RANGE, new ElasticsearchRangeAggregation.Factory<>( codec ) );
builder.queryElementFactory( AggregationTypeKeys.SUM,
new ElasticsearchMetricFieldAggregation.Factory<>( codec, "sum" ) );
builder.queryElementFactory( AggregationTypeKeys.MIN,
new ElasticsearchMetricFieldAggregation.Factory<>( codec, "min" ) );
builder.queryElementFactory( AggregationTypeKeys.MAX,
new ElasticsearchMetricFieldAggregation.Factory<>( codec, "max" ) );
builder.queryElementFactory( AggregationTypeKeys.COUNT,
new ElasticsearchMetricLongAggregation.Factory<>( codec, "value_count" ) );
builder.queryElementFactory( AggregationTypeKeys.COUNT_DISTINCT,
new ElasticsearchMetricLongAggregation.Factory<>( codec, "cardinality" ) );
builder.queryElementFactory( AggregationTypeKeys.AVG,
new ElasticsearchMetricFieldAggregation.Factory<>( codec, "avg" ) );
}
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
/*
* SPDX-License-Identifier: Apache-2.0
* Copyright Red Hat Inc. and Hibernate Authors
*/
package org.hibernate.search.integrationtest.backend.elasticsearch.tmp;

import static org.assertj.core.api.Assertions.assertThat;

import java.time.LocalDate;
import java.time.Month;

import org.hibernate.search.engine.backend.common.DocumentReference;
import org.hibernate.search.engine.backend.document.DocumentElement;
import org.hibernate.search.engine.backend.document.IndexFieldReference;
import org.hibernate.search.engine.backend.document.IndexObjectFieldReference;
import org.hibernate.search.engine.backend.document.model.dsl.IndexSchemaElement;
import org.hibernate.search.engine.backend.document.model.dsl.IndexSchemaObjectField;
import org.hibernate.search.engine.backend.types.Aggregable;
import org.hibernate.search.engine.backend.types.ObjectStructure;
import org.hibernate.search.engine.search.aggregation.AggregationKey;
import org.hibernate.search.engine.search.query.SearchQuery;
import org.hibernate.search.engine.search.query.SearchResult;
import org.hibernate.search.integrationtest.backend.tck.testsupport.util.extension.SearchSetupHelper;
import org.hibernate.search.util.impl.integrationtest.mapper.stub.BulkIndexer;
import org.hibernate.search.util.impl.integrationtest.mapper.stub.SimpleMappedIndex;
import org.hibernate.search.util.impl.integrationtest.mapper.stub.StubMappingScope;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;

public class ElasticsearchMetricTemporalFieldsAggregationsIT {

@RegisterExtension
public static final SearchSetupHelper setupHelper = SearchSetupHelper.create();

private final SimpleMappedIndex<IndexBinding> mainIndex = SimpleMappedIndex.of( IndexBinding::new ).name( "main" );
private final AggregationKey<LocalDate> sumDates = AggregationKey.of( "sumDates" );

@BeforeEach
void setup() {
setupHelper.start().withIndexes( mainIndex ).setup().integration();
initData();
}

@Test
public void test_filteringResults() {
StubMappingScope scope = mainIndex.createScope();
SearchQuery<DocumentReference> query = scope.query()
.where( f -> f.match().field( "style" ).matching( "bla" ) )
.aggregation( sumDates, f -> f.sum().field( "date", LocalDate.class ) )
.toQuery();

SearchResult<DocumentReference> result = query.fetch( 0 );
assertThat( result.aggregation( sumDates ) ).isEqualTo( LocalDate.of( 2204, Month.SEPTEMBER, 16 ) );
}

@Test
public void test_allResults() {
StubMappingScope scope = mainIndex.createScope();
SearchQuery<DocumentReference> query = scope.query()
.where( f -> f.matchAll() )
.aggregation( sumDates, f -> f.sum().field( "date", LocalDate.class ) )
.toQuery();

SearchResult<DocumentReference> result = query.fetch( 0 );
assertThat( result.aggregation( sumDates ) ).isEqualTo( LocalDate.of( 2439, Month.JUNE, 6 ) );
}

private void initData() {
LocalDate baseDate = LocalDate.of( 2016, Month.DECEMBER, 6 );
int[] integers = new int[] { 9, 18, 3, 18, 7, -10, 3, 0, 7, 0 };
String[] styles = new String[] { "bla", "aaa" };

BulkIndexer bulkIndexer = mainIndex.bulkIndexer();
for ( int i = 0; i < integers.length; i++ ) {
int value = integers[i];
String style = styles[i % 2];
String id = i + ":" + value + ":" + style;
LocalDate date = baseDate.plusDays( i );

bulkIndexer.add( id, document -> {
document.addValue( mainIndex.binding().date, date );
document.addValue( mainIndex.binding().converted, date );
document.addValue( mainIndex.binding().style, style );

DocumentElement object = document.addObject( mainIndex.binding().object );
object.addValue( mainIndex.binding().nestedDate, date );
} );
}
bulkIndexer.add( "empty", document -> {} ).join();
}

@SuppressWarnings("unused")
private static class IndexBinding {
final IndexFieldReference<LocalDate> date;
final IndexFieldReference<LocalDate> converted;
final IndexFieldReference<String> style;
final IndexObjectFieldReference object;
final IndexFieldReference<LocalDate> nestedDate;

IndexBinding(IndexSchemaElement root) {
date = root.field( "date", f -> f.asLocalDate().aggregable( Aggregable.YES ) ).toReference();
converted = root.field( "converted", f -> f.asLocalDate().aggregable( Aggregable.YES )
.projectionConverter( String.class, (value, context) -> value.toString() ) ).toReference();
style = root.field( "style", f -> f.asString() ).toReference();

IndexSchemaObjectField nested = root.objectField( "object", ObjectStructure.NESTED );
object = nested.toReference();
nestedDate = nested.field( "nestedDate", f -> f.asLocalDate().aggregable( Aggregable.YES ) )
.toReference();
}
}
}

0 comments on commit 18982c6

Please sign in to comment.