HHH-16123 Add another test and fix rendering the pruned subquery in UnionSubclassEntityPersister

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

@@ -8,6 +8,7 @@ package org.hibernate.persister.entity;
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -39,6 +40,7 @@ import org.hibernate.mapping.PersistentClass;
 import org.hibernate.mapping.Subclass;
 import org.hibernate.mapping.Table;
 import org.hibernate.metamodel.mapping.EntityDiscriminatorMapping;
+import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.SelectableConsumer;
 import org.hibernate.metamodel.mapping.SelectableMapping;
 import org.hibernate.metamodel.mapping.TableDetails;
@@ -528,26 +530,16 @@ public class UnionSubclassEntityPersister extends AbstractEntityPersister {
 
 		// Collect all selectables of every entity subtype and group by selection expression as well as table name
 		final LinkedHashMap<String, Map<String, SelectableMapping>> selectables = new LinkedHashMap<>();
-		final SelectableConsumer selectableConsumer = (i, selectable) ->
-			selectables.computeIfAbsent( selectable.getSelectionExpression(), k -> new HashMap<>() )
-					.put( selectable.getContainingTableExpression(), selectable );
 		// Collect the concrete subclass table names for the treated entity names
 		final Set<String> treatedTableNames = new HashSet<>( treated.size() );
 		for ( String subclassName : treated ) {
 			final UnionSubclassEntityPersister subPersister =
 					(UnionSubclassEntityPersister) metamodel.getEntityDescriptor( subclassName );
-			for ( String subclassTableName : subPersister.getSubclassTableNames() ) {
+			// Collect all the real (non-abstract) table names
-				if ( ArrayHelper.indexOf( subclassSpaces, subclassTableName ) != -1 ) {
+			treatedTableNames.addAll( Arrays.asList( subPersister.getConstraintOrderedTableNameClosure() ) );
-					treatedTableNames.add( subclassTableName );
+			// Collect selectables grouped by the table names in which they appear
-				}
+			// TODO: we could cache this
-			}
+			subPersister.collectSelectableOwners( selectables );
-			subPersister.getIdentifierMapping().forEachSelectable( selectableConsumer );
-			if ( subPersister.getVersionMapping() != null ) {
-				subPersister.getVersionMapping().forEachSelectable( selectableConsumer );
-			}
-			subPersister.visitSubTypeAttributeMappings(
-					attributeMapping -> attributeMapping.forEachSelectable( selectableConsumer )
-			);
 		}
 
 		// Create a union sub-query for the table names, like generateSubquery(PersistentClass model, Mapping mapping)
@@ -555,8 +547,8 @@ public class UnionSubclassEntityPersister extends AbstractEntityPersister {
 				.append( "( " );
 		final StringBuilderSqlAppender sqlAppender = new StringBuilderSqlAppender( buf );
 
-		for ( String name : getSubclassEntityNames() ) {
+		for ( EntityMappingType mappingType : getSubMappingTypes() ) {
-			final AbstractEntityPersister persister = (AbstractEntityPersister) metamodel.findEntityDescriptor( name );
+			final AbstractEntityPersister persister = (AbstractEntityPersister) mappingType;
 			final String subclassTableName = persister.getTableName();
 			if ( treatedTableNames.contains( subclassTableName ) ) {
 				if ( buf.length() > 2 ) {
@@ -587,6 +579,34 @@ public class UnionSubclassEntityPersister extends AbstractEntityPersister {
 		return buf.append( " )" ).toString();
 	}
 
+	private void collectSelectableOwners(LinkedHashMap<String, Map<String, SelectableMapping>> selectables) {
+		if ( isAbstract() ) {
+			for ( EntityMappingType subMappingType : getSubMappingTypes() ) {
+				if ( !subMappingType.isAbstract() ) {
+					( (UnionSubclassEntityPersister) subMappingType ).collectSelectableOwners( selectables );
+				}
+			}
+		}
+		else {
+			final SelectableConsumer selectableConsumer = (i, selectable) -> {
+				Map<String, SelectableMapping> selectableMapping = selectables.computeIfAbsent(
+						selectable.getSelectionExpression(),
+						k -> new HashMap<>()
+				);
+				selectableMapping.put( getTableName(), selectable );
+			};
+			getIdentifierMapping().forEachSelectable( selectableConsumer );
+			if ( getVersionMapping() != null ) {
+				getVersionMapping().forEachSelectable( selectableConsumer );
+			}
+			final AttributeMappingsList attributeMappings = getAttributeMappings();
+			final int size = attributeMappings.size();
+			for ( int i = 0; i < size; i++ ) {
+				attributeMappings.get( i ).forEachSelectable( selectableConsumer );
+			}
+		}
+	}
+
 	@Override
 	protected String[] getSubclassTableKeyColumns(int j) {
 		if ( j != 0 ) {

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/criteria/TreatAbstractSuperclassTest.java

@@ -28,12 +28,14 @@ import jakarta.persistence.JoinColumn;
 import jakarta.persistence.ManyToOne;
 import jakarta.persistence.MappedSuperclass;
 import jakarta.persistence.OneToMany;
+import jakarta.persistence.Tuple;
 import jakarta.persistence.criteria.CriteriaBuilder;
 import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.Join;
 import jakarta.persistence.criteria.Root;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 
 /**
  * @author Marco Belladelli
@@ -42,6 +44,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 @DomainModel(annotatedClasses = {
 		TreatAbstractSuperclassTest.LongBook.class,
 		TreatAbstractSuperclassTest.ShortBook.class,
+		TreatAbstractSuperclassTest.Article.class,
 		TreatAbstractSuperclassTest.Author.class,
 		TreatAbstractSuperclassTest.AuthorParticipation.class,
 })
@@ -85,6 +88,25 @@ public class TreatAbstractSuperclassTest {
 		} );
 	}
 
+	@Test
+	public void testTreatMultiple(SessionFactoryScope scope) {
+		scope.inTransaction( session -> {
+			final CriteriaBuilder cb = session.getCriteriaBuilder();
+			final CriteriaQuery<Tuple> criteria = cb.createTupleQuery();
+			final Root<Publication> publicationRoot = criteria.from( Publication.class );
+			// Treat as nested abstract superclass
+			final Root<Book> bookRoot = cb.treat( publicationRoot, Book.class );
+			final Root<Article> articleRoot = cb.treat( publicationRoot, Article.class );
+			criteria.multiselect(
+					bookRoot.get( "title" ),
+					articleRoot.get( "reference" )
+			);
+			final Tuple tuple = session.createQuery( criteria ).getSingleResult();
+			assertEquals( "Dune", tuple.get( 0 ) );
+			assertNull( tuple.get( 1 ) );
+		} );
+	}
+
 	@Test
 	public void testSubclassJoin(SessionFactoryScope scope) {
 		scope.inTransaction( session -> {
@@ -122,8 +144,29 @@ public class TreatAbstractSuperclassTest {
 		}
 	}
 
+	@Entity(name = "Article")
+	public static class Article extends Publication {
+		private String reference;
+
+		public Article() {
+		}
+
+		public Article(String title) {
+			super( title );
+		}
+
+		public String getReference() {
+			return reference;
+		}
+
+		public void setReference(String reference) {
+			this.reference = reference;
+		}
+	}
+
 	@Entity(name = "Book")
 	public static abstract class Book extends Publication {
+		private String isbn;
 		@OneToMany(mappedBy = "book", cascade = CascadeType.REMOVE)
 		private List<AuthorParticipation> participations = new ArrayList<>();
 
@@ -141,6 +184,7 @@ public class TreatAbstractSuperclassTest {
 
 	@Entity(name = "LongBook")
 	public static class LongBook extends Book {
+		private int pageCount;
 		public LongBook() {
 		}
 
@@ -151,6 +195,7 @@ public class TreatAbstractSuperclassTest {
 
 	@Entity(name = "ShortBook")
 	public static class ShortBook extends Book {
+		private int readTime;
 		public ShortBook() {
 		}