perf(entity-manager): preload a page's single references in one query

src/EntityManager.php

@@ -5,6 +5,8 @@
 readonly class EntityManager
 {
 
+	private const PRELOAD_CHUNK_SIZE = 1000;
+
 	public function __construct(
 		private \Spameri\Elastic\Model\Insert $insert,
 		private \Spameri\Elastic\Model\GetAllBy $getAllBy,
@@ -124,6 +126,8 @@ protected function createCollection(
 		\Spameri\ElasticQuery\Response\ResultSearch $resultSearch,
 	): \Spameri\Elastic\Entity\Collection\ElasticEntityCollection
 	{
+		$this->preloadReferences($class, $resultSearch);
+
 		$entities = [];
 		foreach ($resultSearch->hits() as $hit) {
 			try {
@@ -315,6 +319,89 @@ public function preloadRelations(
 	}
 
 
+	/**
+	 * Fetches every document the page's hits reference as a single entity in one query per
+	 * class, so hydrating them hits the identity map instead of running one find() per reference.
+	 * The preloaded pages preload their own references the same way, so a chain of references
+	 * costs one query per level. Ids already being preloaded further up are skipped, which ends
+	 * reference cycles.
+	 *
+	 * @param class-string $class
+	 */
+	private function preloadReferences(
+		string $class,
+		\Spameri\ElasticQuery\Response\ResultSearch $resultSearch,
+	): void
+	{
+		/** @var array<class-string, array<string, true>> $idsByClass */
+		$idsByClass = [];
+
+		foreach ($resultSearch->hits() as $hit) {
+			$hitClass = $hit->getValue(\Spameri\Elastic\Model\Insert\PrepareEntityArray::ENTITY_CLASS) ?? $class;
+			$idsByClass = \array_replace_recursive(
+				$idsByClass,
+				$this->entityFactory->collectReferences($hit, $hitClass),
+			);
+		}
+
+		foreach ($idsByClass as $referencedClass => $ids) {
+			$idsToLoad = [];
+
+			foreach (\array_keys($ids) as $id) {
+				$id = (string) $id;
+
+				if (
+					$this->identityMap->get($referencedClass, $id) === null
+					&& isset($this->identityMap->preloadingList[$referencedClass][$id]) === false
+					&& isset($this->identityMap->creatingEntityList[$referencedClass][$id]) === false
+				) {
+					$idsToLoad[] = $id;
+				}
+			}
+
+			foreach (\array_chunk($idsToLoad, self::PRELOAD_CHUNK_SIZE) as $chunk) {
+				$this->preloadChunk($referencedClass, $chunk);
+			}
+		}
+	}
+
+
+	/**
+	 * A failed preload is not fatal: whatever it did not load is still fetched one by one with
+	 * find() while hydrating, exactly as without preloading.
+	 *
+	 * @param class-string $class
+	 * @param list<string> $ids
+	 */
+	private function preloadChunk(
+		string $class,
+		array $ids,
+	): void
+	{
+		foreach ($ids as $id) {
+			$this->identityMap->preloadingList[$class][$id] = true;
+		}
+
+		$elasticQuery = new \Spameri\ElasticQuery\ElasticQuery();
+		$elasticQuery->addMustQuery(
+			new \Spameri\ElasticQuery\Query\Terms('_id', $ids),
+		);
+		$elasticQuery->options()->changeSize(\count($ids));
+
+		try {
+			$this->findBy($elasticQuery, $class);
+
+		} catch (\Spameri\Elastic\Exception\ElasticSearch) {
+			// findBy() has logged it; hydration falls back to find().
+
+		} finally {
+			foreach ($ids as $id) {
+				unset($this->identityMap->preloadingList[$class][$id]);
+			}
+		}
+	}
+
+
 	public function clear(): void
 	{
 		$this->identityMap->clear();

src/Factory/EntityFactory.php

@@ -68,6 +68,101 @@ class: $class,
 		return $entity;
 	}
 
+	/**
+	 * Lists the documents a hit references as single entities, walking its properties the way
+	 * resolveProperties() does, so a page can fetch them in one query before it is hydrated.
+	 * Lazy collections are left out; their references load when they are first read.
+	 *
+	 * @return array<class-string, array<string, true>> referenced ids keyed by entity class
+	 * @throws \ReflectionException
+	 */
+	public function collectReferences(
+		\Spameri\ElasticQuery\Response\Result\Hit $hit,
+		string $class,
+		string|null $parentFieldName = null,
+	): array
+	{
+		$references = [];
+		$reflection = $this->reflection->createReflection($class);
+
+		foreach ($this->reflection->getProperties($reflection) as $property) {
+			$hitKey = $parentFieldName === null ? $property->getName() : $parentFieldName . '.' . $property->getName();
+			$reflectionPropertyType = $this->reflection->getPropertyType($property);
+			$value = $hit->getValue($hitKey);
+
+			if ($reflectionPropertyType === null || $value === null) {
+				continue;
+			}
+
+			$propertyTypeName = $reflectionPropertyType->getName();
+
+			if (
+				$propertyTypeName === \Spameri\Elastic\Entity\Property\Date::class
+				|| $propertyTypeName === \Spameri\Elastic\Entity\Property\DateTime::class
+			) {
+				continue;
+			}
+
+			if (\count($property->getAttributes()) > 0) {
+				foreach ($property->getAttributes() as $attribute) {
+					$nested = [];
+
+					if ($attribute->getName() === \Spameri\Elastic\Mapping\Entity::class) {
+						/** @var array{class: class-string} $arguments */
+						$arguments = $attribute->getArguments();
+
+						if ($arguments['class'] !== \Spameri\Elastic\Entity\Property\ElasticId::class) {
+							$nested = $this->collectReferences($hit, $propertyTypeName, $hitKey);
+						}
+
+					} elseif ($attribute->getName() === \Spameri\Elastic\Mapping\STIEntity::class) {
+						$nested = $this->collectReferences($hit, $propertyTypeName, $hitKey);
+
+					} elseif (
+						$attribute->getName() === \Spameri\Elastic\Mapping\STIElasticEntity::class
+						&& \is_array($value)
+						&& isset($value[\Spameri\Elastic\Model\Insert\PrepareEntityArray::ENTITY_CLASS], $value[\Spameri\Elastic\Model\Insert\PrepareEntityArray::ENTITY_ID])
+						&& \class_exists($value[\Spameri\Elastic\Model\Insert\PrepareEntityArray::ENTITY_CLASS])
+					) {
+						$references[$value[\Spameri\Elastic\Model\Insert\PrepareEntityArray::ENTITY_CLASS]][(string) $value[\Spameri\Elastic\Model\Insert\PrepareEntityArray::ENTITY_ID]] = true;
+					}
+
+					$references = \array_replace_recursive($references, $nested);
+				}
+
+				continue;
+			}
+
+			if (\class_exists($propertyTypeName) === false || \enum_exists($propertyTypeName) === true) {
+				continue;
+			}
+
+			$interfaces = \class_implements($propertyTypeName);
+
+			if (isset($interfaces[\Spameri\Elastic\Entity\ValueInterface::class]) === true) {
+				continue;
+			}
+
+			if (isset($interfaces[\Spameri\Elastic\Entity\ElasticEntityInterface::class]) === true) {
+				if (\is_string($value) === true && $value !== '') {
+					$references[$propertyTypeName][$value] = true;
+				}
+
+				continue;
+			}
+
+			if ($this->container->getByType($propertyTypeName, false) === null) {
+				$references = \array_replace_recursive(
+					$references,
+					$this->collectReferences($hit, $propertyTypeName, $hitKey),
+				);
+			}
+		}
+
+		return $references;
+	}
+
+
 	/**
 	 * @return array<mixed>
 	 * @throws \ReflectionException

src/Model/IdentityMap.php

@@ -27,6 +27,13 @@ class IdentityMap
 	 */
 	public array $uninitializedEntityList = [];
 
+	/**
+	 * Referenced documents a page is already fetching in one batch, so nested pages don't fetch them again.
+	 *
+	 * @var array<class-string, array<string, true>>
+	 */
+	public array $preloadingList = [];
+
 
 	public function add(
 		\Spameri\Elastic\Entity\AbstractElasticEntity $entity,
@@ -128,6 +135,7 @@ public function clear(): void
 		$this->persisted = [];
 		$this->creatingEntityList = [];
 		$this->uninitializedEntityList = [];
+		$this->preloadingList = [];
 	}
 
 }