3333import com .carrotsearch .hppc .cursors .IntObjectCursor ;
3434import io .evitadb .api .EntityCollectionContract ;
3535import io .evitadb .api .EvitaSessionContract ;
36- import io .evitadb .api .query .Constraint ;
3736import io .evitadb .api .query .FilterConstraint ;
37+ import io .evitadb .api .query .OrderConstraint ;
3838import io .evitadb .api .query .filter .EntityHaving ;
39+ import io .evitadb .api .query .filter .EntityLocaleEquals ;
3940import io .evitadb .api .query .filter .EntityPrimaryKeyInSet ;
4041import io .evitadb .api .query .filter .FilterBy ;
4142import io .evitadb .api .query .filter .ReferenceHaving ;
43+ import io .evitadb .api .query .order .EntityPrimaryKeyExact ;
44+ import io .evitadb .api .query .order .EntityPrimaryKeyInFilter ;
45+ import io .evitadb .api .query .order .EntityPrimaryKeyNatural ;
4246import io .evitadb .api .query .order .EntityProperty ;
4347import io .evitadb .api .query .order .OrderBy ;
4448import io .evitadb .api .query .require .DefaultPrefetchRequirementCollector ;
141145
142146import static io .evitadb .api .query .QueryConstraints .and ;
143147import static io .evitadb .api .query .QueryConstraints .entityPrimaryKeyInSet ;
148+ import static io .evitadb .api .query .QueryConstraints .filterBy ;
149+ import static io .evitadb .api .query .QueryConstraints .orderBy ;
144150import static io .evitadb .core .query .extraResult .translator .hierarchyStatistics .AbstractHierarchyTranslator .stopAtConstraintToPredicate ;
145151import static java .util .Optional .ofNullable ;
146152
@@ -264,12 +270,17 @@ private static Map<Integer, ServerEntityDecorator> fetchReferencedEntities(
264270 private static FilterBy unwrapFilterBy (@ Nullable FilterBy entityFilterBy ) {
265271 final FilterBy unwrappedEntityFilterBy ;
266272 if (entityFilterBy != null ) {
267- final Constraint <?> entityHavingConstraint = FinderVisitor .findConstraint (entityFilterBy , EntityHaving .class ::isInstance , EntityHaving .class ::isInstance );
268- if (entityHavingConstraint instanceof EntityHaving eh ) {
269- unwrappedEntityFilterBy = new FilterBy (eh .getChildren ());
270- } else {
271- unwrappedEntityFilterBy = null ;
272- }
273+ final List <FilterConstraint > entityConstraints = FinderVisitor .findConstraints (
274+ entityFilterBy ,
275+ it -> it instanceof EntityHaving || it instanceof EntityPrimaryKeyInSet || it instanceof EntityLocaleEquals ,
276+ EntityHaving .class ::isInstance
277+ );
278+ unwrappedEntityFilterBy = filterBy (
279+ entityConstraints
280+ .stream ()
281+ .flatMap (it -> it instanceof EntityHaving eh ? Arrays .stream (eh .getChildren ()) : Stream .of (it ))
282+ .toArray (FilterConstraint []::new )
283+ );
273284 } else {
274285 unwrappedEntityFilterBy = null ;
275286 }
@@ -288,12 +299,20 @@ private static FilterBy unwrapFilterBy(@Nullable FilterBy entityFilterBy) {
288299 private static OrderBy unwrapOrderBy (@ Nullable OrderBy entityOrderBy ) {
289300 final OrderBy unwrappedEntityFilterBy ;
290301 if (entityOrderBy != null ) {
291- final Constraint <?> entityProperty = FinderVisitor .findConstraint (entityOrderBy , EntityProperty .class ::isInstance , EntityProperty .class ::isInstance );
292- if (entityProperty instanceof EntityProperty ep ) {
293- unwrappedEntityFilterBy = new OrderBy (ep .getChildren ());
294- } else {
295- unwrappedEntityFilterBy = null ;
296- }
302+ final List <OrderConstraint > entityConstraints = FinderVisitor .findConstraints (
303+ entityOrderBy ,
304+ it -> it instanceof EntityProperty ||
305+ it instanceof EntityPrimaryKeyExact ||
306+ it instanceof EntityPrimaryKeyNatural ||
307+ it instanceof EntityPrimaryKeyInFilter ,
308+ EntityProperty .class ::isInstance
309+ );
310+ unwrappedEntityFilterBy = orderBy (
311+ entityConstraints
312+ .stream ()
313+ .flatMap (it -> it instanceof EntityProperty eh ? Arrays .stream (eh .getChildren ()) : Stream .of (it ))
314+ .toArray (OrderConstraint []::new )
315+ );
297316 } else {
298317 unwrappedEntityFilterBy = null ;
299318 }
0 commit comments