<commit_before>################################################################################
#  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.
################################################################################
import datetime
import decimal
import os
import uuid
from typing import Collection, Iterable

from pyflink.common import Row
from pyflink.common.serializer import TypeSerializer
from pyflink.common.typeinfo import Types
from pyflink.common.watermark_strategy import WatermarkStrategy, TimestampAssigner
from pyflink.datastream import (TimeCharacteristic, RuntimeContext, WindowAssigner, Trigger,
                                TriggerResult, CountWindow, SlotSharingGroup)
from pyflink.datastream.data_stream import DataStream
from pyflink.datastream.functions import (AggregateFunction, CoMapFunction, CoFlatMapFunction,
                                          MapFunction, FilterFunction, FlatMapFunction,
                                          KeyedCoProcessFunction, KeyedProcessFunction, KeySelector,
                                          ProcessFunction, ProcessWindowFunction, ReduceFunction,
                                          WindowFunction)
from pyflink.datastream.state import (ValueStateDescriptor, ListStateDescriptor, MapStateDescriptor,
                                      ReducingStateDescriptor, ReducingState, AggregatingState,
                                      AggregatingStateDescriptor)
from pyflink.datastream.window import (CountWindowSerializer, MergingWindowAssigner, TimeWindow,
                                       TimeWindowSerializer)
from pyflink.java_gateway import get_gateway
from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction
from pyflink.testing.test_case_utils import PyFlinkBatchTestCase, PyFlinkStreamingTestCase
from pyflink.util.java_utils import get_j_env_configuration


