HHH-16409 Rework entity valued path expansion for group by and order by

Author: mbellade

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java

@@ -751,6 +751,10 @@ public Stack<Clause> getCurrentClauseStack() {
 		return currentClauseStack;
 	}
 
+	@Override
+	public SqmQueryPart<?> getCurrentSqmQueryPart() {
+		return currentSqmQueryPart;
+	}
 
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 	// Statements
@@ -2395,46 +2399,6 @@ else if ( selectableNode instanceof SqmJpaCompoundSelection<?> ) {
 		return -( offset + selections.size() );
 	}
 
-	private boolean selectClauseContains(SqmFrom<?, ?> from) {
-		final SqmQuerySpec<?> sqmQuerySpec = (SqmQuerySpec<?>) currentSqmQueryPart;
-		final List<SqmSelection<?>> selections = sqmQuerySpec.getSelectClause() == null
-				? Collections.emptyList()
-				: sqmQuerySpec.getSelectClause().getSelections();
-		if ( selections.isEmpty() && from instanceof SqmRoot<?> ) {
-			return true;
-		}
-		for ( SqmSelection<?> selection : selections ) {
-			if ( selectableNodeContains( selection.getSelectableNode(), from ) ) {
-				return true;
-			}
-		}
-		return false;
-	}
-
-	private boolean selectableNodeContains(SqmSelectableNode<?> selectableNode, SqmFrom<?, ?> from) {
-		if ( selectableNode == from ) {
-			return true;
-		}
-		else if ( selectableNode instanceof SqmDynamicInstantiation ) {
-			for ( SqmDynamicInstantiationArgument<?> argument : ( (SqmDynamicInstantiation<?>) selectableNode ).getArguments() ) {
-				if ( selectableNodeContains( argument.getSelectableNode(), from ) ) {
-					return true;
-				}
-			}
-		}
-		return false;
-	}
-
-	private boolean groupByClauseContains(SqmFrom<?, ?> from) {
-		final SqmQuerySpec<?> sqmQuerySpec = (SqmQuerySpec<?>) currentSqmQueryPart;
-		for ( SqmExpression<?> expression : sqmQuerySpec.getGroupByClauseExpressions() ) {
-			if ( expression == from ) {
-				return true;
-			}
-		}
-		return false;
-	}
-
 	@Override
 	public List<Expression> visitGroupByClause(List<SqmExpression<?>> groupByClauseExpressions) {
 		if ( !groupByClauseExpressions.isEmpty() ) {
@@ -3755,8 +3719,8 @@ private Expression visitTableGroup(TableGroup tableGroup, SqmFrom<?, ?> path) {
 			final EntityValuedModelPart interpretationModelPart;
 			final TableGroup tableGroupToUse;
 			if ( inferredEntityMapping == null ) {
-				// When the inferred mapping is null, we try to resolve to the FK by default,
-				// which is fine because expansion to all target columns for select and group by clauses is handled below
+				// When the inferred mapping is null, we try to resolve to the FK by default, which is fine because
+				// expansion to all target columns for select and group by clauses is handled in EntityValuedPathInterpretation
 				if ( entityValuedModelPart instanceof EntityAssociationMapping && ( (EntityAssociationMapping) entityValuedModelPart ).isFkOptimizationAllowed() ) {
 					// If the table group uses an association mapping that is not a one-to-many,
 					// we make use of the FK model part
@@ -3883,22 +3847,6 @@ else if ( entityValuedModelPart instanceof AnonymousTupleEntityValuedModelPart )
 				tableGroupToUse = null;
 			}
 
-			final boolean expandToAllColumns;
-			if ( currentClauseStack.getCurrent() == Clause.GROUP ) {
-				// When the table group is known to be fetched i.e. a fetch join
-				// but also when the from clause is part of the select clause
-				// we need to expand to all columns, as we also expand this to all columns in the select clause
-				expandToAllColumns = tableGroup.isFetched() || selectClauseContains( path );
-			}
-			else if ( currentClauseStack.getCurrent() == Clause.ORDER ) {
-				// We must ensure that the order by expression be expanded if the group by
-				// contained the same expression, and that was expanded as well
-				expandToAllColumns = groupByClauseContains( path ) && ( tableGroup.isFetched() || selectClauseContains( path ) );
-			}
-			else {
-				expandToAllColumns = false;
-			}
-
 			final EntityMappingType treatedMapping;
 			if ( path instanceof SqmTreatedPath ) {
 				treatedMapping = creationContext.getSessionFactory()
@@ -3913,7 +3861,7 @@ else if ( currentClauseStack.getCurrent() == Clause.ORDER ) {
 			result = EntityValuedPathInterpretation.from(
 					navigablePath,
 					tableGroupToUse == null ? tableGroup : tableGroupToUse,
-					expandToAllColumns ? null : resultModelPart,
+					resultModelPart,
 					interpretationModelPart,
 					treatedMapping,
 					this

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/FakeSqmToSqlAstConverter.java

@@ -18,6 +18,7 @@
 import org.hibernate.query.sqm.tree.expression.SqmExpression;
 import org.hibernate.query.sqm.tree.expression.SqmParameter;
 import org.hibernate.query.sqm.tree.predicate.SqmPredicate;
+import org.hibernate.query.sqm.tree.select.SqmQueryPart;
 import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.spi.FromClauseAccess;
 import org.hibernate.sql.ast.spi.SqlAliasBaseGenerator;
@@ -86,6 +87,11 @@ public Stack<Clause> getCurrentClauseStack() {
 		return null;
 	}
 
+	@Override
+	public SqmQueryPart<?> getCurrentSqmQueryPart() {
+		return null;
+	}
+
 	@Override
 	public void registerQueryTransformer(QueryTransformer transformer) {
 	}

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/SqmToSqlAstConverter.java

@@ -16,6 +16,7 @@
 import org.hibernate.query.sqm.tree.expression.SqmExpression;
 import org.hibernate.query.sqm.tree.expression.SqmParameter;
 import org.hibernate.query.sqm.tree.predicate.SqmPredicate;
+import org.hibernate.query.sqm.tree.select.SqmQueryPart;
 import org.hibernate.sql.ast.spi.SqlAstCreationState;
 import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.tree.expression.Expression;
@@ -30,6 +31,8 @@
 public interface SqmToSqlAstConverter extends SemanticQueryWalker<Object>, SqlAstCreationState {
 	Stack<Clause> getCurrentClauseStack();
 
+	SqmQueryPart<?> getCurrentSqmQueryPart();
+
 	void registerQueryTransformer(QueryTransformer transformer);
 
 	/**

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/internal/EntityValuedPathInterpretation.java

@@ -27,7 +27,15 @@
 import org.hibernate.query.derived.AnonymousTupleEntityValuedModelPart;
 import org.hibernate.query.sqm.sql.SqmToSqlAstConverter;
 import org.hibernate.query.sqm.tree.domain.SqmEntityValuedSimplePath;
+import org.hibernate.query.sqm.tree.domain.SqmPath;
+import org.hibernate.query.sqm.tree.expression.SqmExpression;
+import org.hibernate.query.sqm.tree.select.SqmDynamicInstantiation;
+import org.hibernate.query.sqm.tree.select.SqmDynamicInstantiationArgument;
+import org.hibernate.query.sqm.tree.select.SqmQuerySpec;
+import org.hibernate.query.sqm.tree.select.SqmSelectableNode;
+import org.hibernate.query.sqm.tree.select.SqmSelection;
 import org.hibernate.spi.NavigablePath;
+import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.SqlAstWalker;
 import org.hibernate.sql.ast.spi.FromClauseAccess;
 import org.hibernate.sql.ast.spi.SqlAstCreationState;
@@ -284,10 +292,28 @@ public static <T> EntityValuedPathInterpretation<T> from(
 			EntityValuedModelPart mapping,
 			EntityValuedModelPart treatedMapping,
 			SqmToSqlAstConverter sqlAstCreationState) {
+		final boolean expandToAllColumns;
+		final Clause currentClause = sqlAstCreationState.getCurrentClauseStack().getCurrent();
+		if ( currentClause == Clause.GROUP || currentClause == Clause.ORDER ) {
+			final SqmQuerySpec<?> querySpec = (SqmQuerySpec<?>) sqlAstCreationState.getCurrentSqmQueryPart();
+			if ( currentClause == Clause.ORDER && !groupByClauseContains( navigablePath, querySpec ) ) {
+				// We must ensure that the order by expression be expanded but only if the group by
+				// contained the same expression, and that was expanded as well
+				expandToAllColumns = false;
+			}
+			else {
+				// When the table group is selected and the navigablePath is selected we need to expand
+				// to all columns, as we also expand this to all columns in the select clause
+				expandToAllColumns = isSelected( tableGroup, navigablePath, querySpec );
+			}
+		}
+		else {
+			expandToAllColumns = false;
+		}
+
 		final SqlExpressionResolver sqlExprResolver = sqlAstCreationState.getSqlExpressionResolver();
 		final Expression sqlExpression;
-
-		if ( resultModelPart == null ) {
+		if ( expandToAllColumns ) {
 			// Expand to all columns of the entity mapping type, as we already did for the selection
 			final EntityMappingType entityMappingType = mapping.getEntityMappingType();
 			final EntityIdentifierMapping identifierMapping = entityMappingType.getIdentifierMapping();
@@ -352,6 +378,47 @@ public static <T> EntityValuedPathInterpretation<T> from(
 		);
 	}
 
+	private static boolean isSelected(TableGroup tableGroup, NavigablePath path, SqmQuerySpec<?> sqmQuerySpec) {
+		// If the table group is selected (initialized), check if the entity valued
+		// navigable path or any child path appears in the select clause
+		return tableGroup.isInitialized() && selectClauseContains( path, sqmQuerySpec );
+	}
+
+	private static boolean selectClauseContains(NavigablePath path, SqmQuerySpec<?> sqmQuerySpec) {
+		final List<SqmSelection<?>> selections = sqmQuerySpec.getSelectClause() == null
+				? Collections.emptyList()
+				: sqmQuerySpec.getSelectClause().getSelections();
+		for ( SqmSelection<?> selection : selections ) {
+			if ( selectableNodeContains( selection.getSelectableNode(), path ) ) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+	private static boolean selectableNodeContains(SqmSelectableNode<?> selectableNode, NavigablePath path) {
+		if ( selectableNode instanceof SqmPath && path.isParentOrEqual( ( (SqmPath<?>) selectableNode ).getNavigablePath() ) ) {
+			return true;
+		}
+		else if ( selectableNode instanceof SqmDynamicInstantiation ) {
+			for ( SqmDynamicInstantiationArgument<?> argument : ( (SqmDynamicInstantiation<?>) selectableNode ).getArguments() ) {
+				if ( selectableNodeContains( argument.getSelectableNode(), path ) ) {
+					return true;
+				}
+			}
+		}
+		return false;
+	}
+
+	private static boolean groupByClauseContains(NavigablePath path, SqmQuerySpec<?> sqmQuerySpec) {
+		for ( SqmExpression<?> expression : sqmQuerySpec.getGroupByClauseExpressions() ) {
+			if ( expression instanceof SqmPath && ( (SqmPath<?>) expression ).getNavigablePath() == path ) {
+				return true;
+			}
+		}
+		return false;
+	}
+
 	private final Expression sqlExpression;
 
 	public EntityValuedPathInterpretation(