diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
index e5ad742ad..01cdbff6a 100644
--- a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
+++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
@@ -16,9 +16,6 @@
 
 package org.springframework.cloud.stream.function;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -48,7 +45,6 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 	@SuppressWarnings("rawtypes")
 	private final Function function;
 
-	@SuppressWarnings("rawtypes")
 	private final Function<Object, Object> outputMessageEnricher;
 
 	PartitionAwareFunctionWrapper(Function<?, ?> function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
@@ -59,24 +55,13 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 			PartitionHandler partitionHandler = new PartitionHandler(evaluationContext, producerProperties, context.getBeanFactory());
 
 			this.outputMessageEnricher = output -> {
-				if (ObjectUtils.isArray(output) && !(output instanceof byte[])) {
-					output = Arrays.asList(output);
-				}
-				if (output instanceof Iterable) {
-					Iterable elements = (Iterable) output;
-					List<Message> messages = new ArrayList<>();
-					for (Object element : elements) {
-						if (!(element instanceof Message)) {
-							element = MessageBuilder.withPayload(element).build();
-						}
-						messages.add(toMessageWithPartitionHeader((Message) element, partitionHandler));
-					}
-					return messages;
+				if ((ObjectUtils.isArray(output) && !(output instanceof byte[])) || output instanceof Iterable) {
+					return output;
 				}
 				else if (!(output instanceof Message)) {
 					output = MessageBuilder.withPayload(output).build();
 				}
-				return toMessageWithPartitionHeader((Message) output, partitionHandler);
+				return toMessageWithPartitionHeader((Message<?>) output, partitionHandler);
 			};
 		}
 		else {
diff --git a/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/function/ImplicitFunctionBindingTests.java b/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/function/ImplicitFunctionBindingTests.java
index 75ac02af0..74fda4ce5 100644
--- a/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/function/ImplicitFunctionBindingTests.java
+++ b/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/function/ImplicitFunctionBindingTests.java
@@ -22,6 +22,7 @@ import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
@@ -763,6 +764,29 @@ public class ImplicitFunctionBindingTests {
 		}
 	}
 
+	@Test
+	public void partitionOnOutputPayloadAsListTest() {
+		System.clearProperty("spring.cloud.function.definition");
+		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(TestChannelBinderConfiguration
+				.getCompleteConfiguration(PojoFunctionConfiguration.class))
+						.web(WebApplicationType.NONE).run("--spring.cloud.function.definition=persons",
+								"--spring.cloud.stream.bindings.persons-out-0.producer.partitionKeyExpression=payload.id",
+								"--spring.cloud.stream.bindings.persons-out-0.producer.partitionCount=5",
+								"--spring.jmx.enabled=false")) {
+
+			InputDestination inputDestination = context.getBean(InputDestination.class);
+			OutputDestination outputDestination = context.getBean(OutputDestination.class);
+
+			Message<byte[]> inputMessage = MessageBuilder.withPayload("Jim Lahey".getBytes()).build();
+
+			inputDestination.send(inputMessage, "persons-in-0");
+
+			assertThat(outputDestination.receive(100, "persons-out-0").getHeaders().get("scst_partition")).isEqualTo(3);
+
+			assertThat(outputDestination.receive(100)).isNull();
+		}
+	}
+
 	@Test
 	public void partitionOnOutputPayloadTestReactive() {
 		System.clearProperty("spring.cloud.function.definition");
@@ -1490,6 +1514,16 @@ public class ImplicitFunctionBindingTests {
 			};
 		}
 
+		@Bean
+		public Function<String, List<Message<Person>>> persons() {
+			return x -> {
+				Person person = new Person();
+				person.setName(x);
+				person.setId(3);
+				return Collections.singletonList(MessageBuilder.withPayload(person).build());
+			};
+		}
+
 		@Bean
 		public Function<Flux<String>, Flux<Person>> funcReactive() {
 			return flux -> flux.map(value -> {