ISPN-15390 Support score projections

infinispan · Jan 15, 2024 · bd8f6bc · bd8f6bc
1 parent 85bfe72
commit bd8f6bc
Show file tree

Hide file tree

Showing 13 changed files with 86 additions and 29 deletions.
diff --git a/...lter/src/main/java/org/infinispan/objectfilter/impl/syntax/parser/IckleParsingResult.java b/...lter/src/main/java/org/infinispan/objectfilter/impl/syntax/parser/IckleParsingResult.java
@@ -6,6 +6,7 @@
 import org.infinispan.objectfilter.SortField;
 import org.infinispan.objectfilter.impl.ql.PropertyPath;
 import org.infinispan.objectfilter.impl.syntax.BooleanExpr;
+import org.infinispan.objectfilter.impl.syntax.parser.projection.ScorePropertyPath;
 
 /**
  * @param <TypeMetadata> is either {@link java.lang.Class} or {@link org.infinispan.protostream.descriptors.Descriptor}
@@ -133,6 +134,14 @@ public TypeMetadata getTargetEntityMetadata() {
       return targetEntityMetadata;
    }
 
+   public boolean hasScoreProjection() {
+      if (projectedPaths == null) {
+         return false;
+      }
+
+      return Arrays.stream(projectedPaths).sequential().anyMatch(propertyPath -> ScorePropertyPath.SCORE_PROPERTY_NAME.equals(propertyPath.asStringPath()));
+   }
+
    /**
     * Returns the projections of the parsed query, represented as dot paths in case of references to fields of embedded
     * entities, e.g. {@code ["foo", "bar.qaz"]}.

diff --git a/.../src/main/java/org/infinispan/objectfilter/impl/syntax/parser/ProtobufPropertyHelper.java b/.../src/main/java/org/infinispan/objectfilter/impl/syntax/parser/ProtobufPropertyHelper.java
@@ -9,6 +9,7 @@
 import org.infinispan.objectfilter.impl.logging.Log;
 import org.infinispan.objectfilter.impl.syntax.IndexedFieldProvider;
 import org.infinispan.objectfilter.impl.syntax.parser.projection.CacheValuePropertyPath;
+import org.infinispan.objectfilter.impl.syntax.parser.projection.ScorePropertyPath;
 import org.infinispan.objectfilter.impl.syntax.parser.projection.VersionPropertyPath;
 import org.infinispan.objectfilter.impl.util.StringHelper;
 import org.infinispan.protostream.SerializationContext;
@@ -37,6 +38,9 @@ public final class ProtobufPropertyHelper extends ObjectPropertyHelper<Descripto
    public static final String VALUE = CacheValuePropertyPath.VALUE_PROPERTY_NAME;
    public static final int VALUE_FIELD_ATTRIBUTE_ID = 150_001;
 
+   public static final String SCORE = ScorePropertyPath.SCORE_PROPERTY_NAME;
+   public static final int SCORE_FIELD_ATTRIBUTE_ID = 150_002;
+
    private final EntityNameResolver<Descriptor> entityNameResolver;
 
    private final IndexedFieldProvider<Descriptor> indexedFieldProvider;
@@ -69,6 +73,9 @@ public List<?> mapPropertyNamePathToFieldIdPath(Descriptor messageDescriptor, St
          if (propertyPath[0].equals(VALUE)) {
             return Arrays.asList(VALUE_FIELD_ATTRIBUTE_ID);
          }
+         if (propertyPath[0].equals(SCORE)) {
+            return Arrays.asList(SCORE_FIELD_ATTRIBUTE_ID);
+         }
       }
 
       List<Integer> translatedPath = new ArrayList<>(propertyPath.length);
@@ -94,6 +101,9 @@ public Class<?> getPrimitivePropertyType(Descriptor entityType, String[] propert
          if (propertyPath[0].equals(VALUE)) {
             return Object.class;
          }
+         if (propertyPath[0].equals(SCORE)) {
+            return Float.class;
+         }
       }
 
       FieldDescriptor field = getField(entityType, propertyPath);
@@ -158,7 +168,8 @@ private FieldDescriptor getField(Descriptor entityType, String[] propertyPath) {
 
    @Override
    public boolean hasProperty(Descriptor entityType, String[] propertyPath) {
-      if (propertyPath.length == 1 && (propertyPath[0].equals(VERSION) || propertyPath[0].equals(VALUE))) {
+      if (propertyPath.length == 1 && (propertyPath[0].equals(VERSION) || propertyPath[0].equals(VALUE) ||
+            propertyPath[0].equals(SCORE))) {
          return true;
       }
 

diff --git a/...c/main/java/org/infinispan/objectfilter/impl/syntax/parser/QueryRendererDelegateImpl.java b/...c/main/java/org/infinispan/objectfilter/impl/syntax/parser/QueryRendererDelegateImpl.java
@@ -18,6 +18,7 @@
 import org.infinispan.objectfilter.impl.syntax.ConstantValueExpr;
 import org.infinispan.objectfilter.impl.syntax.IndexedFieldProvider;
 import org.infinispan.objectfilter.impl.syntax.parser.projection.CacheValuePropertyPath;
+import org.infinispan.objectfilter.impl.syntax.parser.projection.ScorePropertyPath;
 import org.infinispan.objectfilter.impl.syntax.parser.projection.VersionPropertyPath;
 import org.jboss.logging.Logger;
 
@@ -639,8 +640,7 @@ public void projectScore() {
          projectedNullMarkers = new ArrayList<>(ARRAY_INITIAL_LENGTH);
       }
 
-      // TODO ISPN-15390 Project the score (not the version!)
-      projections.add(new VersionPropertyPath<>());
+      projections.add(new ScorePropertyPath<>());
       projectedTypes.add(Object.class); // Usually a core module EntryVersion
       projectedNullMarkers.add(null);
    }

diff --git a/...ain/java/org/infinispan/objectfilter/impl/syntax/parser/projection/ScorePropertyPath.java b/...ain/java/org/infinispan/objectfilter/impl/syntax/parser/projection/ScorePropertyPath.java
@@ -0,0 +1,14 @@
+package org.infinispan.objectfilter.impl.syntax.parser.projection;
+
+import java.util.Collections;
+
+import org.infinispan.objectfilter.impl.ql.PropertyPath;
+
+public class ScorePropertyPath<TypeDescriptor> extends PropertyPath<TypeDescriptor> {
+
+   public static final String SCORE_PROPERTY_NAME = "__ISPN_Score";
+
+   public ScorePropertyPath() {
+      super(Collections.singletonList(new PropertyReference<>(SCORE_PROPERTY_NAME, null, true)));
+   }
+}
diff --git a/query-core/src/main/java/org/infinispan/query/core/impl/MappingEntryIterator.java b/query-core/src/main/java/org/infinispan/query/core/impl/MappingEntryIterator.java
@@ -1,23 +1,23 @@
 package org.infinispan.query.core.impl;
 
 import java.util.NoSuchElementException;
-import java.util.function.BiFunction;
+import java.util.function.Function;
 
 import org.infinispan.commons.api.query.EntityEntry;
 import org.infinispan.commons.util.CloseableIterator;
 
 public class MappingEntryIterator<K, S, T> implements CloseableIterator<T> {
 
    private final CloseableIterator<EntityEntry<K, S>> entryIterator;
-   private final BiFunction<K, S, T> mapper;
+   private final Function<EntityEntry<K, S>, T> mapper;
 
    private long skip = 0;
    private long max = -1;
 
    private T current;
    private long index;
 
-   public MappingEntryIterator(CloseableIterator<EntityEntry<K, S>> entryIterator, BiFunction<K, S, T> mapper) {
+   public MappingEntryIterator(CloseableIterator<EntityEntry<K, S>> entryIterator, Function<EntityEntry<K, S>, T> mapper) {
       this.entryIterator = entryIterator;
       this.mapper = mapper;
    }
@@ -42,7 +42,7 @@ public T next() {
    private void updateNext() {
       while (current == null && entryIterator.hasNext()) {
          EntityEntry<K, S> next = entryIterator.next();
-         T mapped = transform(next.key(), next.value());
+         T mapped = transform(next);
          if (mapped != null) {
             index++;
          }
@@ -52,15 +52,15 @@ private void updateNext() {
       }
    }
 
-   private T transform(K k, S s) {
-      if (s == null) {
+   private T transform(EntityEntry<K, S> next) {
+      if (next == null) {
          return null;
       }
       if (mapper == null) {
-         return (T) s;
+         return (T) next.value();
       }
 
-      return mapper.apply(k, s);
+      return mapper.apply(next);
    }
 
    public MappingEntryIterator<K, S, T> skip(long skip) {

diff --git a/query-core/src/main/java/org/infinispan/query/core/impl/MetadataHybridQuery.java b/query-core/src/main/java/org/infinispan/query/core/impl/MetadataHybridQuery.java
@@ -7,6 +7,7 @@
 import org.infinispan.commons.util.CloseableIterator;
 import org.infinispan.objectfilter.ObjectFilter;
 import org.infinispan.objectfilter.impl.syntax.parser.IckleParsingResult;
+import org.infinispan.objectfilter.impl.syntax.parser.projection.ScorePropertyPath;
 import org.infinispan.query.core.stats.impl.LocalQueryStatistics;
 import org.infinispan.query.dsl.Query;
 import org.infinispan.query.dsl.QueryFactory;
@@ -24,6 +25,22 @@ public MetadataHybridQuery(QueryFactory queryFactory, AdvancedCache<?, ?> cache,
    protected CloseableIterator<ObjectFilter.FilterResult> getInternalIterator() {
       CloseableIterator<EntityEntry<Object, S>> iterator = baseQuery
             .startOffset(0).maxResults(Integer.MAX_VALUE).local(local).entryIterator();
-      return new MappingEntryIterator<>(iterator, objectFilter::filter);
+      return new MappingEntryIterator<>(iterator, this::filter);
+   }
+
+   private ObjectFilter.FilterResult filter(EntityEntry entry) {
+      ObjectFilter.FilterResult filter = objectFilter.filter(entry.key(), entry.value());
+      String[] projection = objectFilter.getProjection();
+      if (projection == null) {
+         return filter;
+      }
+
+      for (int i=0; i<projection.length; i++) {
+         if (ScorePropertyPath.SCORE_PROPERTY_NAME.equals(projection[i])) {
+            filter.getProjection()[i] = entry.score();
+         }
+      }
+
+      return filter;
    }
 }
diff --git a/query/src/main/java/org/infinispan/query/clustered/DistributedEntryIterator.java b/query/src/main/java/org/infinispan/query/clustered/DistributedEntryIterator.java
@@ -23,7 +23,7 @@ class DistributedEntryIterator<K, V> extends DistributedIterator<EntityEntry<K,
    }
 
    @Override
-   protected EntityEntry<K, V> decorate(Object key, Object value) {
-      return new EntityEntry<>((K) key, (V) value, 0f);
+   protected EntityEntry<K, V> decorate(Object key, Object value, float score) {
+      return new EntityEntry<>((K) key, (V) value, score);
    }
 }
diff --git a/query/src/main/java/org/infinispan/query/clustered/DistributedIterator.java b/query/src/main/java/org/infinispan/query/clustered/DistributedIterator.java
@@ -129,13 +129,13 @@ public final T next() {
 
       if (queryStatistics.isEnabled()) queryStatistics.entityLoaded(System.nanoTime() - start);
 
-      return decorate(key, value);
+      return decorate(key, value, scoreDoc.score);
    }
 
    /**
     * Extension point for subclasses.
     */
