diff --git a/hibernate-core/src/main/java/org/hibernate/persister/entity/JoinedSubclassEntityPersister.java b/hibernate-core/src/main/java/org/hibernate/persister/entity/JoinedSubclassEntityPersister.java
index 8c727e7a19..58afae93cc 100644
--- a/hibernate-core/src/main/java/org/hibernate/persister/entity/JoinedSubclassEntityPersister.java
+++ b/hibernate-core/src/main/java/org/hibernate/persister/entity/JoinedSubclassEntityPersister.java
@@ -1392,7 +1392,10 @@ public class JoinedSubclassEntityPersister extends AbstractEntityPersister {
 					applied = applyDiscriminatorPredicate( join, join.getJoinedTableReference(), entityNameUses, metamodel );
 				}
 				assert applied : "Could not apply treat discriminator predicate to root table join";
-				assert i == 0 || retainedTableReferences.contains( tableReferenceJoins.get( i - 1 ).getJoinedTableReference() );
+				if ( i != 0 ) {
+					// Always retain the root table reference join where the discriminator was applied
+					retainedTableReferences.add( tableReferenceJoins.get( i - 1 ).getJoinedTableReference() );
+				}
 			}
 		}
 		if ( tableReferenceJoins.isEmpty() ) {
diff --git a/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java b/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java
index 4f9f5ce274..0a98c28258 100644
--- a/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java
+++ b/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java
@@ -3105,7 +3105,6 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 			actualTableGroup.resolveTableReference( null, persister.getTableName() );
 		}
 
-		final EntityNameUse.UseKind useKind = finalEntityNameUse.getKind();
 		if ( projection ) {
 			EntityMappingType superMappingType = persister;
 			while ( ( superMappingType = superMappingType.getSuperMappingType() ) != null ) {
@@ -3116,8 +3115,10 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 				);
 			}
 		}
+
 		// If we encounter a treat or projection use, we also want register the use for all subtypes.
 		// We do this here to not have to expand entity name uses during pruning later on
+		final EntityNameUse.UseKind useKind = finalEntityNameUse.getKind();
 		if ( useKind == EntityNameUse.UseKind.TREAT ) {
 			for ( EntityMappingType subType : persister.getSubMappingTypes() ) {
 				entityNameUses.compute(
@@ -3125,10 +3126,6 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 						(s, existingUse) -> finalEntityNameUse.stronger( existingUse )
 				);
 			}
-			if ( persister.isInherited() && persister.needsDiscriminator() ) {
-				// Make sure the table group includes the root table when needed for TREAT
-				actualTableGroup.resolveTableReference( persister.getRootTableName() );
-			}
 		}
 		else if ( useKind == EntityNameUse.UseKind.PROJECTION ) {
 			for ( EntityMappingType subType : persister.getSubMappingTypes() ) {
@@ -3374,7 +3371,7 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 				// This is a non-treated join with an entity which is an inheritance subtype,
 				// register a TREAT entity name use to filter only the entities of the correct type.
 				registerEntityNameUsage(
-						getActualTableGroup( joinedTableGroup, sqmJoin ),
+						elementTableGroup,
 						EntityNameUse.TREAT,
 						entityDomainType.getHibernateEntityName()
 				);