Resolve table references before pruning, when the access to the processing state is still there

hibernate-core/src/main/java/org/hibernate/persister/entity/JoinedSubclassEntityPersister.java

@@ -58,6 +58,7 @@ import org.hibernate.sql.ast.tree.from.NamedTableReference;
 import org.hibernate.sql.ast.tree.from.TableGroup;
 import org.hibernate.sql.ast.tree.from.TableReference;
 import org.hibernate.sql.ast.tree.from.TableReferenceJoin;
+import org.hibernate.sql.ast.tree.from.UnknownTableReferenceException;
 import org.hibernate.sql.model.ast.builder.MutationGroupBuilder;
 import org.hibernate.sql.model.ast.builder.TableInsertBuilder;
 import org.hibernate.sql.results.graph.DomainResult;
@@ -1336,10 +1337,12 @@ public class JoinedSubclassEntityPersister extends AbstractEntityPersister {
 			}
 			// Add the table references for all table names of the treated entities as we have to retain these table references.
 			// Table references not appearing in this set can later be pruned away
-			// todo (6.0): no need to resolve all table references, only the ones needed for cardinality
 			for ( String subclassTableName : subclassTableNames ) {
 				final TableReference tableReference =
-						tableGroup.resolveTableReference( null, subclassTableName, false );
+						tableGroup.getTableReference( null, subclassTableName, false, false );
+				if ( tableReference == null ) {
+					throw new UnknownTableReferenceException( getRootTableName(), "Couldn't find table reference" );
+				}
 				retainedTableReferences.add( tableReference );
 			}
 		}

hibernate-core/src/main/java/org/hibernate/persister/entity/UnionSubclassEntityPersister.java

@@ -59,6 +59,7 @@ import org.hibernate.sql.ast.tree.from.NamedTableReference;
 import org.hibernate.sql.ast.tree.from.TableGroup;
 import org.hibernate.sql.ast.tree.from.UnionTableGroup;
 import org.hibernate.sql.ast.tree.from.UnionTableReference;
+import org.hibernate.sql.ast.tree.from.UnknownTableReferenceException;
 import org.hibernate.sql.ast.tree.predicate.Predicate;
 import org.hibernate.type.BasicType;
 import org.hibernate.type.StandardBasicTypes;
@@ -416,7 +417,10 @@ public class UnionSubclassEntityPersister extends AbstractEntityPersister {
 
 	@Override
 	public void pruneForSubclasses(TableGroup tableGroup, Set<String> treatedEntityNames) {
-		final NamedTableReference tableReference = (NamedTableReference) tableGroup.resolveTableReference( getRootTableName() );
+		final NamedTableReference tableReference = (NamedTableReference) tableGroup.getTableReference( getRootTableName() );
+		if ( tableReference == null ) {
+			throw new UnknownTableReferenceException( getRootTableName(), "Couldn't find table reference" );
+		}
 		// Replace the default union sub-query with a specially created one that only selects the tables for the treated entity names
 		tableReference.setPrunedTableExpression( generateSubquery( treatedEntityNames ) );
 	}

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

@@ -2930,6 +2930,20 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 				tg -> new HashSet<>( 1 )
 		);
 		treatedEntityNames.add( treatedType.getHibernateEntityName() );
+
+		// Resolve the table references for the tables that the treated type touches
+		final AbstractEntityPersister persister = (AbstractEntityPersister) creationContext.getSessionFactory()
+				.getRuntimeMetamodels()
+				.getEntityMappingType( treatedType.getHibernateEntityName() );
+		// Avoid doing this for single table entity persisters, as the table span includes secondary tables,
+		// which we don't want to resolve, though we know that there is only a single table anyway
+		if ( persister instanceof SingleTableEntityPersister ) {
+			return;
+		}
+		final int subclassTableSpan = persister.getSubclassTableSpan();
+		for ( int i = 0; i < subclassTableSpan; i++ ) {
+			tableGroup.resolveTableReference( null, persister.getSubclassTableName( i ), false );
+		}
 	}
 
 	protected void registerTypeUsage(DiscriminatorSqmPath path) {