class DataStreamTests(object):

    def setUp(self) -> None:
        super(DataStreamTests, self).setUp()
        config = get_j_env_configuration(self.env._j_stream_execution_environment)
        config.setString("akka.ask.timeout", "20 s")
        self.test_sink = DataStreamTestSinkFunction()

    def tearDown(self) -> None:
        self.test_sink.clear()

    def assert_equals_sorted(self, expected, actual):
        expected.sort()
        actual.sort()
        self.assertEqual(expected, actual)

    def test_basic_operations(self):
        ds = self.env.from_collection([('ab', decimal.Decimal(1)),
                                       ('bdc', decimal.Decimal(2)),
                                       ('cfgs', decimal.Decimal(3)),
                                       ('deeefg', decimal.Decimal(4))],
                                      type_info=Types.ROW([Types.STRING(), Types.BIG_DEC()]))

        class MyMapFunction(MapFunction):
            def map(self, value):
                return Row(value[0], value[1] + 1, value[2])

        class MyFlatMapFunction(FlatMapFunction):
            def flat_map(self, value):
                if value[1] % 2 == 0:
                    yield value

        class MyFilterFunction(FilterFunction):
            def filter(self, value):
                return value[1] > 2

        (ds.map(lambda i: (i[0], len(i[0]), i[1]),
                output_type=Types.TUPLE([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
           .flat_map(MyFlatMapFunction(),
                     output_type=Types.TUPLE([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
           .filter(MyFilterFunction())
           .map(MyMapFunction(),
                output_type=Types.ROW([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
           .add_sink(self.test_sink))
        self.env.execute('test_basic_operations')
        results = self.test_sink.get_results()
        expected = ["+I[cfgs, 5, 3]",
                    "+I[deeefg, 7, 4]"]
        self.assert_equals_sorted(expected, results)

    def test_basic_operations_with_callable_function(self):
        ds = self.env.from_collection([('ab', decimal.Decimal(1)),
                                       ('bdc', decimal.Decimal(2)),
                                       ('cfgs', decimal.Decimal(3)),
                                       ('deeefg', decimal.Decimal(4))],
                                      type_info=Types.ROW([Types.STRING(), Types.BIG_DEC()]))

        def add_1(value):
            return Row(value[0], value[1] + 1, value[2])

        def reserve_even(value):
            if value[1] % 2 == 0:
                yield value

        (ds.map(lambda i: (i[0], len(i[0]), i[1]))
           .flat_map(reserve_even)
           .filter(lambda i: i[1] > 2)
           .map(add_1)
           .add_sink(self.test_sink))
        self.env.execute('test_basic_operations_with_callable_function')
        results = self.test_sink.get_results(True)
        expected = ["<Row('cfgs', 5, Decimal('3'))>",
                    "<Row('deeefg', 7, Decimal('4'))>"]
        self.assert_equals_sorted(expected, results)

    def test_basic_co_operations(self):
        python_file_dir = os.path.join(self.tempdir, "python_file_dir_" + str(uuid.uuid4()))
        os.mkdir(python_file_dir)
        python_file_path = os.path.join(python_file_dir, "test_stream_dependency_manage_lib.py")
        with open(python_file_path, 'w') as f:
            f.write("def add_two(a):\n    return a + 2")

        class MyCoFlatMapFunction(CoFlatMapFunction):

            def flat_map1(self, value):
                yield value + 1

            def flat_map2(self, value):
                yield value - 1

        class MyCoMapFunction(CoMapFunction):

            def map1(self, value):
                from test_stream_dependency_manage_lib import add_two
                return add_two(value)

            def map2(self, value):
                return value + 1

        self.env.add_python_file(python_file_path)
        ds_1 = self.env.from_collection([1, 2, 3, 4, 5])
        ds_2 = ds_1.map(lambda x: x * 2)

        (ds_1.connect(ds_2).flat_map(MyCoFlatMapFunction())
             .connect(ds_2).map(MyCoMapFunction())
             .add_sink(self.test_sink))
        self.env.execute("test_basic_co_operations")
        results = self.test_sink.get_results(True)
        expected = ['4', '5', '6', '7', '8', '3', '5', '7', '9', '11', '3', '5', '7', '9', '11']
        self.assert_equals_sorted(expected, results)

    def test_basic_co_operations_with_output_type(self):
        class MyCoMapFunction(CoMapFunction):

            def map1(self, value):
                return value + 2

            def map2(self, value):
                return value + 1

        class MyCoFlatMapFunction(CoFlatMapFunction):

            def flat_map1(self, value):
                yield value + 1

            def flat_map2(self, value):
                yield value - 1

        ds_1 = self.env.from_collection([1, 2, 3, 4, 5])
        ds_2 = ds_1.map(lambda x: x * 2)

        (ds_1.connect(ds_2).flat_map(MyCoFlatMapFunction(), output_type=Types.INT())
             .connect(ds_2).map(MyCoMapFunction(), output_type=Types.INT())
             .add_sink(self.test_sink))
        self.env.execute("test_basic_co_operations_with_output_type")
        results = self.test_sink.get_results()
        expected = ['4', '5', '6', '7', '8', '3', '5', '7', '9', '11', '3', '5', '7', '9', '11']
        self.assert_equals_sorted(expected, results)

    def test_keyed_co_process(self):
        ds1 = self.env.from_collection([("a", 1), ("b", 2), ("c", 3)],
                                       type_info=Types.ROW([Types.STRING(), Types.INT()]))
        ds2 = self.env.from_collection([("b", 2), ("c", 3), ("d", 4)],
                                       type_info=Types.ROW([Types.STRING(), Types.INT()]))
        ds1 = ds1.assign_timestamps_and_watermarks(
            WatermarkStrategy.for_monotonous_timestamps().with_timestamp_assigner(
                SecondColumnTimestampAssigner()))
        ds2 = ds2.assign_timestamps_and_watermarks(
            WatermarkStrategy.for_monotonous_timestamps().with_timestamp_assigner(
                SecondColumnTimestampAssigner()))
        ds1.connect(ds2) \
            .key_by(lambda x: x[0], lambda x: x[0]) \
            .process(MyKeyedCoProcessFunction()) \
            .map(lambda x: Row(x[0], x[1] + 1)) \
            .add_sink(self.test_sink)
        self.env.execute('test_keyed_co_process_function')
        results = self.test_sink.get_results(True)
        expected = ["<Row('a', 2)>",
                    "<Row('b', 2)>",
                    "<Row('b', 3)>",
                    "<Row('c', 2)>",
                    "<Row('c', 3)>",
                    "<Row('d', 2)>",
                    "<Row('on_timer', 4)>"]
        self.assert_equals_sorted(expected, results)<commit_msg>[FLINK-23742][python] Ignore test_keyed_co_process for the time being<commit_after>################################################################################
#  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.
################################################################################
import datetime
import decimal
import os
import unittest
import uuid
from typing import Collection, Iterable

from pyflink.common import Row
from pyflink.common.serializer import TypeSerializer
from pyflink.common.typeinfo import Types
from pyflink.common.watermark_strategy import WatermarkStrategy, TimestampAssigner
from pyflink.datastream import (TimeCharacteristic, RuntimeContext, WindowAssigner, Trigger,
                                TriggerResult, CountWindow, SlotSharingGroup)
from pyflink.datastream.data_stream import DataStream
from pyflink.datastream.functions import (AggregateFunction, CoMapFunction, CoFlatMapFunction,
                                          MapFunction, FilterFunction, FlatMapFunction,
                                          KeyedCoProcessFunction, KeyedProcessFunction, KeySelector,
                                          ProcessFunction, ProcessWindowFunction, ReduceFunction,
                                          WindowFunction)
from pyflink.datastream.state import (ValueStateDescriptor, ListStateDescriptor, MapStateDescriptor,
                                      ReducingStateDescriptor, ReducingState, AggregatingState,
                                      AggregatingStateDescriptor)
from pyflink.datastream.window import (CountWindowSerializer, MergingWindowAssigner, TimeWindow,
                                       TimeWindowSerializer)
from pyflink.java_gateway import get_gateway
from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction
from pyflink.testing.test_case_utils import PyFlinkBatchTestCase, PyFlinkStreamingTestCase
from pyflink.util.java_utils import get_j_env_configuration


class DataStreamTests(object):

    def setUp(self) -> None:
        super(DataStreamTests, self).setUp()
        config = get_j_env_configuration(self.env._j_stream_execution_environment)
        config.setString("akka.ask.timeout", "20 s")
        self.test_sink = DataStreamTestSinkFunction()

    def tearDown(self) -> None:
        self.test_sink.clear()

    def assert_equals_sorted(self, expected, actual):
        expected.sort()
        actual.sort()
        self.assertEqual(expected, actual)

    def test_basic_operations(self):
        ds = self.env.from_collection([('ab', decimal.Decimal(1)),
                                       ('bdc', decimal.Decimal(2)),
                                       ('cfgs', decimal.Decimal(3)),
                                       ('deeefg', decimal.Decimal(4))],
                                      type_info=Types.ROW([Types.STRING(), Types.BIG_DEC()]))

        class MyMapFunction(MapFunction):
            def map(self, value):
                return Row(value[0], value[1] + 1, value[2])

        class MyFlatMapFunction(FlatMapFunction):
            def flat_map(self, value):
                if value[1] % 2 == 0:
                    yield value

        class MyFilterFunction(FilterFunction):
            def filter(self, value):
                return value[1] > 2

        (ds.map(lambda i: (i[0], len(i[0]), i[1]),
                output_type=Types.TUPLE([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
           .flat_map(MyFlatMapFunction(),
                     output_type=Types.TUPLE([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
           .filter(MyFilterFunction())
           .map(MyMapFunction(),
                output_type=Types.ROW([Types.STRING(), Types.INT(), Types.BIG_DEC()]))
           .add_sink(self.test_sink))
        self.env.execute('test_basic_operations')
        results = self.test_sink.get_results()
        expected = ["+I[cfgs, 5, 3]",
                    "+I[deeefg, 7, 4]"]
        self.assert_equals_sorted(expected, results)

    def test_basic_operations_with_callable_function(self):
        ds = self.env.from_collection([('ab', decimal.Decimal(1)),
                                       ('bdc', decimal.Decimal(2)),
                                       ('cfgs', decimal.Decimal(3)),
                                       ('deeefg', decimal.Decimal(4))],
                                      type_info=Types.ROW([Types.STRING(), Types.BIG_DEC()]))

        def add_1(value):
            return Row(value[0], value[1] + 1, value[2])

        def reserve_even(value):
            if value[1] % 2 == 0:
                yield value

        (ds.map(lambda i: (i[0], len(i[0]), i[1]))
           .flat_map(reserve_even)
           .filter(lambda i: i[1] > 2)
           .map(add_1)
           .add_sink(self.test_sink))
        self.env.execute('test_basic_operations_with_callable_function')
        results = self.test_sink.get_results(True)
        expected = ["<Row('cfgs', 5, Decimal('3'))>",
                    "<Row('deeefg', 7, Decimal('4'))>"]
        self.assert_equals_sorted(expected, results)

    def test_basic_co_operations(self):
        python_file_dir = os.path.join(self.tempdir, "python_file_dir_" + str(uuid.uuid4()))
        os.mkdir(python_file_dir)
        python_file_path = os.path.join(python_file_dir, "test_stream_dependency_manage_lib.py")
        with open(python_file_path, 'w') as f:
            f.write("def add_two(a):\n    return a + 2")

        class MyCoFlatMapFunction(CoFlatMapFunction):

            def flat_map1(self, value):
                yield value + 1

            def flat_map2(self, value):
                yield value - 1

        class MyCoMapFunction(CoMapFunction):

            def map1(self, value):
                from test_stream_dependency_manage_lib import add_two
                return add_two(value)

            def map2(self, value):
                return value + 1

        self.env.add_python_file(python_file_path)
        ds_1 = self.env.from_collection([1, 2, 3, 4, 5])
        ds_2 = ds_1.map(lambda x: x * 2)

        (ds_1.connect(ds_2).flat_map(MyCoFlatMapFunction())
             .connect(ds_2).map(MyCoMapFunction())
             .add_sink(self.test_sink))
        self.env.execute("test_basic_co_operations")
        results = self.test_sink.get_results(True)
        expected = ['4', '5', '6', '7', '8', '3', '5', '7', '9', '11', '3', '5', '7', '9', '11']
        self.assert_equals_sorted(expected, results)

    def test_basic_co_operations_with_output_type(self):
        class MyCoMapFunction(CoMapFunction):

            def map1(self, value):
                return value + 2

            def map2(self, value):
                return value + 1

        class MyCoFlatMapFunction(CoFlatMapFunction):

            def flat_map1(self, value):
                yield value + 1

            def flat_map2(self, value):
                yield value - 1

        ds_1 = self.env.from_collection([1, 2, 3, 4, 5])
        ds_2 = ds_1.map(lambda x: x * 2)

        (ds_1.connect(ds_2).flat_map(MyCoFlatMapFunction(), output_type=Types.INT())
             .connect(ds_2).map(MyCoMapFunction(), output_type=Types.INT())
             .add_sink(self.test_sink))
        self.env.execute("test_basic_co_operations_with_output_type")
        results = self.test_sink.get_results()
        expected = ['4', '5', '6', '7', '8', '3', '5', '7', '9', '11', '3', '5', '7', '9', '11']
        self.assert_equals_sorted(expected, results)

    @unittest.skip("FLINK-23742")
    def test_keyed_co_process(self):
        ds1 = self.env.from_collection([("a", 1), ("b", 2), ("c", 3)],
                                       type_info=Types.ROW([Types.STRING(), Types.INT()]))
        ds2 = self.env.from_collection([("b", 2), ("c", 3), ("d", 4)],
                                       type_info=Types.ROW([Types.STRING(), Types.INT()]))
        ds1 = ds1.assign_timestamps_and_watermarks(
            WatermarkStrategy.for_monotonous_timestamps().with_timestamp_assigner(
                SecondColumnTimestampAssigner()))
        ds2 = ds2.assign_timestamps_and_watermarks(
            WatermarkStrategy.for_monotonous_timestamps().with_timestamp_assigner(
                SecondColumnTimestampAssigner()))
        ds1.connect(ds2) \
            .key_by(lambda x: x[0], lambda x: x[0]) \
            .process(MyKeyedCoProcessFunction()) \
            .map(lambda x: Row(x[0], x[1] + 1)) \
            .add_sink(self.test_sink)
        self.env.execute('test_keyed_co_process_function')
        results = self.test_sink.get_results(True)
        expected = ["<Row('a', 2)>",
                    "<Row('b', 2)>",
                    "<Row('b', 3)>",
                    "<Row('c', 2)>",
                    "<Row('c', 3)>",
                    "<Row('d', 2)>",
                    "<Row('on_timer', 4)>"]
        self.assert_equals_sorted(expected, results)
{'fraction_non_alphanumeric': 0.13183949973944764, 'fraction_numerical': 0.0160500260552371, 'mean_word_length': 2.281073687809882, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 46, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6199553', 'n_tokens_mistral': 5703, 'n_tokens_neox': 5446, 'n_words': 1154}