Address review feedback.

Author: ctailor2

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java

@@ -224,7 +224,7 @@ <T> List<T> populateIdsIfNecessary() {
 
 		StagedValues cascadingValues = new StagedValues();
 
-		List<T> roots = new ArrayList<>();
+		List<T> roots = new ArrayList<>(reverseResults.size());
 
 		for (DbActionExecutionResult result : reverseResults) {
 

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateTemplate.java

@@ -16,6 +16,7 @@
 package org.springframework.data.jdbc.core;
 
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -32,7 +33,7 @@
 import org.springframework.data.mapping.callback.EntityCallbacks;
 import org.springframework.data.relational.core.conversion.AggregateChange;
 import org.springframework.data.relational.core.conversion.AggregateChangeWithRoot;
-import org.springframework.data.relational.core.conversion.MergedAggregateChange;
+import org.springframework.data.relational.core.conversion.BatchingAggregateChange;
 import org.springframework.data.relational.core.conversion.MutableAggregateChange;
 import org.springframework.data.relational.core.conversion.RelationalEntityDeleteWriter;
 import org.springframework.data.relational.core.conversion.RelationalEntityInsertWriter;
@@ -47,8 +48,6 @@
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 
-import static java.util.Collections.*;
-
 /**
  * {@link JdbcAggregateOperations} implementation, storing aggregates in and obtaining them from a JDBC data store.
  *
@@ -138,14 +137,15 @@ public <T> T save(T instance) {
 
 		Assert.notNull(instance, "Aggregate instance must not be null!");
 
-		return StreamSupport.stream(saveAll(singletonList(instance)).spliterator(), false).findFirst()
-				.orElseThrow(() -> new IllegalStateException(
-						String.format("Unable to retrieve the result of executing aggregate change for instance: %s", instance)));
+		return performSave(instance, changeCreatorSelectorForSave(instance));
 	}
 
 	@Override
 	public <T> Iterable<T> saveAll(Iterable<T> instances) {
-		return performSaveChange(instances, this::changeCreatorSelectorForSave);
+
+		Assert.isTrue(instances.iterator().hasNext(), "Aggregate instances must not be empty!");
+
+		return performSaveAll(instances);
 	}
 
 	/**
@@ -160,9 +160,7 @@ public <T> T insert(T instance) {
 
 		Assert.notNull(instance, "Aggregate instance must not be null!");
 
-		return performSaveChange(singletonList(instance), this::changeCreatorSelectorForInsert).stream().findFirst()
-				.orElseThrow(() -> new IllegalStateException(
-						String.format("Unable to retrieve the result of executing aggregate change for instance: %s", instance)));
+		return performSave(instance, entity -> createInsertChange(prepareVersionForInsert(entity)));
 	}
 
 	/**
@@ -177,9 +175,7 @@ public <T> T update(T instance) {
 
 		Assert.notNull(instance, "Aggregate instance must not be null!");
 
-		return performSaveChange(singletonList(instance), this::changeCreatorSelectorForUpdate).stream().findFirst()
-				.orElseThrow(() -> new IllegalStateException(
-						String.format("Unable to retrieve the result of executing aggregate change for instance: %s", instance)));
+		return performSave(instance, entity -> createUpdateChange(prepareVersionForUpdate(entity)));
 	}
 
 	@Override
@@ -318,36 +314,50 @@ private <T> void deleteTree(Object id, @Nullable T entity, Class<T> domainType)
 		triggerAfterDelete(entity, id, change);
 	}
 
-	private <T> List<T> performSaveChange(Iterable<T> instances, Function<T, Function<T, AggregateChangeWithRoot<T>>> changeCreatorSelector) {
-
-		ArrayList<T> instancesList = new ArrayList<>();
-		instances.forEach(instancesList::add);
-		Assert.notEmpty(instancesList, "Aggregate instances must not be empty!");
+	private <T> T performSave(T instance, Function<T, AggregateChangeWithRoot<T>> changeCreator) {
 
 		// noinspection unchecked
-		MergedAggregateChange<T, AggregateChangeWithRoot<T>> mergedAggregateChange = instancesList.stream() //
-				.map(instance -> beforeExecute(instance, changeCreatorSelector.apply(instance))) //
-				.reduce(MutableAggregateChange.mergedSave((Class<T>) ClassUtils.getUserClass(instancesList.get(0))),
-						MergedAggregateChange::merge, (left, right) -> right);
+		BatchingAggregateChange<T, AggregateChangeWithRoot<T>> batchingAggregateChange = //
+				BatchingAggregateChange.forSave((Class<T>) ClassUtils.getUserClass(instance));
+		batchingAggregateChange.add(beforeExecute(instance, changeCreator));
 
-		return executor.executeSave(mergedAggregateChange).stream()
-				.map(entityAfterExecution -> afterExecute(mergedAggregateChange, entityAfterExecution))
-				.collect(Collectors.toList());
-	}
+		Iterator<T> afterExecutionIterator = executor.executeSave(batchingAggregateChange).iterator();
 
-	private <T> Function<T, AggregateChangeWithRoot<T>> changeCreatorSelectorForSave(T instance) {
+		Assert.isTrue(afterExecutionIterator.hasNext(), "Instances after execution must not be empty!");
 
-		return context.getRequiredPersistentEntity(instance.getClass()).isNew(instance)
-				? changeCreatorSelectorForInsert(instance)
-				: changeCreatorSelectorForUpdate(instance);
+		return afterExecute(batchingAggregateChange, afterExecutionIterator.next());
 	}
 
-	private <T> Function<T, AggregateChangeWithRoot<T>> changeCreatorSelectorForInsert(T instance) {
-		return entity -> createInsertChange(prepareVersionForInsert(entity));
+	private <T> List<T> performSaveAll(Iterable<T> instances) {
+
+		Iterator<T> iterator = instances.iterator();
+		T firstInstance = iterator.next();
+
+		// noinspection unchecked
+		BatchingAggregateChange<T, AggregateChangeWithRoot<T>> batchingAggregateChange = //
+				BatchingAggregateChange.forSave((Class<T>) ClassUtils.getUserClass(firstInstance));
+		batchingAggregateChange.add(beforeExecute(firstInstance, changeCreatorSelectorForSave(firstInstance)));
+
+		while (iterator.hasNext()) {
+			T instance = iterator.next();
+			batchingAggregateChange.add(beforeExecute(instance, changeCreatorSelectorForSave(instance)));
+		}
+
+		List<T> instancesAfterExecution = executor.executeSave(batchingAggregateChange);
+
+		ArrayList<T> results = new ArrayList<>(instancesAfterExecution.size());
+		for (T instance : instancesAfterExecution) {
+			results.add(afterExecute(batchingAggregateChange, instance));
+		}
+
+		return results;
 	}
 
-	private <T> Function<T, AggregateChangeWithRoot<T>> changeCreatorSelectorForUpdate(T instance) {
-		return entity -> createUpdateChange(prepareVersionForUpdate(entity));
+	private <T> Function<T, AggregateChangeWithRoot<T>> changeCreatorSelectorForSave(T instance) {
+
+		return context.getRequiredPersistentEntity(instance.getClass()).isNew(instance)
+				? entity -> createInsertChange(prepareVersionForInsert(entity))
+				: entity -> createUpdateChange(prepareVersionForUpdate(entity));
 	}
 
 	private <T> AggregateChangeWithRoot<T> createInsertChange(T instance) {

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIntegrationTests.java

@@ -601,7 +601,7 @@ void manyInsertsWithNestedEntities() {
 
 		List<Root> savedRoots = rootRepository.saveAll(asList(root1, root2));
 
-		List<Root> reloadedRoots = rootRepository.findAll();
+		List<Root> reloadedRoots = rootRepository.findAllByOrderByIdAsc();
 		assertThat(reloadedRoots).isEqualTo(savedRoots);
 		assertThat(reloadedRoots).hasSize(2);
 		assertIsEqualToWithNonNullIds(reloadedRoots.get(0), root1);
@@ -628,7 +628,7 @@ void manyUpdatesWithNestedEntities() {
 
 		List<Root> updatedRoots = rootRepository.saveAll(asList(updatedRoot1, updatedRoot2));
 
-		List<Root> reloadedRoots = rootRepository.findAll();
+		List<Root> reloadedRoots = rootRepository.findAllByOrderByIdAsc();
 		assertThat(reloadedRoots).isEqualTo(updatedRoots);
 		assertThat(reloadedRoots).containsExactly(updatedRoot1, updatedRoot2);
 	}
@@ -645,7 +645,7 @@ void manyInsertsAndUpdatesWithNestedEntities() {
 		Root root2 = createRoot("root2");
 		List<Root> savedRoots = rootRepository.saveAll(asList(updatedRoot1, root2));
 
-		List<Root> reloadedRoots = rootRepository.findAll();
+		List<Root> reloadedRoots = rootRepository.findAllByOrderByIdAsc();
 		assertThat(reloadedRoots).isEqualTo(savedRoots);
 		assertThat(reloadedRoots.get(0)).isEqualTo(updatedRoot1);
 		assertIsEqualToWithNonNullIds(reloadedRoots.get(1), root2);
@@ -795,7 +795,9 @@ MyEventListener eventListener() {
 		}
 	}
 
-	interface RootRepository extends ListCrudRepository<Root, Long> {}
+	interface RootRepository extends ListCrudRepository<Root, Long> {
+		List<Root> findAllByOrderByIdAsc();
+	}
 
 	@Value
 	static class Root {

spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/MergedAggregateChange.java renamed to spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/BatchingAggregateChange.java

@@ -15,20 +15,36 @@
  */
 package org.springframework.data.relational.core.conversion;
 
+import org.springframework.util.Assert;
+
 /**
  * Represents the changes happening to one or more aggregates (as used in the context of Domain Driven Design) as a
- * whole. This change allows additional {@link MutableAggregateChange} of a particular kind to be merged into it to
- * broadly represent the changes to multiple aggregates across all such merged changes.
+ * whole. This change allows additional {@link MutableAggregateChange} of a particular kind to be added to it to
+ * broadly represent the changes to multiple aggregates across all such added changes.
  *
  * @author Chirag Tailor
  * @since 3.0
  */
-public interface MergedAggregateChange<T, C extends MutableAggregateChange<T>> extends AggregateChange<T> {
+public interface BatchingAggregateChange<T, C extends MutableAggregateChange<T>> extends AggregateChange<T> {
 	/**
-	 * Merges a {@code MutableAggregateChange} into this {@code MergedAggregateChange}.
+	 * Adds a {@code MutableAggregateChange} into this {@code BatchingAggregateChange}.
 	 *
 	 * @param aggregateChange must not be {@literal null}.
-	 * @return the change resulting from the merge. Guaranteed to be not {@literal null}.
 	 */
-	MergedAggregateChange<T, C> merge(C aggregateChange);
+	void add(C aggregateChange);
+
+	/**
+	 * Factory method to create a {@link BatchingAggregateChange} for saving entities.
+	 *
+	 * @param entityClass aggregate root type.
+	 * @param <T> entity type.
+	 * @return the {@link BatchingAggregateChange} for saving root entities.
+	 * @since 3.0
+	 */
+	static <T> BatchingAggregateChange<T, AggregateChangeWithRoot<T>> forSave(Class<T> entityClass) {
+
+		Assert.notNull(entityClass, "Entity class must not be null");
+
+		return new SaveBatchingAggregateChange<>(entityClass);
+	}
 }

spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/MutableAggregateChange.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.relational.core.conversion;
 
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 
@@ -34,7 +35,7 @@ public interface MutableAggregateChange<T> extends AggregateChange<T> {
 	 * @param entity aggregate root to save.
 	 * @param <T> entity type.
 	 * @return the {@link AggregateChangeWithRoot} for saving the root {@code entity}.
-	 * @since 2.4
+	 * @since 1.2
 	 */
 	@SuppressWarnings("unchecked")
 	static <T> AggregateChangeWithRoot<T> forSave(T entity) {
@@ -44,6 +45,24 @@ static <T> AggregateChangeWithRoot<T> forSave(T entity) {
 		return new DefaultAggregateChangeWithRoot<>(Kind.SAVE, (Class<T>) ClassUtils.getUserClass(entity));
 	}
 
+	/**
+	 * Factory method to create an {@link AggregateChangeWithRoot} for saving entities.
+	 *
+	 * @param entity aggregate root to save.
+	 * @param previousVersion the previous version assigned to the instance being saved. May be {@literal null}.
+	 * @param <T> entity type.
+	 * @return the {@link AggregateChangeWithRoot} for saving the root {@code entity}.
+	 * @since 2.4
+	 * @deprecated since 3.0, use {@link #forSave(Object)} instead.
+	 */
+	@Deprecated
+	static <T> AggregateChangeWithRoot<T> forSave(T entity, @Nullable Number previousVersion) {
+
+		Assert.notNull(entity, "Entity must not be null");
+
+		return forSave(entity);
+	}
+
 	/**
 	 * Factory method to create an {@link MutableAggregateChange} for deleting entities.
 	 *
@@ -61,33 +80,36 @@ static <T> MutableAggregateChange<T> forDelete(T entity) {
 	}
 
 	/**
-	 * Factory method to create a {@link MergedAggregateChange} for saving entities.
+	 * Factory method to create an {@link MutableAggregateChange} for deleting entities.
 	 *
 	 * @param entityClass aggregate root type.
 	 * @param <T> entity type.
-	 * @return the {@link MergedAggregateChange} for saving root entities.
-	 * @since 3.0
+	 * @return the {@link MutableAggregateChange} for deleting the root {@code entity}.
+	 * @since 1.2
 	 */
-	static <T> MergedAggregateChange<T, AggregateChangeWithRoot<T>> mergedSave(Class<T> entityClass) {
+	static <T> MutableAggregateChange<T> forDelete(Class<T> entityClass) {
 
 		Assert.notNull(entityClass, "Entity class must not be null");
 
-		return new SaveMergedAggregateChange<>(entityClass);
+		return new DefaultAggregateChange<>(Kind.DELETE, entityClass);
 	}
 
 	/**
 	 * Factory method to create an {@link MutableAggregateChange} for deleting entities.
 	 *
 	 * @param entityClass aggregate root type.
+	 * @param previousVersion the previous version assigned to the instance being saved. May be {@literal null}.
 	 * @param <T> entity type.
 	 * @return the {@link MutableAggregateChange} for deleting the root {@code entity}.
 	 * @since 2.4
+	 * @deprecated since 3.0, use {@link #forDelete(Class)} instead.
 	 */
-	static <T> MutableAggregateChange<T> forDelete(Class<T> entityClass) {
+	@Deprecated
+	static <T> MutableAggregateChange<T> forDelete(Class<T> entityClass, @Nullable Number previousVersion) {
 
 		Assert.notNull(entityClass, "Entity class must not be null");
 
-		return new DefaultAggregateChange<>(Kind.DELETE, entityClass);
+		return forDelete(entityClass);
 	}
 
 	/**

spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/SaveMergedAggregateChange.java renamed to spring-data-relational/src/main/java/org/springframework/data/relational/core/conversion/SaveBatchingAggregateChange.java

@@ -19,7 +19,6 @@
 
 import java.util.ArrayList;
 import java.util.Comparator;
-import java.util.EnumMap;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -30,27 +29,27 @@
 import org.springframework.util.Assert;
 
 /**
- * A {@link org.springframework.data.relational.core.conversion.MergedAggregateChange} implementation for save changes
+ * A {@link BatchingAggregateChange} implementation for save changes
  * that can contain actions for any mix of insert and update operations. When consumed, actions are yielded in the
  * appropriate entity tree order with inserts carried out from root to leaves and deletes in reverse. All insert
  * operations are grouped into batches to offer the ability for an optimized batch operation to be used.
  *
  * @author Chirag Tailor
  * @since 3.0
  */
-public class SaveMergedAggregateChange<T> implements MergedAggregateChange<T, AggregateChangeWithRoot<T>> {
+public class SaveBatchingAggregateChange<T> implements BatchingAggregateChange<T, AggregateChangeWithRoot<T>> {
 
 	private static final Comparator<PersistentPropertyPath<RelationalPersistentProperty>> pathLengthComparator = //
 			Comparator.comparing(PersistentPropertyPath::getLength);
 
 	private final Class<T> entityType;
 	private final List<DbAction.WithRoot<?>> rootActions = new ArrayList<>();
-	private final Map<PersistentPropertyPath<RelationalPersistentProperty>, EnumMap<IdValueSource, List<DbAction.Insert<Object>>>> insertActions = //
+	private final Map<PersistentPropertyPath<RelationalPersistentProperty>, Map<IdValueSource, List<DbAction.Insert<Object>>>> insertActions = //
 			new HashMap<>();
 	private final Map<PersistentPropertyPath<RelationalPersistentProperty>, List<DbAction.Delete<?>>> deleteActions = //
 			new HashMap<>();
 
-	public SaveMergedAggregateChange(Class<T> entityType) {
+	public SaveBatchingAggregateChange(Class<T> entityType) {
 		this.entityType = entityType;
 	}
 
@@ -78,7 +77,7 @@ public void forEachAction(Consumer<? super DbAction<?>> consumer) {
 	}
 
 	@Override
-	public MergedAggregateChange<T, AggregateChangeWithRoot<T>> merge(AggregateChangeWithRoot<T> aggregateChange) {
+	public void add(AggregateChangeWithRoot<T> aggregateChange) {
 
 		aggregateChange.forEachAction(action -> {
 			if (action instanceof DbAction.WithRoot<?> rootAction) {
@@ -90,21 +89,20 @@ public MergedAggregateChange<T, AggregateChangeWithRoot<T>> merge(AggregateChang
 				addDelete(deleteAction);
 			}
 		});
-		return this;
 	}
 
 	private void addInsert(DbAction.Insert<Object> action) {
 
 		PersistentPropertyPath<RelationalPersistentProperty> propertyPath = action.getPropertyPath();
 		insertActions.merge(propertyPath,
-				new EnumMap<>(singletonMap(action.getIdValueSource(), new ArrayList<>(singletonList(action)))),
-				(enumMap, enumMapDefaultValue) -> {
-					enumMap.merge(action.getIdValueSource(), new ArrayList<>(singletonList(action)),
+				new HashMap<>(singletonMap(action.getIdValueSource(), new ArrayList<>(singletonList(action)))),
+				(map, mapDefaultValue) -> {
+					map.merge(action.getIdValueSource(), new ArrayList<>(singletonList(action)),
 							(actions, listDefaultValue) -> {
 								actions.add(action);
 								return actions;
 							});
-					return enumMap;
+					return map;
 				});
 	}