-   protected T decorate(Object key, Object value) {
+   protected T decorate(Object key, Object value, float score) {
       return (T) value;
    }
 

diff --git a/...y/src/main/java/org/infinispan/query/dsl/embedded/impl/HibernateSearchPropertyHelper.java b/...y/src/main/java/org/infinispan/query/dsl/embedded/impl/HibernateSearchPropertyHelper.java
@@ -15,6 +15,7 @@
 import org.infinispan.objectfilter.impl.syntax.parser.EntityNameResolver;
 import org.infinispan.objectfilter.impl.syntax.parser.ReflectionPropertyHelper;
 import org.infinispan.objectfilter.impl.syntax.parser.projection.CacheValuePropertyPath;
+import org.infinispan.objectfilter.impl.syntax.parser.projection.ScorePropertyPath;
 import org.infinispan.objectfilter.impl.syntax.parser.projection.VersionPropertyPath;
 import org.infinispan.objectfilter.impl.util.StringHelper;
 import org.infinispan.search.mapper.mapping.SearchIndexedEntity;
@@ -25,6 +26,7 @@ public class HibernateSearchPropertyHelper extends ReflectionPropertyHelper {
    public static final String KEY = "__ISPN_Key";
    public static final String VALUE = CacheValuePropertyPath.VALUE_PROPERTY_NAME;
    public static final String VERSION = VersionPropertyPath.VERSION_PROPERTY_NAME;
+   public static final String SCORE = ScorePropertyPath.SCORE_PROPERTY_NAME;
 
    private final SearchMapping searchMapping;
 
@@ -54,8 +56,13 @@ public Object convertToPropertyType(Class<?> entityType, String[] propertyPath,
 
    @Override
    public Class<?> getPrimitivePropertyType(Class<?> entityType, String[] propertyPath) {
-      if (propertyPath.length == 1 && propertyPath[0].equals(VERSION)) {
-         return EntryVersion.class;
+      if (propertyPath.length == 1) {
+         if (propertyPath[0].equals(VERSION)) {
+            return EntryVersion.class;
+         }
+         if (propertyPath[0].equals(SCORE)) {
+            return Float.class;
+         }
       }
 
       IndexValueFieldDescriptor fieldDescriptor = getValueFieldDescriptor(entityType, propertyPath);
@@ -107,7 +114,7 @@ public boolean hasProperty(Class<?> entityType, String[] propertyPath) {
       }
 
       if (propertyPath.length == 1 && (propertyPath[0].equals(KEY) || propertyPath[0].equals(VALUE) ||
-            propertyPath[0].equals(VERSION)) ) {
+            propertyPath[0].equals(VERSION) || propertyPath[0].equals(SCORE)) ) {
             return true;
       }
 

diff --git a/query/src/main/java/org/infinispan/query/dsl/embedded/impl/SearchQueryBuilder.java b/query/src/main/java/org/infinispan/query/dsl/embedded/impl/SearchQueryBuilder.java
@@ -12,7 +12,6 @@
 import org.hibernate.search.backend.lucene.LuceneExtension;
 import org.hibernate.search.backend.lucene.search.query.LuceneSearchQuery;
 import org.hibernate.search.backend.lucene.search.query.dsl.LuceneSearchQueryOptionsStep;
-import org.hibernate.search.engine.backend.common.DocumentReference;
 import org.hibernate.search.engine.common.EntityReference;
 import org.hibernate.search.engine.search.aggregation.AggregationKey;
 import org.hibernate.search.engine.search.predicate.SearchPredicate;
@@ -72,15 +71,12 @@ public LuceneSearchQuery<List<Object>> keyAndEntity() {
       SearchProjectionFactory<EntityReference, ?> projectionFactory = scope.projection();
       SearchProjection<?>[] searchProjections = new SearchProjection<?>[]{
             projectionFactory.entityReference().toProjection(),
-            projectionFactory.entity().toProjection()
+            projectionFactory.entity().toProjection(),
+            projectionFactory.score().toProjection()
       };
       return build((SearchProjection<List<Object>>) SearchProjectionInfo.composite(projectionFactory, searchProjections).getProjection());
    }
 
-   public LuceneSearchQuery<DocumentReference> documentReference() {
-      return build(scope.projection().documentReference().toProjection());
-   }
-
    public void routeOnSegments(BitSet segments) {
       routingKeys = segments.stream()
             .mapToObj(String::valueOf)

diff --git a/query/src/main/java/org/infinispan/query/dsl/embedded/impl/SearchQueryMaker.java b/query/src/main/java/org/infinispan/query/dsl/embedded/impl/SearchQueryMaker.java
@@ -2,6 +2,7 @@
 
 import static org.infinispan.query.dsl.embedded.impl.HibernateSearchPropertyHelper.KEY;
 import static org.infinispan.query.dsl.embedded.impl.HibernateSearchPropertyHelper.VALUE;
+import static org.infinispan.query.dsl.embedded.impl.HibernateSearchPropertyHelper.SCORE;
 import static org.infinispan.query.logging.Log.CONTAINER;
 
 import java.io.IOException;
@@ -190,6 +191,8 @@ private SearchProjectionInfo makeProjection(TypeMetadata typeMetadata, SearchPro
             searchProjections[i] = projectionFactory.entity().toProjection();
          } else if (KEY.equals(projections[i])) {
             searchProjections[i] = projectionFactory.entityReference().toProjection();
+         } else if (SCORE.equals(projections[i])) {
+            searchProjections[i] = projectionFactory.score().toProjection();
          } else {
             boolean isMultiField = propertyHelper.isRepeatedProperty(typeMetadata, FieldPaths.split(projections[i]));
             FieldProjectionValueStep<?, ?> projectionStep = projectionFactory.field(projections[i], projectedTypes[i]);

diff --git a/query/src/main/java/org/infinispan/query/impl/IndexedQueryImpl.java b/query/src/main/java/org/infinispan/query/impl/IndexedQueryImpl.java
@@ -138,7 +138,7 @@ public <K> CloseableIterator<EntityEntry<K, E>> entryIterator() {
    }
 
    private <K, V> EntityEntry<K, V> mapToEntry(List<Object> projection) {
-      return new EntityEntry<>((K) ((EntityReference) projection.get(0)).id(), (V) projection.get(1), 0f);
+      return new EntityEntry<>((K) ((EntityReference) projection.get(0)).id(), (V) projection.get(1), (float) projection.get(2));
    }
 
    @Override

diff --git a/query/src/test/java/org/infinispan/query/projection/ScoreProjectionTest.java b/query/src/test/java/org/infinispan/query/projection/ScoreProjectionTest.java
@@ -56,7 +56,7 @@ public void entityAndScore() {
       query = cache.query("select g, score(g) from org.infinispan.query.model.Game g where g.description : 'description'~2");
       games = query.list();
       assertThat(games).extracting(objects -> objects[0]).extracting("name").containsExactly("1", "2", "3", "4");
-      //TODO ISPN-15390 Implement score projection
+      assertThat(games).extracting(objects -> objects[1]).hasOnlyElementsOfType(Float.class).isNotNull();
    }
 
    @Test
@@ -67,7 +67,7 @@ public void fieldAndScore() {
       query = cache.query("select score(g), g.name from org.infinispan.query.model.Game g where g.description : 'description'~2");
       games = query.list();
       assertThat(games).extracting(objects -> objects[1]).containsExactly("1", "2", "3", "4");
-      //TODO ISPN-15390 Implement score projection
+      assertThat(games).extracting(objects -> objects[0]).hasOnlyElementsOfType(Float.class).isNotNull();
    }
 
    @Test
@@ -78,8 +78,8 @@ public void entityScoreAndField() {
       query = cache.query("select g.name, score(g), g from org.infinispan.query.model.Game g where g.description : 'description'~2");
       games = query.list();
       assertThat(games).extracting(objects -> objects[0]).containsExactly("1", "2", "3", "4");
+      assertThat(games).extracting(objects -> objects[1]).hasOnlyElementsOfType(Float.class).isNotNull();
       assertThat(games).extracting(objects -> objects[2]).extracting("name").containsExactly("1", "2", "3", "4");
-      //TODO ISPN-15390 Implement score projection
    }
 
    @Test