KAFKA-18692: Consider to unify KStreamImpl "repartitionRequired" with GraphNode "keyChangingOperation"

Author: appchemist

streams/src/main/java/org/apache/kafka/streams/kstream/internals/BranchedKStreamImpl.java

@@ -33,17 +33,15 @@ public class BranchedKStreamImpl<K, V> implements BranchedKStream<K, V> {
     private static final String BRANCH_NAME = "KSTREAM-BRANCH-";
 
     private final KStreamImpl<K, V> source;
-    private final boolean repartitionRequired;
     private final String splitterName;
     private final Map<String, KStream<K, V>> outputBranches = new HashMap<>();
 
     private final List<Predicate<? super K, ? super V>> predicates = new ArrayList<>();
     private final List<String> childNames = new ArrayList<>();
     private final ProcessorGraphNode<K, V> splitterNode;
 
-    BranchedKStreamImpl(final KStreamImpl<K, V> source, final boolean repartitionRequired, final NamedInternal named) {
+    BranchedKStreamImpl(final KStreamImpl<K, V> source, final NamedInternal named) {
         this.source = source;
-        this.repartitionRequired = repartitionRequired;
         this.splitterName = named.orElseGenerateWithPrefix(source.builder, BRANCH_NAME);
 
         // predicates and childNames are passed by reference so when the user adds a branch they get added to
@@ -86,7 +84,7 @@ private void createBranch(final Branched<K, V> branched, final int index) {
         source.builder.addGraphNode(splitterNode, branchChildNode);
         final KStreamImpl<K, V> branch = new KStreamImpl<>(branchChildName, source.keySerde,
                 source.valueSerde, source.subTopologySourceNodes,
-                repartitionRequired, branchChildNode, source.builder);
+                branchChildNode, source.builder);
         process(branch, branchChildName, branchedInternal);
     }
 

streams/src/main/java/org/apache/kafka/streams/kstream/internals/CogroupedStreamAggregateBuilder.java

@@ -210,7 +210,7 @@ private void processRepartitions(final Map<KGroupedStreamImpl<K, ?>, Aggregator<
                                      final String queryableName) {
         for (final KGroupedStreamImpl<K, ?> repartitionReqs : groupPatterns.keySet()) {
 
-            if (repartitionReqs.repartitionRequired) {
+            if (repartitionReqs.repartitionRequired()) {
 
                 final OptimizableRepartitionNodeBuilder<K, ?> repartitionNodeBuilder = optimizableRepartitionNodeBuilder();
 

streams/src/main/java/org/apache/kafka/streams/kstream/internals/GroupedStreamAggregateBuilder.java

@@ -36,7 +36,6 @@ class GroupedStreamAggregateBuilder<K, V> {
     private final InternalStreamsBuilder builder;
     private final Serde<K> keySerde;
     private final Serde<V> valueSerde;
-    private final boolean repartitionRequired;
     private final String userProvidedRepartitionTopicName;
     private final Set<String> subTopologySourceNodes;
     private final String name;
@@ -51,15 +50,13 @@ class GroupedStreamAggregateBuilder<K, V> {
 
     GroupedStreamAggregateBuilder(final InternalStreamsBuilder builder,
                                   final GroupedInternal<K, V> groupedInternal,
-                                  final boolean repartitionRequired,
                                   final Set<String> subTopologySourceNodes,
                                   final String name,
                                   final GraphNode graphNode) {
 
         this.builder = builder;
         this.keySerde = groupedInternal.keySerde();
         this.valueSerde = groupedInternal.valueSerde();
-        this.repartitionRequired = repartitionRequired;
         this.subTopologySourceNodes = subTopologySourceNodes;
         this.name = name;
         this.graphNode = graphNode;
@@ -124,7 +121,7 @@ private <KR, VR> KTable<KR, VR> build(final String aggFunctionName,
         String sourceName = this.name;
         GraphNode parentNode = graphNode;
 
-        if (repartitionRequired) {
+        if (builder.isRepartitionRequired(graphNode)) {
             final OptimizableRepartitionNodeBuilder<K, V> repartitionNodeBuilder = optimizableRepartitionNodeBuilder();
 
             final String repartitionTopicPrefix = userProvidedRepartitionTopicName != null ? userProvidedRepartitionTopicName : storeName;

streams/src/main/java/org/apache/kafka/streams/kstream/internals/InternalStreamsBuilder.java

@@ -100,14 +100,14 @@ public <K, V> KStream<K, V> stream(final Collection<String> topics,
 
         final String name = new NamedInternal(consumed.name()).orElseGenerateWithPrefix(this, KStreamImpl.SOURCE_NAME);
         final StreamSourceNode<K, V> streamSourceNode = new StreamSourceNode<>(name, topics, consumed);
+        streamSourceNode.repartitionByKey();
 
         addGraphNode(root, streamSourceNode);
 
         return new KStreamImpl<>(name,
                                  consumed.keySerde(),
                                  consumed.valueSerde(),
                                  Collections.singleton(name),
-                                 false,
                                  streamSourceNode,
                                  this);
     }
@@ -116,14 +116,14 @@ public <K, V> KStream<K, V> stream(final Pattern topicPattern,
                                        final ConsumedInternal<K, V> consumed) {
         final String name = new NamedInternal(consumed.name()).orElseGenerateWithPrefix(this, KStreamImpl.SOURCE_NAME);
         final StreamSourceNode<K, V> streamPatternSourceNode = new StreamSourceNode<>(name, topicPattern, consumed);
+        streamPatternSourceNode.repartitionByKey();
 
         addGraphNode(root, streamPatternSourceNode);
 
         return new KStreamImpl<>(name,
                                  consumed.keySerde(),
                                  consumed.valueSerde(),
                                  Collections.singleton(name),
-                                 false,
                                  streamPatternSourceNode,
                                  this);
     }
@@ -614,6 +614,14 @@ private GraphNode getKeyChangingParentNode(final GraphNode repartitionNode) {
         return null;
     }
 
+    protected boolean isRepartitionRequired(final GraphNode node) {
+        final GraphNode find = findParentNodeMatching(node, GraphNode::canResolveRepartition);
+        if (find == null) {
+            return false;
+        }
+        return find.isRepartitionRequired();
+    }
+
     private String getFirstRepartitionTopicName(final Collection<OptimizableRepartitionNode<?, ?>> repartitionNodes) {
         return repartitionNodes.iterator().next().repartitionTopic();
     }

streams/src/main/java/org/apache/kafka/streams/kstream/internals/KGroupedStreamImpl.java

@@ -45,28 +45,28 @@ class KGroupedStreamImpl<K, V> extends AbstractStream<K, V> implements KGroupedS
     static final String AGGREGATE_NAME = "KSTREAM-AGGREGATE-";
 
     private final GroupedStreamAggregateBuilder<K, V> aggregateBuilder;
-    final boolean repartitionRequired;
     final String userProvidedRepartitionTopicName;
 
     KGroupedStreamImpl(final String name,
                        final Set<String> subTopologySourceNodes,
                        final GroupedInternal<K, V> groupedInternal,
-                       final boolean repartitionRequired,
                        final GraphNode graphNode,
                        final InternalStreamsBuilder builder) {
         super(name, groupedInternal.keySerde(), groupedInternal.valueSerde(), subTopologySourceNodes, graphNode, builder);
-        this.repartitionRequired = repartitionRequired;
         this.userProvidedRepartitionTopicName = groupedInternal.name();
         this.aggregateBuilder = new GroupedStreamAggregateBuilder<>(
             builder,
             groupedInternal,
-            repartitionRequired,
             subTopologySourceNodes,
             name,
             graphNode
         );
     }
 
+    public boolean repartitionRequired() {
+        return builder.isRepartitionRequired(graphNode);
+    }
+
     @Override
     public KTable<K, V> reduce(final Reducer<V> reducer) {
         return reduce(reducer, Materialized.with(keySerde, valueSerde));