[FLINK-22348][python] Fix the Python operators of Python DataStream API doesn't use managed memory in execute_and_collect

Author: HuangXingBo

flink-python/pyflink/datastream/data_stream.py

@@ -641,6 +641,8 @@ def execute_and_collect(self, job_execution_name: str = None, limit: int = None)
         :param job_execution_name: The name of the job execution.
         :param limit: The limit for the collected elements.
         """
+        JPythonConfigUtil = get_gateway().jvm.org.apache.flink.python.util.PythonConfigUtil
+        JPythonConfigUtil.declareManagedMemory(self._j_data_stream.getExecutionEnvironment())
         if job_execution_name is None and limit is None:
             return CloseableIterator(self._j_data_stream.executeAndCollect(), self.get_type())
         elif job_execution_name is not None and limit is None:

flink-python/pyflink/datastream/tests/test_data_stream.py

@@ -387,11 +387,11 @@ def test_execute_and_collect(self):
                       decimal.Decimal('2000000000000000000.061111111111111'
                                       '11111111111111'))]
         expected = test_data
-        ds = self.env.from_collection(test_data)
+        ds = self.env.from_collection(test_data).map(lambda a: a)
         with ds.execute_and_collect() as results:
-            actual = []
-            for result in results:
-                actual.append(result)
+            actual = [result for result in results]
+            actual.sort()
+            expected.sort()
             self.assertEqual(expected, actual)
 
     def test_key_by_map(self):

flink-python/src/main/java/org/apache/flink/python/util/PythonConfigUtil.java

@@ -102,38 +102,7 @@ public static Configuration getEnvironmentConfig(StreamExecutionEnvironment env)
         return (Configuration) getConfigurationMethod.invoke(env);
     }
 
-    /**
-     * Configure the {@link OneInputPythonFunctionOperator} to be chained with the
-     * upstream/downstream operator by setting their parallelism, slot sharing group, co-location
-     * group to be the same, and applying a {@link ForwardPartitioner}. 1. operator with name
-     * "_keyed_stream_values_operator" should align with its downstream operator. 2. operator with
-     * name "_stream_key_by_map_operator" should align with its upstream operator.
-     */
-    private static void alignStreamNode(StreamNode streamNode, StreamGraph streamGraph) {
-        if (streamNode.getOperatorName().equals(KEYED_STREAM_VALUE_OPERATOR_NAME)) {
-            StreamEdge downStreamEdge = streamNode.getOutEdges().get(0);
-            StreamNode downStreamNode = streamGraph.getStreamNode(downStreamEdge.getTargetId());
-            chainStreamNode(downStreamEdge, streamNode, downStreamNode);
-            downStreamEdge.setPartitioner(new ForwardPartitioner());
-        }
-
-        if (streamNode.getOperatorName().equals(STREAM_KEY_BY_MAP_OPERATOR_NAME)
-                || streamNode.getOperatorName().equals(STREAM_PARTITION_CUSTOM_MAP_OPERATOR_NAME)) {
-            StreamEdge upStreamEdge = streamNode.getInEdges().get(0);
-            StreamNode upStreamNode = streamGraph.getStreamNode(upStreamEdge.getSourceId());
-            chainStreamNode(upStreamEdge, streamNode, upStreamNode);
-        }
-    }
-
-    private static void chainStreamNode(
-            StreamEdge streamEdge, StreamNode firstStream, StreamNode secondStream) {
-        streamEdge.setPartitioner(new ForwardPartitioner<>());
-        firstStream.setParallelism(secondStream.getParallelism());
-        firstStream.setCoLocationGroup(secondStream.getCoLocationGroup());
-        firstStream.setSlotSharingGroup(secondStream.getSlotSharingGroup());
-    }
-
-    /** Set Python Operator Use Managed Memory. */
+    /** Set Table Python Operator Use Managed Memory. */
     public static void declareManagedMemory(
             Transformation<?> transformation,
             StreamExecutionEnvironment env,
@@ -144,14 +113,12 @@ public static void declareManagedMemory(
         }
     }
 
-    private static void declareManagedMemory(Transformation<?> transformation) {
-        if (isPythonOperator(transformation)) {
-            transformation.declareManagedMemoryUseCaseAtSlotScope(ManagedMemoryUseCase.PYTHON);
-        }
-        List<Transformation<?>> inputTransformations = transformation.getInputs();
-        for (Transformation inputTransformation : inputTransformations) {
-            declareManagedMemory(inputTransformation);
-        }
+    /** Set DataStream Python Operator Use Managed Memory. */
+    public static void declareManagedMemory(StreamExecutionEnvironment env)
+            throws IllegalAccessException, NoSuchMethodException, InvocationTargetException,
+                    NoSuchFieldException {
+        Configuration config = getEnvConfigWithDependencies(env);
+        declareManagedMemory(env, config);
     }
 
     /**
@@ -169,17 +136,7 @@ public static StreamGraph generateStreamGraphWithDependencies(
 
         boolean executedInBatchMode = isExecuteInBatchMode(env, mergedConfig);
 
-        if (mergedConfig.getBoolean(PythonOptions.USE_MANAGED_MEMORY)) {
-            Field transformationsField =
-                    StreamExecutionEnvironment.class.getDeclaredField("transformations");
-            transformationsField.setAccessible(true);
-            for (Transformation transform :
-                    (List<Transformation<?>>) transformationsField.get(env)) {
-                if (isPythonOperator(transform)) {
-                    transform.declareManagedMemoryUseCaseAtSlotScope(ManagedMemoryUseCase.PYTHON);
-                }
-            }
-        }
+        declareManagedMemory(env, mergedConfig);
 
         String jobName =
                 getEnvironmentConfig(env)
@@ -217,6 +174,70 @@ public static StreamGraph generateStreamGraphWithDependencies(
         return streamGraph;
     }
 
+    public static Configuration getMergedConfig(
+            StreamExecutionEnvironment env, TableConfig tableConfig) {
+        try {
+            Configuration config = new Configuration(getEnvironmentConfig(env));
+            PythonDependencyUtils.merge(config, tableConfig.getConfiguration());
+            Configuration mergedConfig =
+                    PythonDependencyUtils.configurePythonDependencies(env.getCachedFiles(), config);
+            mergedConfig.setString("table.exec.timezone", tableConfig.getLocalTimeZone().getId());
+            return mergedConfig;
+        } catch (IllegalAccessException | NoSuchMethodException | InvocationTargetException e) {
+            throw new TableException("Method getMergedConfig failed.", e);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    public static Configuration getMergedConfig(ExecutionEnvironment env, TableConfig tableConfig) {
+        try {
+            Field field = ExecutionEnvironment.class.getDeclaredField("cacheFile");
+            field.setAccessible(true);
+            Configuration config = new Configuration(env.getConfiguration());
+            PythonDependencyUtils.merge(config, tableConfig.getConfiguration());
+            Configuration mergedConfig =
+                    PythonDependencyUtils.configurePythonDependencies(
+                            (List<Tuple2<String, DistributedCache.DistributedCacheEntry>>)
+                                    field.get(env),
+                            config);
+            mergedConfig.setString("table.exec.timezone", tableConfig.getLocalTimeZone().getId());
+            return mergedConfig;
+        } catch (NoSuchFieldException | IllegalAccessException e) {
+            throw new TableException("Method getMergedConfig failed.", e);
+        }
+    }
+
+    /**
+     * Configure the {@link OneInputPythonFunctionOperator} to be chained with the
+     * upstream/downstream operator by setting their parallelism, slot sharing group, co-location
+     * group to be the same, and applying a {@link ForwardPartitioner}. 1. operator with name
+     * "_keyed_stream_values_operator" should align with its downstream operator. 2. operator with
+     * name "_stream_key_by_map_operator" should align with its upstream operator.
+     */
+    private static void alignStreamNode(StreamNode streamNode, StreamGraph streamGraph) {
+        if (streamNode.getOperatorName().equals(KEYED_STREAM_VALUE_OPERATOR_NAME)) {
+            StreamEdge downStreamEdge = streamNode.getOutEdges().get(0);
+            StreamNode downStreamNode = streamGraph.getStreamNode(downStreamEdge.getTargetId());
+            chainStreamNode(downStreamEdge, streamNode, downStreamNode);
+            downStreamEdge.setPartitioner(new ForwardPartitioner());
+        }
+
+        if (streamNode.getOperatorName().equals(STREAM_KEY_BY_MAP_OPERATOR_NAME)
+                || streamNode.getOperatorName().equals(STREAM_PARTITION_CUSTOM_MAP_OPERATOR_NAME)) {
+            StreamEdge upStreamEdge = streamNode.getInEdges().get(0);
+            StreamNode upStreamNode = streamGraph.getStreamNode(upStreamEdge.getSourceId());
+            chainStreamNode(upStreamEdge, streamNode, upStreamNode);
+        }
+    }
+
+    private static void chainStreamNode(
+            StreamEdge streamEdge, StreamNode firstStream, StreamNode secondStream) {
+        streamEdge.setPartitioner(new ForwardPartitioner<>());
+        firstStream.setParallelism(secondStream.getParallelism());
+        firstStream.setCoLocationGroup(secondStream.getCoLocationGroup());
+        firstStream.setSlotSharingGroup(secondStream.getSlotSharingGroup());
+    }
+
     private static boolean isPythonOperator(StreamOperatorFactory streamOperatorFactory) {
         if (streamOperatorFactory instanceof SimpleOperatorFactory) {
             return ((SimpleOperatorFactory) streamOperatorFactory).getOperator()
@@ -295,36 +316,29 @@ private static boolean isExecuteInBatchMode(
         return !existsUnboundedSource;
     }
 
-    public static Configuration getMergedConfig(
-            StreamExecutionEnvironment env, TableConfig tableConfig) {
-        try {
-            Configuration config = new Configuration(getEnvironmentConfig(env));
-            PythonDependencyUtils.merge(config, tableConfig.getConfiguration());
-            Configuration mergedConfig =
-                    PythonDependencyUtils.configurePythonDependencies(env.getCachedFiles(), config);
-            mergedConfig.setString("table.exec.timezone", tableConfig.getLocalTimeZone().getId());
-            return mergedConfig;
-        } catch (IllegalAccessException | NoSuchMethodException | InvocationTargetException e) {
-            throw new TableException("Method getMergedConfig failed.", e);
+    @SuppressWarnings("unchecked")
+    private static void declareManagedMemory(StreamExecutionEnvironment env, Configuration config)
+            throws NoSuchFieldException, IllegalAccessException {
+        if (config.getBoolean(PythonOptions.USE_MANAGED_MEMORY)) {
+            Field transformationsField =
+                    StreamExecutionEnvironment.class.getDeclaredField("transformations");
+            transformationsField.setAccessible(true);
+            for (Transformation transform :
+                    (List<Transformation<?>>) transformationsField.get(env)) {
+                if (isPythonOperator(transform)) {
+                    transform.declareManagedMemoryUseCaseAtSlotScope(ManagedMemoryUseCase.PYTHON);
+                }
+            }
         }
     }
 
-    @SuppressWarnings("unchecked")
-    public static Configuration getMergedConfig(ExecutionEnvironment env, TableConfig tableConfig) {
-        try {
-            Field field = ExecutionEnvironment.class.getDeclaredField("cacheFile");
-            field.setAccessible(true);
-            Configuration config = new Configuration(env.getConfiguration());
-            PythonDependencyUtils.merge(config, tableConfig.getConfiguration());
-            Configuration mergedConfig =
-                    PythonDependencyUtils.configurePythonDependencies(
-                            (List<Tuple2<String, DistributedCache.DistributedCacheEntry>>)
-                                    field.get(env),
-                            config);
-            mergedConfig.setString("table.exec.timezone", tableConfig.getLocalTimeZone().getId());
-            return mergedConfig;
-        } catch (NoSuchFieldException | IllegalAccessException e) {
-            throw new TableException("Method getMergedConfig failed.", e);
+    private static void declareManagedMemory(Transformation<?> transformation) {
+        if (isPythonOperator(transformation)) {
+            transformation.declareManagedMemoryUseCaseAtSlotScope(ManagedMemoryUseCase.PYTHON);
+        }
+        List<Transformation<?>> inputTransformations = transformation.getInputs();
+        for (Transformation inputTransformation : inputTransformations) {
+            declareManagedMemory(inputTransformation);
         }
     }
 }