[GitHub] [iceberg] yegangy0718 opened a new pull request, #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 opened a new pull request, #6382:
URL: https://github.com/apache/iceberg/pull/6382

   This PR is created as part of issue https://github.com/apache/iceberg/issues/6303 and project https://github.com/apache/iceberg/projects/27
   
   In this PR, we focus on 
   bin packing based on traffic distribution statistics. This works well for skewed data on partition columns (like event time). This requires calculating traffic distribution statistics across partition columns and use the statistics to guide shuffling decision.
   
   Changes:
   1. Implement ShuffleOperator which will be added before Iceberg Writer operator to collect data distribution based on key(generated from provided KeySelector)
   2. Implement ShuffleRecordWrapper which contains either the record for data distribution information
   
   I will have following up PRs to implement ShuffleCoordinator, the data distribution sending and receiving logic between coordinator and operator, and etc. 
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1043246601


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+
+    assertNotNull(
+        stateContext
+            .getOperatorStateStore()
+            .getListState(operator.generateGlobalDataDistributionWeightDescriptor()));
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("a"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("b"));
+    assertEquals(2L, (long) operator.localDataStatisticsMap().get("a"));
+    assertEquals(1L, (long) operator.localDataStatisticsMap().get("b"));
+  }
+
+  // ---------------- helper methods -------------------------
+
+  private StateInitializationContext getStateContext() throws Exception {
+    // Create the state context.
+    OperatorStateStore operatorStateStore = createOperatorStateStore();
+    return new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
+  }
+
+  private OperatorStateStore createOperatorStateStore() throws Exception {
+    MockEnvironment env = new MockEnvironmentBuilder().build();
+    AbstractStateBackend abstractStateBackend = new HashMapStateBackend();
+    CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
+    return abstractStateBackend.createOperatorStateBackend(

Review Comment:
   Maybe put some data here, and check that the state is restored as expected?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1048051789


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+
+    assertNotNull(
+        stateContext
+            .getOperatorStateStore()
+            .getListState(operator.generateGlobalDataDistributionWeightDescriptor()));
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("a"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("b"));
+    assertEquals(2L, (long) operator.localDataStatisticsMap().get("a"));
+    assertEquals(1L, (long) operator.localDataStatisticsMap().get("b"));
+  }
+
+  // ---------------- helper methods -------------------------
+
+  private StateInitializationContext getStateContext() throws Exception {
+    // Create the state context.
+    OperatorStateStore operatorStateStore = createOperatorStateStore();
+    return new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
+  }
+
+  private OperatorStateStore createOperatorStateStore() throws Exception {
+    MockEnvironment env = new MockEnvironmentBuilder().build();
+    AbstractStateBackend abstractStateBackend = new HashMapStateBackend();
+    CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
+    return abstractStateBackend.createOperatorStateBackend(

Review Comment:
   We get globalDataStatistics from coordinator and then save the value during snapshot.  The function `handleOperatorEvent` to receive globalDataStatistics from coordinator is not implemented in this PR yet. Thus we won't be able to test that part. 



##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+
+    assertNotNull(
+        stateContext
+            .getOperatorStateStore()
+            .getListState(operator.generateGlobalDataDistributionWeightDescriptor()));
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("a"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("b"));
+    assertEquals(2L, (long) operator.localDataStatisticsMap().get("a"));
+    assertEquals(1L, (long) operator.localDataStatisticsMap().get("b"));
+  }
+
+  // ---------------- helper methods -------------------------
+
+  private StateInitializationContext getStateContext() throws Exception {
+    // Create the state context.
+    OperatorStateStore operatorStateStore = createOperatorStateStore();
+    return new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
+  }
+
+  private OperatorStateStore createOperatorStateStore() throws Exception {
+    MockEnvironment env = new MockEnvironmentBuilder().build();
+    AbstractStateBackend abstractStateBackend = new HashMapStateBackend();
+    CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
+    return abstractStateBackend.createOperatorStateBackend(

Review Comment:
   I will add a test using OneInputStreamOperatorTestHarness in latest commit. After implementing the globalDataStatistics setting logic in next PR, I will add the test to take snapshot and restore from new state to verify state is restored as expected. 



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   My current understanding is, when using broadcast, there will be two streams. One stream will be send to all tasks, and the other stream won’t be send to all tasks. Then connecting them together by calling `connect` . The task can use `getBroadcastState` to get the broadcast stream.
   From concept perspective, yes, it’s similar to what we are dong. Global data statistics will be the broadcast one.
   But for our case, there is no stream/place publish the broadcast state, where the shuffle operator can get BroadcastState? Do you mean the shuffle operator publishes the data statistics once receiving from coordinator and then in function `initializeState` shuffle operator can get BroadcastState?



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   Will update the name to localStatistics. 
   For the local and global DataStatistics definition, thanks for the input. Yes, to support high cardinality, data statistics cannot be collected by count simply. 
   I have concern about defining SimpleCountStatistics and other CountStatistics. Since ShuffleOperator itself is an internal class. If we have different impl for CountStatistics, then  client may need to provide config. But sometimes client may not know stream data distribution since it can change. 
   For now I would prefer to keep the current format and refactor it later when we finalize the high cardinality design. 
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1043214544


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   The global data distribution should be the same (or samish) for every instance of the operator.
   Shall we use `getBroadcastState` instead?
   Based on my understanding that would better approximate our requirements for redistributing the state when rescaling. Even if there are some small discrepancies in the values between the different operation instances, we should not be concerned too much about them.
   
   WDYT?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064083656


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight

Review Comment:
   we should explain why this wrapper class it needed. This is the only way for the shuffle operator to pass the statistics to the custom partitioner.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064080676


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {

Review Comment:
   not sure if this method is necessary for code reuse. not sure why it is necessary to check it in testing.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#issuecomment-1453302730

   Left some comments, and I will reflect on them in the next few days. Thank you.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044077784


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,

Review Comment:
   TypeInformation contains serializer. Maybe we don't need to limit `K` to `Serializable?
   
   ```
       public abstract TypeSerializer<T> createSerializer(ExecutionConfig config);
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044080608


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, Map<K, Long> globalDataDistributionWeight) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataDistributionWeight != null,
+        "A ShuffleRecordWrapper has to contain exactly one of record and stats, not neither or both");

Review Comment:
   nit: `contain exactly one of record and stats` -> `contain either record or stats`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1045449754


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,

Review Comment:
   I see what you're saying. Even though `K` is not serializable itself, user can still use `TypeInformation` to define `TypeSerializer`. 
   I will remove the `Serializable` limitation for K in latest commit. 



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   OK. Let me take some time to investigate if `BroadcastState` can be applied to this case. 



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, Map<K, Long> globalDataDistributionWeight) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataDistributionWeight != null,
+        "A ShuffleRecordWrapper has to contain exactly one of record and stats, not neither or both");

Review Comment:
   will update in next version. 



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, Map<K, Long> globalDataDistributionWeight) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataDistributionWeight != null,
+        "A ShuffleRecordWrapper has to contain exactly one of record and stats, not neither or both");
+    this.globalDataDistributionWeight = globalDataDistributionWeight;
+    this.record = record;
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromDistribution(
+      Map<K, Long> globalDataDistributionWeight) {
+    return new ShuffleRecordWrapper<>(null, globalDataDistributionWeight);
+  }
+
+  boolean hasGlobalDataDistributionWeight() {
+    return globalDataDistributionWeight != null;
+  }
+
+  boolean hasRecord() {
+    return record != null;
+  }
+
+  Map<K, Long> globalDataDistributionWeight() {
+    return globalDataDistributionWeight;
+  }
+
+  T record() {
+    return record;
+  }
+
+  @Override
+  public String toString() {
+    return String.format(

Review Comment:
   Oh, yeah, other Iceberg class use guava. to generate String. I will update in next version



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, Map<K, Long> globalDataDistributionWeight) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataDistributionWeight != null,
+        "A ShuffleRecordWrapper has to contain exactly one of record and stats, not neither or both");
+    this.globalDataDistributionWeight = globalDataDistributionWeight;
+    this.record = record;
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromDistribution(

Review Comment:
   will update in next version.



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {

Review Comment:
   if we don't implement `OneInputStreamOperator` and only implement `Input<IN>`, then we need to implement function `setKeyContextElement`, but `OneInputStreamOperator` provides the default implementation.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1054537560


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   I think the data statistics interface will be very straightforward.
   I would expect something like this:
   ```
   public interface DataStatistic<K> {
       long size();
       void put(K key);
       void merge(DataStatistic<K> other);
   }
   ```
   
   I do not think that the ShuffleOperator needs to know anything about the statistics and the partitioners used.
   My feeling is that the ShuffleOperator is just responsible to:
   - Collect statistics
   - Distribute the statistics to the Partitioners
   
   The `Partitioners` are responsible to decide on the distribution based on the collected `DataStatistics`. Here, I am not sure we can separate the `Partitioner` from the actual strategy (Bin Pack, or Range) used, since fetching a distribution based on a key (Iceberg partition based - Map), or fetching a distribution based on a range list (Iceberg sorted tables - TreeSet) has different performance on the hot code path, so we could not fall back to the more general solution (TreeSet) for performance reasons.
   
   I still feel that having an interface defined for `DataStatistics` would be a good direction even at this stage.
   
   We can implement the `MapStatistics` easily like this:
   ```
   public class MapStatistic<K> implements DataStatistic<K> {
       private Map<K, Long> data = Maps.newHashMap();
   
       @Override
       public long size() {
           return data.size();
       }
   
       @Override
       public void put(K key) {
           data.put(key, data.getOrDefault(key, 0L) + 1);
       }
   
       @Override
       public void merge(DataStatistic<K> other) {
           Preconditions.checkArgument(other instanceof MapStatistic, "Can not merge this type of statistics: " + other);
           MapStatistic mapStatistic = (MapStatistic<K>) other;
           Map<K, Long> otherData = mapStatistic.data;
           otherData.forEach((key, count) -> data.put(key, data.getOrDefault(key, 0L) + count));
       }
   }
   ```
   
   WTYT?
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#issuecomment-1439558920

   Hi @stevenzwu  if we decide to keep shuffling implementation in the Iceberg repo, could you help to take another look at the PR? Thanks!


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1136558118


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private DataStatistics<K> localStatistics;
+  private DataStatistics<K> globalStatistics;
+  private ListState<DataStatistics<K>> globalStatisticsState;
+
+  public DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(

Review Comment:
   Hi @stevenzwu  as we discussed offline, I update the state from ListState to  UnionListState since ListState uses round robin pattern during rescaling and UnionlistState uses broadcast pattern which can satisfy the DataStatisticsOperator need better(everyone should have the same global data statistics) 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124679471


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsAndRecordWrapper<T, K>>

Review Comment:
   I will be fine either way. this streaming operator is mainly responsible for calculating and reporting statistics periodically. with `finish()`, it is probably not important to report the last statistics.
   
   Note that for real batch jobs, the shuffling will be different. Statistics are probably sampled and calculated based on all the input data. it is a different shuffling compared to stream shuffling.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1046704878


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;

Review Comment:
   We plan to create a StreamPartitioner which is chained with the shuffle operator to takes the global distribution via ShuffleRecordWrapper, and then generate the channel assignment at there



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1050285783


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   SGTM. I can create a followup PR to add scaling up test case to check the behavior for BroadcastState.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064083596


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getListState(generateGlobalDataDistributionWeightDescriptor());
+
+    if (globalDataStatisticsState.get() != null
+        && globalDataStatisticsState.get().iterator().hasNext()) {
+      globalDataStatistics = globalDataStatisticsState.get().iterator().next();
+    } else {
+      globalDataStatistics = dataStatisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    // TODO: handle scaling up
+    if (globalDataStatistics != null && globalDataStatistics.size() > 0) {
+      output.collect(new StreamRecord<>(ShuffleRecordWrapper.fromStatistics(globalDataStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with globalDataDistributionWeight from coordinator and update
+    // globalDataStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localDataStatistics.put(key);
+    output.collect(new StreamRecord<>(ShuffleRecordWrapper.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.debug("Taking shuffle operator snapshot for checkpoint {}", checkpointId);
+
+    // Set the globalDataDistributionWeightState state to the latest global value.
+    if (globalDataStatistics != null && globalDataStatistics.size() > 0) {
+      globalDataStatisticsState.clear();
+      globalDataStatisticsState.add(globalDataStatistics);
+    }
+
+    // TODO: send to coordinator
+    // For now we make it simple to send globalDataStatisticsState at checkpoint
+
+    // Reset the local data count
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();

Review Comment:
   It might make sense to add a `reset` or `clear` method to `DataStatistics` interface. what do you think? @yegangy0718 and @pvary 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1055124544


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   I like @pvary's proposal on the `DataStatistics` interface.
   
   > As for Range partitioning strategy, I would expect that we need to use some kind of Digest statistics.
   
   I am not sure about this. I can see range partition for non-partition columns with cardinality at hundreds or thousands, e.g., country code (string), event type (string).
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1055124977


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;

Review Comment:
   @zinking what do you mean `channel Id`? the downstream channel/subtask id?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1125391014


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.put(key, statistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) otherStatistics;

Review Comment:
   I think if the input is not in the right type, it would be better to throw IllegalArgumentException.
   I checked the other places as well, for example at https://github.com/apache/iceberg/blob/9cf9ca23eef196b3615f14dfa5cf6024dff85f6d/api/src/main/java/org/apache/iceberg/SortOrder.java#L245
   and https://github.com/apache/iceberg/blob/9cf9ca23eef196b3615f14dfa5cf6024dff85f6d/flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/data/FlinkSchemaVisitor.java#L89
   It's common that first checking type and then convert it into specific type.



##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestDataStatisticsOperator.java:
##########
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestDataStatisticsOperator {

Review Comment:
   yes, for now in the current `DataStatisticsOperator` implementation, we set `globalStatisticsState` in function `DataStatisticsOperator#snapshotState`. But to get `globalStatisticsState`, it needs to implement the function `DataStatisticsOperator#handleOperatorEvent` which we plan to do in another PR.  



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.put(key, statistics.getOrDefault(key, 0L) + 1L);

Review Comment:
   will update in latest commit



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsAndRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for data statistics and record. It is the only way for data statistics operator to send
+ * global data statistics to custom partitioner to distribute data based on statistics
+ *
+ * <p>DataStatisticsAndRecordWrapper is sent from {@link DataStatisticsOperator} to partitioner. It
+ * contains either data statistics(globally aggregated) or a record. Once partitioner receives the data
+ * statistics, it will use that to decide the coming record should send to which writer subtask. After
+ * shuffling, a filter and mapper are required to filter out the data distribution weight, unwrap the
+ * object and extract the original record type T.
+ */
+public class DataStatisticsAndRecordWrapper<T, K> implements Serializable {

Review Comment:
   The reason I rename it is, in `DataStatisticsOperator`, it generates the object which contains either global data statistics or record, while the place where shuffle happens is at the partitioner. The name `DataStatisticsAndRecord` is closer to the usage of the class(transmit global data statistics).
   
   Indeed, like you said, the other class which uses `And` contains both. For example, `RecordAndPosition` class, it contains both record and position. 
   How about naming it `DataStatisticsOrRecord`  even though there is no such kind of `or` class in the current repo :( 
   WDYT
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124250337


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.put(key, statistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) otherStatistics;

Review Comment:
   nit: if the input is not of type `MapDataStatistics` during the type conversion in line 48, will throw an error on its own. Therefore,  whether it is necessary to perform this `checkArgument`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064063897


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatistics.java:
##########
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatistics defines the interface to collect data statistics.
+ *
+ * <p>{@link ShuffleOperator} will store local data statistics and later distribute the global
+ * statistics(received from ShuffleCoordiantor) to Partitioner.
+ */
+interface DataStatistics<K> {

Review Comment:
   let's annotate this as `@Internal`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064081111


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getListState(generateGlobalDataDistributionWeightDescriptor());
+
+    if (globalDataStatisticsState.get() != null
+        && globalDataStatisticsState.get().iterator().hasNext()) {
+      globalDataStatistics = globalDataStatisticsState.get().iterator().next();
+    } else {
+      globalDataStatistics = dataStatisticsFactory.createDataStatistics();

Review Comment:
   this is a cold start problem (no state, rescale, etc.). does it make sense to create an empty Statistics object or we should keep it null to indicate cold start problem?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064210797


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {

Review Comment:
   I would prefer `add`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1067557305


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatistics.java:
##########
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatistics defines the interface to collect data statistics.
+ *
+ * <p>Data statistics tracks traffic volume distribution across data keys. For low-cardinality key,
+ * a simple map of (key, count) can be used. For high-cardinality key, probabilistic data structures
+ * (sketching) can be used.
+ */
+interface DataStatistics<K> {
+  long size();

Review Comment:
   let's add some Javadoc for the methods.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044092125


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+
+    assertNotNull(
+        stateContext
+            .getOperatorStateStore()
+            .getListState(operator.generateGlobalDataDistributionWeightDescriptor()));
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("a"));
+    assertTrue(operator.localDataStatisticsMap().containsKey("b"));
+    assertEquals(2L, (long) operator.localDataStatisticsMap().get("a"));
+    assertEquals(1L, (long) operator.localDataStatisticsMap().get("b"));
+  }
+
+  // ---------------- helper methods -------------------------
+
+  private StateInitializationContext getStateContext() throws Exception {
+    // Create the state context.
+    OperatorStateStore operatorStateStore = createOperatorStateStore();
+    return new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
+  }
+
+  private OperatorStateStore createOperatorStateStore() throws Exception {
+    MockEnvironment env = new MockEnvironmentBuilder().build();
+    AbstractStateBackend abstractStateBackend = new HashMapStateBackend();
+    CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
+    return abstractStateBackend.createOperatorStateBackend(

Review Comment:
   @yegangy0718 maybe we should use `OneInputStreamOperatorTestHarness`. you can refer to `TestIcebergStreamWriter`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1149748083


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());

Review Comment:
   OK. Will remove final



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.info("Taking data statistics operator snapshot for checkpoint {}", checkpointId);
+
+    // Only subtask 0 saves the state so that globalStatisticsState(UnionListState) stores
+    // an exact copy of globalStatistics
+    if (!globalStatistics.isEmpty() && getRuntimeContext().getIndexOfThisSubtask() == 0) {

Review Comment:
   it will change after first checkpoint. globalStatistics is being updated in the function #handleOperatorEvent when operator receives the event from coordinator. 
   The purpose of adding `!globalStatistics.isEmpty` is, in the first round of checkpoint, operator doesn't have any global statistics yet, when taking snapshot, it will store the empty `globalStatistics`  into `globalStatisticsState`
   



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.info("Taking data statistics operator snapshot for checkpoint {}", checkpointId);
+
+    // Only subtask 0 saves the state so that globalStatisticsState(UnionListState) stores
+    // an exact copy of globalStatistics
+    if (!globalStatistics.isEmpty() && getRuntimeContext().getIndexOfThisSubtask() == 0) {
+      globalStatisticsState.clear();
+      LOG.debug("Saving global statistics {} to state", globalStatistics);

Review Comment:
   will update log level to info and add subtask id in latest commit



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1144258901


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOrRecord.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for data statistics and record. It is the only way for data statistics operator to send
+ * global data statistics to custom partitioner to distribute data based on statistics
+ *
+ * <p>DataStatisticsOrRecord is sent from {@link DataStatisticsOperator} to partitioner. It

Review Comment:
   I actually think it's good to keep `DataStatisticsOrRecord is sent from {@link DataStatisticsOperator} to partitioner.` since it tells the readers where the class/object is being used. Once readers get the context, then they understand the later part `Once partitioner receives the data...`



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.

Review Comment:
   you're right. Data clustering is done by partitioner when distributing the data. I will remove this line as suggested



##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestDataStatisticsOperator.java:
##########
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestDataStatisticsOperator {
+  private DataStatisticsOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+    DataStatisticsFactory<String> dataStatisticsFactory = new MapDataStatisticsFactory<>();
+
+    this.operator = new DataStatisticsOperator<>(keySelector, mockGateway, dataStatisticsFactory);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatistics() instanceof MapDataStatistics);
+    MapDataStatistics<String> mapDataStatistics =
+        (MapDataStatistics<String>) operator.localDataStatistics();
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("a"));
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("b"));
+    assertEquals(2L, (long) mapDataStatistics.dataStatistics().get("a"));
+    assertEquals(1L, (long) mapDataStatistics.dataStatistics().get("b"));
+  }
+
+  @Test
+  public void testOperatorOutput() throws Exception {
+    try (OneInputStreamOperatorTestHarness<String, DataStatisticsOrRecord<String, String>>
+        testHarness = createHarness(this.operator)) {
+      testHarness.processElement(new StreamRecord<>("a"));
+      testHarness.processElement(new StreamRecord<>("b"));
+      testHarness.processElement(new StreamRecord<>("b"));
+
+      List<String> recordsOutput =
+          testHarness.extractOutputValues().stream()
+              .filter(DataStatisticsOrRecord::hasRecord)
+              .map(DataStatisticsOrRecord::record)
+              .collect(Collectors.toList());
+      assertThat(recordsOutput)
+          .containsExactlyInAnyOrderElementsOf(ImmutableList.of("a", "b", "b"));
+    }
+  }
+
+  private StateInitializationContext getStateContext() throws Exception {
+    // Create the state context.
+    OperatorStateStore operatorStateStore = createOperatorStateStore();
+    return new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
+  }
+
+  private OperatorStateStore createOperatorStateStore() throws Exception {

Review Comment:
   yes, can be merged



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.merge(key, 1L, Long::sum);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);

Review Comment:
   how about getting the type for statistics from the class name 
   `"Map statistics can not merge with " + otherStatistics.getClass()`
   



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private DataStatistics<K> localStatistics;

Review Comment:
   yes, right. I will update in the latest commit.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1065372357


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsFactory.java:
##########
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatisticsFactory provides the DataStatistics definition for different mode like HASH, RANGE
+ */
+class DataStatisticsFactory<K> {
+
+  DataStatistics<K> createDataStatistics() {

Review Comment:
   The factory will create different types DataStatistics based on different scenarios like `MapStatistics` for low-cardinary and `SketchStatistics`/ `DigestStatistics` for high-cardinary. Thus I would prefer to keep it general. 
   Regarding whether we need this factory, it depends. We can either use ShuffleOperatorFactory to get config(low or high cardinary or range mode) and then pass the right data statistics to operator and coordinator, or use this factory, pass it into operator and coordinator to create DataStatistics by calling this function.  



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.

Review Comment:
   will update



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {
+    dataStatistics.put(key, dataStatistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> other) {
+    Preconditions.checkArgument(
+        other instanceof MapDataStatistics, "Can not merge this type of statistics: " + other);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) other;
+    mapDataStatistic.dataStatistics.forEach(
+        (key, count) -> dataStatistics.put(key, dataStatistics.getOrDefault(key, 0L) + count));
+  }
+
+  @VisibleForTesting

Review Comment:
   yes, right. I will remove the annotation



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =

Review Comment:
   will add the check to line 81
   ```
       if (context.isRestored() && globalStatisticsState.get() != null
           && globalStatisticsState.get().iterator().hasNext())
   ```
   since we still need to initialize the globalDataStatisticsState variable.



##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+    DataStatisticsFactory<String> dataStatisticsFactory = new DataStatisticsFactory<>();
+
+    this.operator = new ShuffleOperator<>(keySelector, mockGateway, dataStatisticsFactory);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+
+    assertNotNull(
+        stateContext
+            .getOperatorStateStore()
+            .getListState(operator.generateGlobalDataDistributionWeightDescriptor()));
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatistics() instanceof MapDataStatistics);
+    MapDataStatistics<String> mapDataStatistics =
+        (MapDataStatistics<String>) operator.localDataStatistics();
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("a"));
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("b"));
+    assertEquals(2L, (long) mapDataStatistics.dataStatistics().get("a"));
+    assertEquals(1L, (long) mapDataStatistics.dataStatistics().get("b"));
+  }
+
+  @Test
+  public void testOperatorOutput() throws Exception {
+    try (OneInputStreamOperatorTestHarness<String, ShuffleRecordWrapper<String, String>>
+        testHarness = createHarness(this.operator)) {
+      testHarness.processElement(new StreamRecord<>("a"));
+      testHarness.processElement(new StreamRecord<>("b"));
+      testHarness.processElement(new StreamRecord<>("b"));
+
+      List<String> recordsOutput =
+          testHarness.extractOutputValues().stream()
+              .filter(ShuffleRecordWrapper::hasRecord)
+              .map(ShuffleRecordWrapper::record)
+              .collect(Collectors.toList());
+      assertThat(recordsOutput)
+          .containsExactlyInAnyOrderElementsOf(ImmutableList.of("a", "b", "b"));
+    }
+  }
+
+  // ---------------- helper methods -------------------------

Review Comment:
   will remove



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {

Review Comment:
   we use it in unit test now to get the state to make sure state is not null like what SourceOperatorTest did https://github.com/apache/flink/blob/68b37fb867374df5a201f0b170e35c21266e5d7b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/SourceOperatorTest.java#L87
   I was thinking we can use it to verify globalStatisticsState was updated as expected after taking snapshot. 
   But actually, we can also create a test-only visible function to return globalStatisticsState directly if we need to check its value. 



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getListState(generateGlobalDataDistributionWeightDescriptor());
+
+    if (globalDataStatisticsState.get() != null
+        && globalDataStatisticsState.get().iterator().hasNext()) {
+      globalDataStatistics = globalDataStatisticsState.get().iterator().next();
+    } else {
+      globalDataStatistics = dataStatisticsFactory.createDataStatistics();

Review Comment:
   We will always check `globalStatistics != null && globalStatistics.size() > 0` to decide whether to send globalStatistics to partitioner. 
   We don't have any special action for cold start yet. Leaving it to empty doesn't bring too much benefit as of now. 



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {

Review Comment:
   will rename the method to `add` in latest commit



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatistics.java:
##########
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatistics defines the interface to collect data statistics.
+ *
+ * <p>{@link ShuffleOperator} will store local data statistics and later distribute the global

Review Comment:
   Make sense to me. Will update in latest commit



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1054570635


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   Also, I am not sure that we want to seamlessly switch between different implementations. @stevenzwu's table describes the following possibilities in the doc where the strategy is based on the table definition and not on the statistics:
   ```
   Partitioned \ Sorted | Unsorted           | Sorted
   ---------------------| ------------------ | ----------------------------
   Unpartitioned        | None               | Range
   ---------------------| ------------------ | ----------------------------
   Partitioned          | BinPacking / Hash  | BinPacking if SortOrder only contains partition columns
                        |                    | Range if SortOder contains non-partition columns
   ```
   
   I think we can we can reasonably expect that we have enough memory to keep the statistics in a map if BinPacking is used. As for Range partitioning strategy, I would expect that we need to use some kind of Digest statistics. I think it would be a very rare corner case when the possible values are fluctuating so much that we need to change dynamically between the Digest and the Map based statistics.
   
   In the meantime we might want to provide a `DataStatisticsFactory` for the `ShuffleOperator` to generate the empty statistics, like this:
   ```
   public interface DataStatisticFactory<K> {
       DataStatistic<K> newDataStatistic();
       TypeInformation type();
   }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064080560


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.

Review Comment:
   nit: we can make the description a little more accurate. e.g.
   
   ```
   Shuffle operator collects traffic distribution statistics. A custom partitioner shall be attached to the shuffle operator
   output. The custom partitioner leverages the statistics to shuffle record to improve data clustering while 
   maintaining relative balanced traffic distribution to downstream subtasks.
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044076581


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   high-cardinality is out of the scope of initial impl. But I see @pvary 's point regarding compatibility for future support. For that, we may want to introduce a `TrafficStatistics` interface. Currently, we can provide a `SimpleCountStatistics` impl.  We can implement a `TrafficStatisticsSerializer`. The serializer can be expanded in the future to handle high-cardinality statistics like `TDigestStatistics` e.g.. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044079308


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   nit: variable name can be simpler as `localStatistics`. `Map` suffix is not necessary.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044083371


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, Map<K, Long> globalDataDistributionWeight) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataDistributionWeight != null,
+        "A ShuffleRecordWrapper has to contain exactly one of record and stats, not neither or both");
+    this.globalDataDistributionWeight = globalDataDistributionWeight;
+    this.record = record;
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromDistribution(
+      Map<K, Long> globalDataDistributionWeight) {
+    return new ShuffleRecordWrapper<>(null, globalDataDistributionWeight);
+  }
+
+  boolean hasGlobalDataDistributionWeight() {
+    return globalDataDistributionWeight != null;
+  }
+
+  boolean hasRecord() {
+    return record != null;
+  }
+
+  Map<K, Long> globalDataDistributionWeight() {
+    return globalDataDistributionWeight;
+  }
+
+  T record() {
+    return record;
+  }
+
+  @Override
+  public String toString() {
+    return String.format(

Review Comment:
   nit: use guava. e.g.
   ```
       return MoreObjects.toStringHelper(this)
           .add("contains_null", containsNull)
           .add("contains_nan", containsNaN)
           .add("lower_bound", lowerBound)
           .add("upper_bound", upperBound)
           .toString();
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064081012


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =

Review Comment:
   we need to check if context is restored
   ```
   if (context.isRestored()) {
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064064551


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {

Review Comment:
   from the implementation, I see `put` is probably not the most intuitive method name. For map count, this operation is more like `increment(K key)`. Not sure if `increment` will be meaningful to digest/probabilistic data structures. Alternative would be `add(K key)` or `add(K key, int count)`. 
   
   cc @pvary 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064064551


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {

Review Comment:
   from the implementation, I see `put` is probably not the most intuitive method name. For map count, this operation is more like `increment(K key)`. Not sure if `increment` will be meaningful to digest/probabilistic data structures. Alternative would be `add(K key)` or `add(K key, int count). 
   
   cc @pvary 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064064970


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatistics.java:
##########
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatistics defines the interface to collect data statistics.
+ *
+ * <p>{@link ShuffleOperator} will store local data statistics and later distribute the global

Review Comment:
   I feel here we don't need to refer to ShuffleOperator. But we do need to expand a little more on what `data statistics` mean? E.g.,
   
   ```
   Data statistics tracks traffic volume distribution across some key or keys. For low-cardinality key, a simple map of (key, count) can be used. For high-cardinality key, probabilistic data structures (sketching) can be used.
   ```
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064064970


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatistics.java:
##########
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatistics defines the interface to collect data statistics.
+ *
+ * <p>{@link ShuffleOperator} will store local data statistics and later distribute the global

Review Comment:
   I feel here we don't need to refer to ShuffleOperator. But we do need to expand a little more on what `data statistics` mean? E.g.,
   
   
   Data statistics tracks traffic volume distribution across some key or keys. For low-cardinality key, a simple map of (key, count) can be used. For high-cardinality key, probabilistic data structures (sketching) can be used.
   
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1065404653


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =

Review Comment:
   will add the check to line 81
   ```
       if (context.isRestored() && globalStatisticsState.get() != null
           && globalStatisticsState.get().iterator().hasNext())
   ```
   since we still need to initialize the globalDataStatisticsState variable.
   If we set globalDataStatisticsState to null when context.isRestored() = false, then when taking snapshot, we cannot execute `globalStatisticsState.add(globalStatistics)`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1043244839


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {

Review Comment:
   Might as well add tests for the rebalance case if it is not too complicated



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1061820186


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   It sounds good to me to add `DataStatistics` interface. 
   For the usage of `DataStatisticsFactory `, I think it will take table/mode configuration to decide which `DataStatistics` to generate. For now,  function `newDataStatistic/createDataStatistic` will always return `MapDataStatistic` since this is the only supported type. 
   I will make the modification in the latest commit. 



##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {

Review Comment:
   The logic to send data statistics from operator to coordinator and data statistics aggregation function in coordinator are not defined in the PR. The global state is empty. 
   I will add the rebalance case in the next PR once the above parts are there. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064080833


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;

Review Comment:
   nit: simplify the names a little by removing `Data`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064084354


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {

Review Comment:
   not sure if this test is necessary. seems too trivial to test.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044053748


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   @pvary discussed this with me offline a few days ago. Flink doc said that broadcast is usually used with broadcast stream. But it may just work for regular stream. This is the part yet to be verified. @yegangy0718 can you experiment the broadcast with some unit test (especially for rescale)?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044082721


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, Map<K, Long> globalDataDistributionWeight) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataDistributionWeight != null,
+        "A ShuffleRecordWrapper has to contain exactly one of record and stats, not neither or both");
+    this.globalDataDistributionWeight = globalDataDistributionWeight;
+    this.record = record;
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K extends Serializable> ShuffleRecordWrapper<T, K> fromDistribution(

Review Comment:
   nit: `fromStatistics` to be consistent with the usage of statistics.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1049262265


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   Let's merge this change, and we can play around the Broadcaststate in a followup PR. WTYT?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124680450


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsAndRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsAndRecordWrapper<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;

Review Comment:
   it will be used in the next PR with the shuffle coordinator to aggregate statistics globally



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124200310


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsAndRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsAndRecordWrapper<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;

Review Comment:
   it doesn't seem to be in use.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124213533


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsAndRecordWrapper<T, K>>

Review Comment:
   a quick question: Do we need to override `finish()` to cover the bounded case?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064211920


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  // TODO: support to store statistics for high cardinality cases
+  private transient DataStatistics<K> localDataStatistics;
+  private transient DataStatistics<K> globalDataStatistics;
+  private transient ListState<DataStatistics<K>> globalDataStatisticsState;
+  private transient DataStatisticsFactory<K> dataStatisticsFactory;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> dataStatisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.dataStatisticsFactory = dataStatisticsFactory;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<DataStatistics<K>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataStatisticsState", TypeInformation.of(new TypeHint<DataStatistics<K>>() {}));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();
+    globalDataStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getListState(generateGlobalDataDistributionWeightDescriptor());
+
+    if (globalDataStatisticsState.get() != null
+        && globalDataStatisticsState.get().iterator().hasNext()) {
+      globalDataStatistics = globalDataStatisticsState.get().iterator().next();
+    } else {
+      globalDataStatistics = dataStatisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    // TODO: handle scaling up
+    if (globalDataStatistics != null && globalDataStatistics.size() > 0) {
+      output.collect(new StreamRecord<>(ShuffleRecordWrapper.fromStatistics(globalDataStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with globalDataDistributionWeight from coordinator and update
+    // globalDataStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localDataStatistics.put(key);
+    output.collect(new StreamRecord<>(ShuffleRecordWrapper.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.debug("Taking shuffle operator snapshot for checkpoint {}", checkpointId);
+
+    // Set the globalDataDistributionWeightState state to the latest global value.
+    if (globalDataStatistics != null && globalDataStatistics.size() > 0) {
+      globalDataStatisticsState.clear();
+      globalDataStatisticsState.add(globalDataStatistics);
+    }
+
+    // TODO: send to coordinator
+    // For now we make it simple to send globalDataStatisticsState at checkpoint
+
+    // Reset the local data count
+    localDataStatistics = dataStatisticsFactory.createDataStatistics();

Review Comment:
   Either `reset` or `clear` method - or a factory for the `DataStatistics`.
   TDigest does not have a `reset` or `clear` like method, but we can hide it if we want to use it later. I think we should decide what creates a cleaner API.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064211285


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsFactory.java:
##########
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatisticsFactory provides the DataStatistics definition for different mode like HASH, RANGE
+ */
+class DataStatisticsFactory<K> {
+
+  DataStatistics<K> createDataStatistics() {

Review Comment:
   I would feel that we either need a `clean` method to the DataStatistics API, or need to abstract away the way to recreate a new one after sending it to the JobManager (at checkpoint). I do not feel that the Operator should know which type of statistics is used.
   
   Alternatively we could use statistics which automatically keep historical data - but I would not bother with them in this phase 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1119645660


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>Shuffle operator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the shuffle operator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>

Review Comment:
   `DataStatisticsOperator` is better than `StatisticsOperator`. I'm OK to update the class name since this operator indeed only collects data statistics. I will rename `ShuffleRecordWrapper` to `DataStatisticsAndRecordWrapper` as well. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Flink: Implement data statistics operator to collect traffic distribution for guiding smart shuffling

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1152944119


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsAndRecordWrapper<T, K>>

Review Comment:
   Ok, I see. Thank you for your explanation.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044053748


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;
+  private transient Map<K, Long> globalDataDistributionWeightMap;
+  private transient ListState<Map<K, Long>> globalDataDistributionWeightState;
+
+  public ShuffleOperator(
+      KeySelector<T, K> keySelector,
+      TypeInformation<K> keyType,
+      OperatorEventGateway operatorEventGateway) {
+    this.keySelector = keySelector;
+    this.keyType = keyType;
+    this.operatorEventGateway = operatorEventGateway;
+  }
+
+  @VisibleForTesting
+  ListStateDescriptor<Map<K, Long>> generateGlobalDataDistributionWeightDescriptor() {
+    return new ListStateDescriptor<>(
+        "globalDataDistributionWeight", new MapTypeInfo<>(keyType, TypeInformation.of(Long.class)));
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localDataStatisticsMap = Maps.newHashMap();
+    globalDataDistributionWeightState =

Review Comment:
   @pvary discussed this with me offline. Flink doc said that broadcast is usually used with broadcast stream. But it may just work for regular stream. This is the part yet to be verified. @yegangy0718 can you experiment the broadcast with some unit test (especially for rescale)?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1073022432


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record

Review Comment:
   removed



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064084432


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+    DataStatisticsFactory<String> dataStatisticsFactory = new DataStatisticsFactory<>();
+
+    this.operator = new ShuffleOperator<>(keySelector, mockGateway, dataStatisticsFactory);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+
+    assertNotNull(
+        stateContext
+            .getOperatorStateStore()
+            .getListState(operator.generateGlobalDataDistributionWeightDescriptor()));
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatistics() instanceof MapDataStatistics);
+    MapDataStatistics<String> mapDataStatistics =
+        (MapDataStatistics<String>) operator.localDataStatistics();
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("a"));
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("b"));
+    assertEquals(2L, (long) mapDataStatistics.dataStatistics().get("a"));
+    assertEquals(1L, (long) mapDataStatistics.dataStatistics().get("b"));
+  }
+
+  @Test
+  public void testOperatorOutput() throws Exception {
+    try (OneInputStreamOperatorTestHarness<String, ShuffleRecordWrapper<String, String>>
+        testHarness = createHarness(this.operator)) {
+      testHarness.processElement(new StreamRecord<>("a"));
+      testHarness.processElement(new StreamRecord<>("b"));
+      testHarness.processElement(new StreamRecord<>("b"));
+
+      List<String> recordsOutput =
+          testHarness.extractOutputValues().stream()
+              .filter(ShuffleRecordWrapper::hasRecord)
+              .map(ShuffleRecordWrapper::record)
+              .collect(Collectors.toList());
+      assertThat(recordsOutput)
+          .containsExactlyInAnyOrderElementsOf(ImmutableList.of("a", "b", "b"));
+    }
+  }
+
+  // ---------------- helper methods -------------------------

Review Comment:
   maybe remove this comment. doesn't seem to be the style that Iceberg code has



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1130256726


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsAndRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for data statistics and record. It is the only way for data statistics operator to send
+ * global data statistics to custom partitioner to distribute data based on statistics
+ *
+ * <p>DataStatisticsAndRecordWrapper is sent from {@link DataStatisticsOperator} to partitioner. It
+ * contains either data statistics(globally aggregated) or a record. Once partitioner receives the data
+ * statistics, it will use that to decide the coming record should send to which writer subtask. After
+ * shuffling, a filter and mapper are required to filter out the data distribution weight, unwrap the
+ * object and extract the original record type T.
+ */
+public class DataStatisticsAndRecordWrapper<T, K> implements Serializable {

Review Comment:
   `DataStatisticsOrRecord` sounds good to me



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1149722119


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.info("Taking data statistics operator snapshot for checkpoint {}", checkpointId);
+
+    // Only subtask 0 saves the state so that globalStatisticsState(UnionListState) stores
+    // an exact copy of globalStatistics
+    if (!globalStatistics.isEmpty() && getRuntimeContext().getIndexOfThisSubtask() == 0) {

Review Comment:
   this doesn't seem correct. we don't need the `!globalStatistics.isEmpty` condition here. otherwise, it will never change after the first checkpoint/add.



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));

Review Comment:
   nit: maybe cache the value at the beginning of this method. `T record = streamRecord.getValue()`. I saw more of this style in the Iceberg code.



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.info("Taking data statistics operator snapshot for checkpoint {}", checkpointId);

Review Comment:
   nit: add subtask index to the log. E.g. `Snapshotting data statistics operator for checkpoint {} in subtask {}`



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.info("Taking data statistics operator snapshot for checkpoint {}", checkpointId);
+
+    // Only subtask 0 saves the state so that globalStatisticsState(UnionListState) stores
+    // an exact copy of globalStatistics
+    if (!globalStatistics.isEmpty() && getRuntimeContext().getIndexOfThisSubtask() == 0) {
+      globalStatisticsState.clear();
+      LOG.debug("Saving global statistics {} to state", globalStatistics);

Review Comment:
   nit: I think this should actually be info level logging. maybe add subtask index to the log.



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());

Review Comment:
   nit: Iceberg coding style doesn't use final for stack local variables.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064064723


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {
+    dataStatistics.put(key, dataStatistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> other) {
+    Preconditions.checkArgument(
+        other instanceof MapDataStatistics, "Can not merge this type of statistics: " + other);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) other;
+    mapDataStatistic.dataStatistics.forEach(
+        (key, count) -> dataStatistics.put(key, dataStatistics.getOrDefault(key, 0L) + count));
+  }
+
+  @VisibleForTesting

Review Comment:
   I feel this will not just for testing purpose. In the end, partitioner would need to retrieve the actual map statistics.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1043242787


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   For high cardinality cases we might want to store the statistics in some `digest` implementation, like [t-digest](https://github.com/tdunning/t-digest) or [q-digest](https://github.com/addthis/stream-lib/blob/master/src/main/java/com/clearspring/analytics/stream/quantile/QDigest.java).
   
   Do not need any specific change ATM, but we might want to keep the road open for this .



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064083893


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+public class ShuffleRecordWrapper<T, K> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final DataStatistics<K> globalDataStatistics;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, DataStatistics<K> globalDataStatistics) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataStatistics != null,
+        "A ShuffleRecordWrapper contain either record and stats, not neither or both");
+    this.globalDataStatistics = globalDataStatistics;
+    this.record = record;
+  }
+
+  static <T, K> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K> ShuffleRecordWrapper<T, K> fromStatistics(DataStatistics<K> globalDataStatistics) {
+    return new ShuffleRecordWrapper<>(null, globalDataStatistics);
+  }
+
+  boolean hasGlobalDataDistributionWeight() {
+    return globalDataStatistics != null;
+  }
+
+  boolean hasRecord() {
+    return record != null;
+  }
+
+  DataStatistics<K> globalDataDistributionWeight() {

Review Comment:
   nit: maybe just `dataStatistics`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064083753


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or

Review Comment:
   > It may contain a record or data distribution weight.
   
   we can be more precise here. `It contains either a record or the data statistics (globally aggregated).`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064064168


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsFactory.java:
##########
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+/**
+ * DataStatisticsFactory provides the DataStatistics definition for different mode like HASH, RANGE
+ */
+class DataStatisticsFactory<K> {
+
+  DataStatistics<K> createDataStatistics() {

Review Comment:
   I am not sure if there is much value of this factory class. even if we want to keep it, this method should be called `createMapStatistics`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124206142


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsAndRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsAndRecordWrapper<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private DataStatistics<K> localStatistics;
+  private DataStatistics<K> globalStatistics;
+  private ListState<DataStatistics<K>> globalStatisticsState;
+
+  public DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =

Review Comment:
   It needs to be inside the block of `if (context.isRestored())`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124227506


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.put(key, statistics.getOrDefault(key, 0L) + 1L);

Review Comment:
   nit
   ```suggestion
       statistics.merge(key, 1L, Long::sum);
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124252381


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.put(key, statistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) otherStatistics;
+    mapDataStatistic.statistics.forEach(

Review Comment:
   also use `merge`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] hililiwei commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["hililiwei (via GitHub)" <gi...@apache.org>]

hililiwei commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124263299


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestDataStatisticsOperator.java:
##########
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestDataStatisticsOperator {

Review Comment:
   It seems that we haven't tested the recovery case?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1148171973


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.merge(key, 1L, Long::sum);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);

Review Comment:
   yes, simple class name works well too.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1149807331


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private transient volatile DataStatistics<K> localStatistics;
+  private transient volatile DataStatistics<K> globalStatistics;
+  private transient ListState<DataStatistics<K>> globalStatisticsState;
+
+  DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.info("Taking data statistics operator snapshot for checkpoint {}", checkpointId);
+
+    // Only subtask 0 saves the state so that globalStatisticsState(UnionListState) stores
+    // an exact copy of globalStatistics
+    if (!globalStatistics.isEmpty() && getRuntimeContext().getIndexOfThisSubtask() == 0) {

Review Comment:
   you are right. I misread the code.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on pull request #6382: Flink: Implement data statistics operator to collect traffic distribution for guiding smart shuffling

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#issuecomment-1486188577

   thanks @yegangy0718 for the contribution and @pvary and @hililiwei for the review


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu merged pull request #6382: Flink: Implement data statistics operator to collect traffic distribution for guiding smart shuffling

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu merged PR #6382:
URL: https://github.com/apache/iceberg/pull/6382


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] zinking commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

zinking commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1045453077


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;

Review Comment:
   wasn't this simply channel Id in the design spec ?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1046704878


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+@Internal
+public class ShuffleRecordWrapper<T, K extends Serializable> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final Map<K, Long> globalDataDistributionWeight;

Review Comment:
   We plan to create a StreamPartitioner which takes the global distribution, and then generate the channel assignment.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064080711


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record

Review Comment:
   nit: these comments seem not relevant anymore



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064083871


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+public class ShuffleRecordWrapper<T, K> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final DataStatistics<K> globalDataStatistics;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, DataStatistics<K> globalDataStatistics) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataStatistics != null,
+        "A ShuffleRecordWrapper contain either record and stats, not neither or both");
+    this.globalDataStatistics = globalDataStatistics;
+    this.record = record;
+  }
+
+  static <T, K> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K> ShuffleRecordWrapper<T, K> fromStatistics(DataStatistics<K> globalDataStatistics) {
+    return new ShuffleRecordWrapper<>(null, globalDataStatistics);
+  }
+
+  boolean hasGlobalDataDistributionWeight() {

Review Comment:
   let's be consistent here. by the variable name, the method name could be `hasDataStatistics`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1064065085


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapDataStatistics.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> dataStatistics = Maps.newHashMap();
+
+  @Override
+  public long size() {
+    return dataStatistics.size();
+  }
+
+  @Override
+  public void put(K key) {
+    dataStatistics.put(key, dataStatistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> other) {
+    Preconditions.checkArgument(
+        other instanceof MapDataStatistics, "Can not merge this type of statistics: " + other);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) other;

Review Comment:
   nit: I would call this `otherStatistics`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1044078604


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {

Review Comment:
   why do we need to implements `OneInputStreamOperator` interface? it doesn't seem that we are overriding this method.
   
   ```
   public interface OneInputStreamOperator<IN, OUT> extends StreamOperator<OUT>, Input<IN> {
       @Override
       default void setKeyContextElement(StreamRecord<IN> record) throws Exception {
           setKeyContextElement1(record);
       }
   }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

yegangy0718 commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1061820225


##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestShuffleOperator.java:
##########
@@ -0,0 +1,132 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestShuffleOperator {
+  private ShuffleOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+
+    this.operator =
+        new ShuffleOperator<>(keySelector, TypeInformation.of(String.class), mockGateway);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testInitializeState() throws Exception {

Review Comment:
   The logic to send data statistics from operator to coordinator and data statistics aggregation function in coordinator are not defined in the PR. Thus the global state is always empty when restoring. 
   I will add the rebalance case in the next PR once the above parts are there. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1061962729


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>It collects the data statistics information, sends to coordinator and gets the global data
+ * distribution weight from coordinator. Then it will ingest the weight into data stream(wrap by a
+ * class{@link ShuffleRecordWrapper}) and send to partitioner.
+ */
+@Internal
+public class ShuffleOperator<T, K extends Serializable>
+    extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>
+    implements OneInputStreamOperator<T, ShuffleRecordWrapper<T, K>>, OperatorEventHandler {
+
+  private static final long serialVersionUID = 1L;
+
+  private final KeySelector<T, K> keySelector;
+  // the type of the key to collect data statistics
+  private final TypeInformation<K> keyType;
+  private final OperatorEventGateway operatorEventGateway;
+  // key is generated by applying KeySelector to record
+  // value is the times key occurs
+  private transient Map<K, Long> localDataStatisticsMap;

Review Comment:
   Thanks for the change @yegangy0718!
   Looks good to me



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] pvary commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

[GitBox <gi...@apache.org>]

pvary commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1061961825


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleRecordWrapper.java:
##########
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for record and data distribution weight
+ *
+ * <p>ShuffleRecordWrapper is sent from ShuffleOperator to partitioner. It may contain a record or
+ * data distribution weight. Once partitioner receives the weight, it will use that to decide the
+ * coming record should send to which writer subtask. After shuffling, a filter and mapper are
+ * required to filter out the data distribution weight, unwrap the object and extract the original
+ * record type T.
+ */
+public class ShuffleRecordWrapper<T, K> implements Serializable {
+
+  private static final long serialVersionUID = 1L;
+
+  private final DataStatistics<K> globalDataStatistics;
+  private final T record;
+
+  private ShuffleRecordWrapper(T record, DataStatistics<K> globalDataStatistics) {
+    Preconditions.checkArgument(
+        record != null ^ globalDataStatistics != null,
+        "A ShuffleRecordWrapper contain either record and stats, not neither or both");
+    this.globalDataStatistics = globalDataStatistics;
+    this.record = record;
+  }
+
+  static <T, K> ShuffleRecordWrapper<T, K> fromRecord(T record) {
+    return new ShuffleRecordWrapper<>(record, null);
+  }
+
+  static <T, K> ShuffleRecordWrapper<T, K> fromStatistics(DataStatistics<K> globalDataStatistics) {
+    return new ShuffleRecordWrapper<>(null, globalDataStatistics);
+  }
+
+  boolean hasGlobalDataDistributionWeight() {
+    return globalDataStatistics != null;
+  }
+
+  boolean hasRecord() {
+    return record != null;
+  }
+
+  DataStatistics<K> globalDataDistributionWeight() {
+    return globalDataStatistics;
+  }
+
+  T record() {
+    return record;
+  }
+
+  @Override
+  public String toString() {
+    return MoreObjects.toStringHelper(this)
+        .add("globalDataStatistics", globalDataStatistics)

Review Comment:
   nit: We might want to add a toString for the `MapDataStatistics` as well



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1117772760


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/ShuffleOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * Shuffle operator can help to improve data clustering based on the key.
+ *
+ * <p>Shuffle operator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the shuffle operator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class ShuffleOperator<T, K> extends AbstractStreamOperator<ShuffleRecordWrapper<T, K>>

Review Comment:
   I have been thinking about the name of this class. This operator technically only does statistics calculation. Hence `ShuffleOperator` sounds like a misleading name. But `StatisticsOperator` is too generic. Maybe `DataStatisticsOperator`?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] yegangy0718 commented on pull request #6382: Implement ShuffleOperator to collect data statistics

["yegangy0718 (via GitHub)" <gi...@apache.org>]

yegangy0718 commented on PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#issuecomment-1469377105

   Hi @hililiwei,  thanks for reviewing the code change.  I have handled and replied to your comments. Let me know if you have more thoughts. 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1143943999


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.

Review Comment:
   nit: this sentence is not accurate. statistics operator itself doesn't help improve data clustering. we can probably remove this sentence as the following sentence seems good as a short summary sentence.



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.merge(key, 1L, Long::sum);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);

Review Comment:
   nit: the error msg can be a little clearer. note that maybe we can add a `String type()` to the `DataStatistics` interface. we shouldn't dump the whole statistics from toString() in the error msg.
   ```
   Map statistics can not merge with statistics type: " + otherStatistics.type()
   ```



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private DataStatistics<K> localStatistics;

Review Comment:
   The two statistics variables probably need to be `transient volatile`.
   
   The other `globalStatisticsState` need to be `transient`



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private DataStatistics<K> localStatistics;
+  private DataStatistics<K> globalStatistics;
+  private ListState<DataStatistics<K>> globalStatisticsState;
+
+  public DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.debug("Taking data statistics operator snapshot for checkpoint {}", checkpointId);
+
+    // Update globalStatisticsState with latest global statistics
+    if (!globalStatistics.isEmpty()) {
+      globalStatisticsState.clear();
+      globalStatisticsState.add(globalStatistics);

Review Comment:
   with union state, we only subtask 0 to checkpoint the state.



##########
flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/sink/shuffle/TestDataStatisticsOperator.java:
##########
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.OperatorStateStore;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.operators.coordination.MockOperatorEventGateway;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateInitializationContextImpl;
+import org.apache.flink.runtime.state.TestTaskStateManager;
+import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.util.MockOutput;
+import org.apache.flink.streaming.util.MockStreamConfig;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.MapDataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestDataStatisticsOperator {
+  private DataStatisticsOperator<String, String> operator;
+
+  private Environment getTestingEnvironment() {
+    return new StreamMockEnvironment(
+        new Configuration(),
+        new Configuration(),
+        new ExecutionConfig(),
+        1L,
+        new MockInputSplitProvider(),
+        1,
+        new TestTaskStateManager());
+  }
+
+  @Before
+  public void before() throws Exception {
+    MockOperatorEventGateway mockGateway = new MockOperatorEventGateway();
+    KeySelector<String, String> keySelector =
+        new KeySelector<String, String>() {
+          private static final long serialVersionUID = 7662520075515707428L;
+
+          @Override
+          public String getKey(String value) {
+            return value;
+          }
+        };
+    DataStatisticsFactory<String> dataStatisticsFactory = new MapDataStatisticsFactory<>();
+
+    this.operator = new DataStatisticsOperator<>(keySelector, mockGateway, dataStatisticsFactory);
+    Environment env = getTestingEnvironment();
+    this.operator.setup(
+        new OneInputStreamTask<String, String>(env),
+        new MockStreamConfig(new Configuration(), 1),
+        new MockOutput<>(Lists.newArrayList()));
+  }
+
+  @After
+  public void clean() throws Exception {
+    operator.close();
+  }
+
+  @Test
+  public void testProcessElement() throws Exception {
+    StateInitializationContext stateContext = getStateContext();
+    operator.initializeState(stateContext);
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("a"));
+    operator.processElement(new StreamRecord<>("b"));
+    assertTrue(operator.localDataStatistics() instanceof MapDataStatistics);
+    MapDataStatistics<String> mapDataStatistics =
+        (MapDataStatistics<String>) operator.localDataStatistics();
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("a"));
+    assertTrue(mapDataStatistics.dataStatistics().containsKey("b"));
+    assertEquals(2L, (long) mapDataStatistics.dataStatistics().get("a"));
+    assertEquals(1L, (long) mapDataStatistics.dataStatistics().get("b"));
+  }
+
+  @Test
+  public void testOperatorOutput() throws Exception {
+    try (OneInputStreamOperatorTestHarness<String, DataStatisticsOrRecord<String, String>>
+        testHarness = createHarness(this.operator)) {
+      testHarness.processElement(new StreamRecord<>("a"));
+      testHarness.processElement(new StreamRecord<>("b"));
+      testHarness.processElement(new StreamRecord<>("b"));
+
+      List<String> recordsOutput =
+          testHarness.extractOutputValues().stream()
+              .filter(DataStatisticsOrRecord::hasRecord)
+              .map(DataStatisticsOrRecord::record)
+              .collect(Collectors.toList());
+      assertThat(recordsOutput)
+          .containsExactlyInAnyOrderElementsOf(ImmutableList.of("a", "b", "b"));
+    }
+  }
+
+  private StateInitializationContext getStateContext() throws Exception {
+    // Create the state context.
+    OperatorStateStore operatorStateStore = createOperatorStateStore();
+    return new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
+  }
+
+  private OperatorStateStore createOperatorStateStore() throws Exception {

Review Comment:
   nit: should this be merged with the method above?



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOrRecord.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for data statistics and record. It is the only way for data statistics operator to send
+ * global data statistics to custom partitioner to distribute data based on statistics
+ *
+ * <p>DataStatisticsOrRecord is sent from {@link DataStatisticsOperator} to partitioner. It

Review Comment:
   it seems that we can remove the comments before `It contains either ...`.
   
   We can also change `It contains either ...` to `This wrapper class contains either ...`



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/statistics/MapDataStatistics.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle.statistics;
+
+import java.util.Map;
+import org.apache.flink.annotation.Internal;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+
+/** MapDataStatistics uses map to count key frequency */
+@Internal
+public class MapDataStatistics<K> implements DataStatistics<K> {
+  private final Map<K, Long> statistics = Maps.newHashMap();
+
+  @Override
+  public boolean isEmpty() {
+    return statistics.size() == 0;
+  }
+
+  @Override
+  public void add(K key) {
+    // increase count of occurrence by one in the dataStatistics map
+    statistics.put(key, statistics.getOrDefault(key, 0L) + 1L);
+  }
+
+  @Override
+  public void merge(DataStatistics<K> otherStatistics) {
+    Preconditions.checkArgument(
+        otherStatistics instanceof MapDataStatistics,
+        "Can not merge this type of statistics: " + otherStatistics);
+    MapDataStatistics<K> mapDataStatistic = (MapDataStatistics<K>) otherStatistics;

Review Comment:
   I agree with @yegangy0718 that it is a little better to check the type IllegalArgumentException



##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOperator.java:
##########
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeHint;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.coordination.OperatorEvent;
+import org.apache.flink.runtime.operators.coordination.OperatorEventGateway;
+import org.apache.flink.runtime.operators.coordination.OperatorEventHandler;
+import org.apache.flink.runtime.state.StateInitializationContext;
+import org.apache.flink.runtime.state.StateSnapshotContext;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatisticsFactory;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+
+/**
+ * DataStatisticsOperator can help to improve data clustering based on the key.
+ *
+ * <p>DataStatisticsOperator collects traffic distribution statistics. A custom partitioner shall be
+ * attached to the DataStatisticsOperator output. The custom partitioner leverages the statistics to
+ * shuffle record to improve data clustering while maintaining relative balanced traffic
+ * distribution to downstream subtasks.
+ */
+class DataStatisticsOperator<T, K> extends AbstractStreamOperator<DataStatisticsOrRecord<T, K>>
+    implements OneInputStreamOperator<T, DataStatisticsOrRecord<T, K>>, OperatorEventHandler {
+  private static final long serialVersionUID = 1L;
+
+  // keySelector will be used to generate key from data for collecting data statistics
+  private final KeySelector<T, K> keySelector;
+  private final OperatorEventGateway operatorEventGateway;
+  private final DataStatisticsFactory<K> statisticsFactory;
+  private DataStatistics<K> localStatistics;
+  private DataStatistics<K> globalStatistics;
+  private ListState<DataStatistics<K>> globalStatisticsState;
+
+  public DataStatisticsOperator(
+      KeySelector<T, K> keySelector,
+      OperatorEventGateway operatorEventGateway,
+      DataStatisticsFactory<K> statisticsFactory) {
+    this.keySelector = keySelector;
+    this.operatorEventGateway = operatorEventGateway;
+    this.statisticsFactory = statisticsFactory;
+  }
+
+  @Override
+  public void initializeState(StateInitializationContext context) throws Exception {
+    localStatistics = statisticsFactory.createDataStatistics();
+    globalStatisticsState =
+        context
+            .getOperatorStateStore()
+            .getUnionListState(
+                new ListStateDescriptor<>(
+                    "globalStatisticsState",
+                    TypeInformation.of(new TypeHint<DataStatistics<K>>() {})));
+
+    if (context.isRestored()) {
+      int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+      if (globalStatisticsState.get() == null
+          || !globalStatisticsState.get().iterator().hasNext()) {
+        LOG.warn("Subtask {} doesn't have global statistics state to restore", subtaskIndex);
+        globalStatistics = statisticsFactory.createDataStatistics();
+      } else {
+        LOG.info("Restoring global statistics state for subtask {}", subtaskIndex);
+        globalStatistics = globalStatisticsState.get().iterator().next();
+      }
+    } else {
+      globalStatistics = statisticsFactory.createDataStatistics();
+    }
+  }
+
+  @Override
+  public void open() throws Exception {
+    if (!globalStatistics.isEmpty()) {
+      output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromDataStatistics(globalStatistics)));
+    }
+  }
+
+  @Override
+  public void handleOperatorEvent(OperatorEvent evt) {
+    // TODO: receive event with aggregated statistics from coordinator and update globalStatistics
+  }
+
+  @Override
+  public void processElement(StreamRecord<T> streamRecord) throws Exception {
+    final K key = keySelector.getKey(streamRecord.getValue());
+    localStatistics.add(key);
+    output.collect(new StreamRecord<>(DataStatisticsOrRecord.fromRecord(streamRecord.getValue())));
+  }
+
+  @Override
+  public void snapshotState(StateSnapshotContext context) throws Exception {
+    long checkpointId = context.getCheckpointId();
+    LOG.debug("Taking data statistics operator snapshot for checkpoint {}", checkpointId);

Review Comment:
   this line can be info level since it is infrequent and small. this line should also log the subtask id too.
   
   We can add another debug or trace level logging that dumps the actual statistics.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1124675531


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsAndRecordWrapper.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for data statistics and record. It is the only way for data statistics operator to send
+ * global data statistics to custom partitioner to distribute data based on statistics
+ *
+ * <p>DataStatisticsAndRecordWrapper is sent from {@link DataStatisticsOperator} to partitioner. It
+ * contains either data statistics(globally aggregated) or a record. Once partitioner receives the data
+ * statistics, it will use that to decide the coming record should send to which writer subtask. After
+ * shuffling, a filter and mapper are required to filter out the data distribution weight, unwrap the
+ * object and extract the original record type T.
+ */
+public class DataStatisticsAndRecordWrapper<T, K> implements Serializable {

Review Comment:
   @yegangy0718  `ShuffleRecord` probably still makes sense for this wrapper class.
   
   I saw some other classes in Iceberg with names like `FooAndBar`. But they typically contains both `Foo` **and** `Bar`. Here is  an `or` relationship. Hence I don't know if `And` is the most accurate name. 
   
   Regardless, I prefer us remove the `Wrapper` suffice from the name. it can be described in the Javadoc



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org

[GitHub] [iceberg] stevenzwu commented on a diff in pull request #6382: Implement ShuffleOperator to collect data statistics

["stevenzwu (via GitHub)" <gi...@apache.org>]

stevenzwu commented on code in PR #6382:
URL: https://github.com/apache/iceberg/pull/6382#discussion_r1148171635


##########
flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/DataStatisticsOrRecord.java:
##########
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.flink.sink.shuffle;
+
+import java.io.Serializable;
+import org.apache.iceberg.flink.sink.shuffle.statistics.DataStatistics;
+import org.apache.iceberg.relocated.com.google.common.base.MoreObjects;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * The wrapper class for data statistics and record. It is the only way for data statistics operator to send
+ * global data statistics to custom partitioner to distribute data based on statistics
+ *
+ * <p>DataStatisticsOrRecord is sent from {@link DataStatisticsOperator} to partitioner. It

Review Comment:
   maybe add the additional context after the sentence of `it contains either ...`. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org