HHH-18583 Joined + discriminator inheritance treat in where clause not restricting to subtype

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

@@ -747,6 +747,15 @@ public AbstractEntityPersister(
 		final ArrayList<Boolean> definedBySubclass = new ArrayList<>();
 		final ArrayList<Boolean> propNullables = new ArrayList<>();
 
+		if ( persistentClass.hasSubclasses() ) {
+			for ( Selectable selectable : persistentClass.getIdentifier().getSelectables() ) {
+				if ( !selectable.isFormula() ) {
+					// Identifier columns are always shared between subclasses
+					sharedColumnNames.add( ( (Column) selectable ).getQuotedName( dialect ) );
+				}
+			}
+		}
+
 		for ( Property prop : persistentClass.getSubclassPropertyClosure() ) {
 			names.add( prop.getName() );
 			types.add( prop.getType() );

hibernate-core/src/test/java/org/hibernate/orm/test/inheritance/discriminator/JoinedInheritanceDiscriminatorSelectionTest.java

@@ -38,6 +38,7 @@
 @SessionFactory( useCollectingStatementInspector = true )
 @Jira( "https://hibernate.atlassian.net/browse/HHH-17727" )
 @Jira( "https://hibernate.atlassian.net/browse/HHH-17806" )
+@Jira( "https://hibernate.atlassian.net/browse/HHH-18583" )
 public class JoinedInheritanceDiscriminatorSelectionTest {
 	@BeforeAll
 	public void setUp(SessionFactoryScope scope) {
@@ -51,7 +52,7 @@ public void setUp(SessionFactoryScope scope) {
 
 	@AfterAll
 	public void tearDown(SessionFactoryScope scope) {
-		scope.inTransaction( session -> session.createMutationQuery( "delete from ParentEntity" ).executeUpdate() );
+		scope.getSessionFactory().getSchemaManager().truncateMappedObjects();
 	}
 
 	@Test
@@ -93,6 +94,23 @@ public void testSelectParentAttribute(SessionFactoryScope scope) {
 					String.class
 			).getResultList() ).containsOnly( "parent", "child_a" );
 			inspector.assertNumberOfJoins( 0, 0 );
+			inspector.clear();
+
+			// With treat() we preserve the join
+
+			assertThat( session.createQuery(
+					"select p.name from ParentEntity p where treat(p as ChildA).id is not null",
+					String.class
+			).getResultList() ).containsExactlyInAnyOrder( "child_a", "sub_child_a" );
+			inspector.assertNumberOfJoins( 0, 1 );
+			inspector.clear();
+
+			assertThat( session.createQuery(
+					"select p.name from ParentEntity p where treat(p as ChildB).id is not null",
+					String.class
+			).getSingleResult() ).isEqualTo( "child_b" );
+			inspector.assertNumberOfJoins( 0, 1 );
+			inspector.clear();
 		} );
 	}
 
@@ -123,8 +141,9 @@ public void testSelectInstance(SessionFactoryScope scope) {
 		inspector.clear();
 
 		scope.inTransaction( session -> {
-			// NOTE: we currently always join all subclasses when selecting the entity instance. We could
-			//  maybe avoid this when we have a physical discriminator column and a type filter
+			// With type filters we still join all subclasses when selecting the entity instance
+			// because we are not aware of the type restriction when processing the selection
+
 			assertThat( session.createQuery(
 					"from ParentEntity p where type(p) = ParentEntity",
 					ParentEntity.class
@@ -144,6 +163,22 @@ public void testSelectInstance(SessionFactoryScope scope) {
 					ParentEntity.class
 			).getResultList() ).hasSize( 1 );
 			inspector.assertNumberOfJoins( 0, 3 );
+			inspector.clear();
+
+			// With treat() we only join the needed subclasses
+
+			assertThat( session.createQuery(
+					"select treat(p as ChildA) from ParentEntity p",
+					ParentEntity.class
+			).getResultList() ).hasSize( 2 );
+			inspector.assertNumberOfJoins( 0, 2 );
+			inspector.clear();
+
+			assertThat( session.createQuery(
+					"select treat(p as ChildB) from ParentEntity p",
+					ParentEntity.class
+			).getResultList() ).hasSize( 1 );
+			inspector.assertNumberOfJoins( 0, 1 );
 		} );
 	}