more efficient loading by @NaturalId

For entities with a single @NaturalId property only.
Uses a unique key EntityLoader instead of two selects.

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultResolveNaturalIdEventListener.java

@@ -7,11 +7,12 @@
 package org.hibernate.event.internal;
 
 import java.io.Serializable;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import org.hibernate.HibernateException;
-import org.hibernate.cache.spi.access.NaturalIdDataAccess;
-import org.hibernate.engine.spi.PersistenceContext;
+import org.hibernate.LockOptions;
+import org.hibernate.WrongClassException;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.event.spi.ResolveNaturalIdEvent;
@@ -19,6 +20,7 @@ import org.hibernate.event.spi.ResolveNaturalIdEventListener;
 import org.hibernate.internal.CoreLogging;
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.persister.entity.UniqueKeyLoadable;
 import org.hibernate.pretty.MessageHelper;
 import org.hibernate.stat.spi.StatisticsImplementor;
 
@@ -96,10 +98,11 @@ public class DefaultResolveNaturalIdEventListener
 	 * @return The entity from the cache, or null.
 	 */
 	protected Serializable resolveFromCache(final ResolveNaturalIdEvent event) {
-		return event.getSession().getPersistenceContextInternal().getNaturalIdHelper().findCachedNaturalIdResolution(
-				event.getEntityPersister(),
-				event.getOrderedNaturalIdValues()
-		);
+		return event.getSession().getPersistenceContextInternal().getNaturalIdHelper()
+				.findCachedNaturalIdResolution(
+						event.getEntityPersister(),
+						event.getOrderedNaturalIdValues()
+				);
 	}
 
 	/**
@@ -120,29 +123,48 @@ public class DefaultResolveNaturalIdEventListener
 			startTime = System.nanoTime();
 		}
 
-		final Serializable pk = event.getEntityPersister().loadEntityIdByNaturalId(
-				event.getOrderedNaturalIdValues(),
-				event.getLockOptions(),
-				session
-		);
+		Object[] naturalIdValues = event.getOrderedNaturalIdValues();
+
+		final Serializable pk;
+		EntityPersister persister = event.getEntityPersister();
+		LockOptions lockOptions = event.getLockOptions();
+		if ( persister instanceof UniqueKeyLoadable
+				&& naturalIdValues.length==1 ) {
+			UniqueKeyLoadable rootPersister = (UniqueKeyLoadable)
+					persister.getFactory().getMetamodel().entityPersister( persister.getRootEntityName() );
+			Map.Entry<String, Object> e = event.getNaturalIdValues().entrySet().iterator().next();
+			Object entity = rootPersister.loadByUniqueKey( e.getKey(), e.getValue(), lockOptions, session );
+			if ( entity == null ) {
+				pk = null;
+			}
+			else {
+				if ( !persister.isInstance(entity) ) {
+					throw new WrongClassException(
+							"loaded object was of wrong class " + entity.getClass(),
+							e.getKey(),
+							persister.getEntityName()
+					);
+				}
+				pk = persister.getIdentifier( entity, session );
+			}
+		}
+		else {
+			pk = persister.loadEntityIdByNaturalId( naturalIdValues, lockOptions, session );
+		}
 
 		if ( stats ) {
 			final long endTime = System.nanoTime();
 			final long milliseconds = TimeUnit.MILLISECONDS.convert( endTime - startTime, TimeUnit.NANOSECONDS );
 			statistics.naturalIdQueryExecuted(
-					event.getEntityPersister().getRootEntityName(),
+					persister.getRootEntityName(),
 					milliseconds
 			);
 		}
 
 		//PK can be null if the entity doesn't exist
 		if (pk != null) {
-			final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
-			persistenceContext.getNaturalIdHelper().cacheNaturalIdCrossReferenceFromLoad(
-					event.getEntityPersister(),
-					pk,
-					event.getOrderedNaturalIdValues()
-			);
+			session.getPersistenceContextInternal().getNaturalIdHelper()
+					.cacheNaturalIdCrossReferenceFromLoad( persister, pk, naturalIdValues );
 		}
 
 		return pk;

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

@@ -2481,6 +2481,23 @@ public abstract class AbstractEntityPersister
 		return getAppropriateUniqueKeyLoader( propertyName, session ).loadByUniqueKey( session, uniqueKey );
 	}
 
+	public Object loadByUniqueKey(
+			String propertyName,
+			Object uniqueKey,
+			LockOptions lockOptions,
+			SharedSessionContractImplementor session) throws HibernateException {
+		//TODO: cache this
+		return new EntityLoader(
+				this,
+				propertyMapping.toColumns( propertyName ),
+				propertyMapping.toType( propertyName ),
+				1,
+				lockOptions,
+				getFactory(),
+				session.getLoadQueryInfluencers()
+		).loadByUniqueKey( session, uniqueKey );
+	}
+
 	private EntityLoader getAppropriateUniqueKeyLoader(String propertyName, SharedSessionContractImplementor session) {
 		final boolean useStaticLoader = !session.getLoadQueryInfluencers().hasEnabledFilters()
 				&& !session.getLoadQueryInfluencers().hasEnabledFetchProfiles()

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

@@ -6,6 +6,7 @@
  */
 package org.hibernate.persister.entity;
 
+import org.hibernate.LockOptions;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 
 /**
@@ -16,7 +17,19 @@ public interface UniqueKeyLoadable extends Loadable {
 	 * Load an instance of the persistent class, by a unique key other
 	 * than the primary key.
 	 */
-	Object loadByUniqueKey(String propertyName, Object uniqueKey, SharedSessionContractImplementor session);
+	Object loadByUniqueKey(
+			String propertyName,
+			Object uniqueKey,
+			SharedSessionContractImplementor session);
+
+	/**
+	 * Load an instance of the persistent class, by a natural id.
+	 */
+	Object loadByUniqueKey(
+			String propertyName,
+			Object uniqueKey,
+			LockOptions lockOptions,
+			SharedSessionContractImplementor session);
 
 	/**
 	 * Get the property number of the unique key property

hibernate-core/src/test/java/org/hibernate/test/naturalid/inheritance/cache/InheritedNaturalIdNoCacheTest.java

@@ -6,6 +6,7 @@
  */
 package org.hibernate.test.naturalid.inheritance.cache;
 
+import org.hibernate.WrongClassException;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Configuration;
 
@@ -14,7 +15,7 @@ import org.junit.Test;
 
 import static org.hibernate.testing.transaction.TransactionUtil.doInHibernate;
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
 
 public class InheritedNaturalIdNoCacheTest extends BaseCoreFunctionalTestCase {
 
@@ -57,10 +58,21 @@ public class InheritedNaturalIdNoCacheTest extends BaseCoreFunctionalTestCase {
 		});
 
 		doInHibernate( this::sessionFactory, session -> {
-			ExtendedEntity user = session.byNaturalId( ExtendedEntity.class )
-				.using( "uid", "base" )
-				.load();
-			assertNull( user );
+			try {
+				session.byNaturalId( ExtendedEntity.class )
+						.using( "uid", "base" )
+						.load();
+				fail( "Expecting WrongClassException" );
+			}
+			catch (WrongClassException expected) {
+				// expected outcome
+			}
+			catch (Exception other) {
+				throw new AssertionError(
+						"Unexpected exception type : " + other.getClass().getName(),
+						other
+				);
+			}
 		});
 	}
 }