File size: 21,231 Bytes
5fe31b7
e7f788e
 
fed4eb4
c4a2537
 
9e30502
c4a2537
9e30502
e7f788e
 
 
 
 
 
9e30502
fed4eb4
 
 
 
 
 
 
 
 
 
 
 
9e30502
 
 
 
 
 
 
 
fed4eb4
 
 
9e30502
fed4eb4
 
9e30502
 
fed4eb4
9e30502
fed4eb4
 
 
9e30502
 
 
fed4eb4
9e30502
 
fed4eb4
 
 
 
 
9e30502
 
 
e7f788e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9e30502
5fe31b7
1fbdf4c
 
cb31369
 
1fbdf4c
 
 
 
 
 
 
 
e7f788e
 
5fe31b7
1fbdf4c
 
 
 
 
 
 
e7f788e
 
36b5223
 
1fbdf4c
36b5223
 
 
cb31369
e7f788e
36b5223
e7f788e
 
 
36b5223
5fe31b7
1fbdf4c
a66b8be
 
36b5223
a66b8be
1fbdf4c
a66b8be
 
1fbdf4c
 
cb31369
fc6c7eb
 
36b5223
a66b8be
 
 
 
e7f788e
 
 
 
 
 
36b5223
5fe31b7
1fbdf4c
36b5223
1fbdf4c
 
 
 
cb31369
fc6c7eb
 
e7f788e
a66b8be
 
 
 
e7f788e
 
 
 
 
 
 
5fe31b7
1fbdf4c
 
 
cb31369
fc6c7eb
 
e7f788e
a66b8be
 
 
 
 
 
 
 
e7f788e
5fe31b7
 
 
e7f788e
5fe31b7
 
 
e7f788e
 
 
 
 
 
916e075
9e30502
916e075
 
 
e7f788e
 
5fe31b7
1fbdf4c
a66b8be
 
 
 
 
 
1fbdf4c
cb31369
5fe31b7
 
 
9e30502
a19f8c1
e7f788e
 
 
5fe31b7
 
 
 
 
 
 
e7f788e
 
 
 
5fe31b7
 
 
 
 
 
 
e7f788e
5fe31b7
a19f8c1
 
5fe31b7
a19f8c1
 
 
 
 
 
 
 
 
5fe31b7
 
 
 
a19f8c1
5fe31b7
 
 
 
e7f788e
 
5fe31b7
e7f788e
 
916e075
 
 
 
e7f788e
1fbdf4c
5fe31b7
e7f788e
a66b8be
 
 
 
 
 
 
1fbdf4c
cb31369
1fbdf4c
cb31369
5fe31b7
 
 
1fbdf4c
 
 
 
 
 
916e075
 
 
 
 
1fbdf4c
e7f788e
1fbdf4c
5fe31b7
1fbdf4c
e7f788e
916e075
 
 
 
e7f788e
 
5fe31b7
1fbdf4c
 
 
cb31369
e7f788e
 
 
 
 
 
 
 
 
fc6c7eb
e7f788e
 
 
 
5fe31b7
1fbdf4c
 
 
cb31369
5fe31b7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e7f788e
 
 
5fe31b7
 
 
e7f788e
 
916e075
 
 
e7f788e
 
5fe31b7
1fbdf4c
 
 
cb31369
e7f788e
 
 
 
5fe31b7
 
 
e7f788e
5fe31b7
 
 
916e075
e7f788e
 
 
5fe31b7
 
 
e7f788e
 
 
5fe31b7
1fbdf4c
 
 
cb31369
e7f788e
 
 
 
 
 
 
 
cb31369
5fe31b7
1fbdf4c
 
 
cb31369
e7f788e
 
 
 
 
 
 
 
 
 
5fe31b7
 
1fbdf4c
cb31369
1fbdf4c
 
 
cb31369
e7f788e
 
 
 
5fe31b7
 
 
 
e7f788e
 
 
 
 
 
 
 
 
 
 
 
5fe31b7
 
1fbdf4c
a66b8be
 
1fbdf4c
cb31369
5fe31b7
 
e7f788e
 
5fe31b7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e7f788e
5fe31b7
e7f788e
 
 
 
5fe31b7
 
1fbdf4c
a66b8be
36b5223
a66b8be
1fbdf4c
cb31369
e7f788e
 
 
 
5fe31b7
 
 
e7f788e
5fe31b7
e7f788e
 
 
 
5fe31b7
 
1fbdf4c
 
 
cb31369
e7f788e
a66b8be
e7f788e
 
5fe31b7
 
 
 
e7f788e
 
5fe31b7
e7f788e
5fe31b7
e7f788e
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
import re
from abc import abstractmethod
from dataclasses import field
from typing import Any, Dict, Generator, List, Optional, Union

from .artifact import Artifact
from .dataclass import InternalField, NonPositionalField
from .stream import MultiStream, Stream
from .utils import is_module_available


class Operator(Artifact):
    pass


class PackageRequirementsMixin(Artifact):
    """Base class used to automatically check for the existence of required python dependencies for an artifact (e.g. Operator or Metric).

    The _requirement list is either a list of required packages
    (e.g. ["torch","sentence_transformers"]) or a dictionary between required packages
    and detailed installation instructions on how how to install each package.
    (e.g. {"torch" : "Install Torch using `pip install torch`", "sentence_transformers" : Install Sentence Transformers using `pip install sentence-transformers`})
    Note that the package names should be specified as they are used in the python import statement for the package.
    """

    _requirements_list: Union[List[str], Dict[str, str]] = InternalField(
        default_factory=list
    )

    def verify(self):
        super().verify()
        self.check_missing_requirements()

    def check_missing_requirements(self, requirements=None):
        if requirements is None:
            requirements = self._requirements_list
        if isinstance(requirements, List):
            requirements = {package: "" for package in requirements}

        missing_packages = []
        installation_instructions = []
        for package, installation_instruction in requirements.items():
            if not is_module_available(package):
                missing_packages.append(package)
                installation_instructions.append(installation_instruction)
        if missing_packages:
            raise MissingRequirementsError(
                self.__class__.__name__, missing_packages, installation_instructions
            )


class MissingRequirementsError(Exception):
    def __init__(self, class_name, missing_packages, installation_instructions):
        self.class_name = class_name
        self.missing_packages = missing_packages
        self.installation_instruction = installation_instructions
        self.message = (
            f"{self.class_name} requires the following missing package(s): {', '.join(self.missing_packages)}. "
            + "\n".join(self.installation_instruction)
        )
        super().__init__(self.message)


class OperatorError(Exception):
    def __init__(self, exception: Exception, operators: List[Operator]):
        super().__init__(
            "This error was raised by the following operators: "
            + ",\n".join([str(operator) for operator in operators])
            + "."
        )
        self.exception = exception
        self.operators = operators

    @classmethod
    def from_operator_error(cls, exception: Exception, operator: Operator):
        return cls(exception.exception, [*exception.operators, operator])

    @classmethod
    def from_exception(cls, exception: Exception, operator: Operator):
        return cls(exception, [operator])


class StreamingOperator(Operator, PackageRequirementsMixin):
    """Base class for all stream operators in the streaming model.

    Stream operators are a key component of the streaming model and are responsible for processing continuous data streams.
    They perform operations such as transformations, aggregations, joins, windowing and more on these streams.
    There are several types of stream operators, including source operators, processing operators, etc.

    As a `StreamingOperator`, this class is responsible for performing operations on a stream, and must be implemented by all other specific types of stream operators in the system.
    When called, a `StreamingOperator` must return a MultiStream.

    As a subclass of `Artifact`, every `StreamingOperator` can be saved in a catalog for further usage or reference.

    """

    @abstractmethod
    def __call__(self, streams: Optional[MultiStream] = None) -> MultiStream:
        """Abstract method that performs operations on the stream.

        Args:
            streams (Optional[MultiStream]): The input MultiStream, which can be None.

        Returns:
            MultiStream: The output MultiStream resulting from the operations performed on the input.
        """


class SideEffectOperator(StreamingOperator):
    """Base class for operators that does not affect the stream."""

    def __call__(self, streams: Optional[MultiStream] = None) -> MultiStream:
        self.process()
        return streams

    @abstractmethod
    def process() -> None:
        pass


class SourceOperator(StreamingOperator):
    """A class representing a source operator in the streaming system.

    A source operator is responsible for generating the data stream from some source, such as a database or a file.
    This is the starting point of a stream processing pipeline.
    The `SourceOperator` class is a type of `SourceOperator`, which is a special type of `StreamingOperator`
    that generates an output stream but does not take any input streams.

    When called, a `SourceOperator` invokes its `process` method, which should be implemented by all subclasses
    to generate the required `MultiStream`.

    """

    caching: bool = NonPositionalField(default=None)

    def __call__(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
        multi_stream = self.process()
        if self.caching is not None:
            multi_stream.set_caching(self.caching)
        return multi_stream

    @abstractmethod
    def process(self) -> MultiStream:
        pass


class StreamInitializerOperator(SourceOperator):
    """A class representing a stream initializer operator in the streaming system.

    A stream initializer operator is a special type of `SourceOperator` that is capable of taking parameters during the stream generation process. This can be useful in situations where the stream generation process needs to be customized or configured based on certain parameters.

    When called, a `StreamInitializerOperator` invokes its `process` method, passing any supplied arguments and keyword arguments. The `process` method should be implemented by all subclasses to generate the required `MultiStream` based on the given arguments and keyword arguments.

    """

    caching: bool = NonPositionalField(default=None)

    def __call__(self, *args, **kwargs) -> MultiStream:
        multi_stream = self.process(*args, **kwargs)
        if self.caching is not None:
            multi_stream.set_caching(self.caching)
        return self.process(*args, **kwargs)

    @abstractmethod
    def process(self, *args, **kwargs) -> MultiStream:
        pass


class MultiStreamOperator(StreamingOperator):
    """A class representing a multi-stream operator in the streaming system.

    A multi-stream operator is a type of `StreamingOperator` that operates on an entire MultiStream object at once. It takes a `MultiStream` as input and produces a `MultiStream` as output. The `process` method should be implemented by subclasses to define the specific operations to be performed on the input `MultiStream`.
    """

    caching: bool = NonPositionalField(default=None)

    def __call__(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
        self.before_process_multi_stream()
        result = self._process_multi_stream(multi_stream)
        if self.caching is not None:
            result.set_caching(self.caching)
        return result

    def before_process_multi_stream(self):
        pass

    def _process_multi_stream(
        self, multi_stream: Optional[MultiStream] = None
    ) -> MultiStream:
        result = self.process(multi_stream)
        assert isinstance(
            result, MultiStream
        ), "MultiStreamOperator must return a MultiStream"
        return result

    @abstractmethod
    def process(self, multi_stream: MultiStream) -> MultiStream:
        pass

    def process_instance(self, instance, stream_name="tmp"):
        multi_stream = MultiStream({stream_name: [instance]})
        processed_multi_stream = self(multi_stream)
        return next(iter(processed_multi_stream[stream_name]))


class SingleStreamOperator(MultiStreamOperator):
    """A class representing a single-stream operator in the streaming system.

    A single-stream operator is a type of `MultiStreamOperator` that operates on individual
    `Stream` objects within a `MultiStream`. It iterates through each `Stream` in the `MultiStream`
    and applies the `process` method.
    The `process` method should be implemented by subclasses to define the specific operations
    to be performed on each `Stream`.

    """

    apply_to_streams: List[str] = NonPositionalField(
        default=None
    )  # None apply to all streams
    dont_apply_to_streams: List[str] = NonPositionalField(default=None)

    def _process_multi_stream(self, multi_stream: MultiStream) -> MultiStream:
        result = {}
        for stream_name, stream in multi_stream.items():
            if self._is_should_be_processed(stream_name):
                stream = self._process_single_stream(stream, stream_name)
            else:
                stream = stream
            assert isinstance(
                stream, Stream
            ), "SingleStreamOperator must return a Stream"
            result[stream_name] = stream

        return MultiStream(result)

    def _process_single_stream(
        self, stream: Stream, stream_name: Optional[str] = None
    ) -> Stream:
        return Stream(
            self._process_stream,
            gen_kwargs={"stream": stream, "stream_name": stream_name},
        )

    def _is_should_be_processed(self, stream_name):
        if (
            self.apply_to_streams is not None
            and self.dont_apply_to_streams is not None
            and stream_name in self.apply_to_streams
            and stream_name in self.dont_apply_to_streams
        ):
            raise ValueError(
                f"Stream '{stream_name}' can be in either apply_to_streams or dont_apply_to_streams not both."
            )

        return (
            self.apply_to_streams is None or stream_name in self.apply_to_streams
        ) and (
            self.dont_apply_to_streams is None
            or stream_name not in self.dont_apply_to_streams
        )

    def _process_stream(
        self, stream: Stream, stream_name: Optional[str] = None
    ) -> Generator:
        yield from self.process(stream, stream_name)

    @abstractmethod
    def process(self, stream: Stream, stream_name: Optional[str] = None) -> Generator:
        pass

    def process_instance(self, instance, stream_name="tmp"):
        processed_stream = self._process_single_stream([instance], stream_name)
        return next(iter(processed_stream))


class PagedStreamOperator(SingleStreamOperator):
    """A class representing a paged-stream operator in the streaming system.

    A paged-stream operator is a type of `SingleStreamOperator` that operates on a page of instances
    in a `Stream` at a time, where a page is a subset of instances.
    The `process` method should be implemented by subclasses to define the specific operations
    to be performed on each page.

    Args:
        page_size (int): The size of each page in the stream. Defaults to 1000.
    """

    page_size: int = 1000

    def _process_stream(
        self, stream: Stream, stream_name: Optional[str] = None
    ) -> Generator:
        page = []
        for instance in stream:
            page.append(instance)
            if len(page) >= self.page_size:
                yield from self.process(page, stream_name)
                page = []
        yield from self._process_page(page, stream_name)

    def _process_page(
        self, page: List[Dict], stream_name: Optional[str] = None
    ) -> Generator:
        yield from self.process(page, stream_name)

    @abstractmethod
    def process(self, page: List[Dict], stream_name: Optional[str] = None) -> Generator:
        pass

    def process_instance(self, instance, stream_name="tmp"):
        processed_stream = self._process_page([instance], stream_name)
        return next(iter(processed_stream))


class SingleStreamReducer(StreamingOperator):
    """A class representing a single-stream reducer in the streaming system.

    A single-stream reducer is a type of `StreamingOperator` that operates on individual `Stream` objects within a `MultiStream` and reduces each `Stream` to a single output value. The `process` method should be implemented by subclasses to define the specific reduction operation to be performed on each `Stream`.
    """

    def __call__(self, multi_stream: Optional[MultiStream] = None) -> Dict[str, Any]:
        result = {}
        for stream_name, stream in multi_stream.items():
            stream = self.process(stream)
            result[stream_name] = stream

        return result

    @abstractmethod
    def process(self, stream: Stream) -> Stream:
        pass


class StreamInstanceOperator(SingleStreamOperator):
    """A class representing a stream instance operator in the streaming system.

    A stream instance operator is a type of `SingleStreamOperator` that operates on individual instances within a `Stream`. It iterates through each instance in the `Stream` and applies the `process` method. The `process` method should be implemented by subclasses to define the specific operations to be performed on each instance.
    """

    def _process_stream(
        self, stream: Stream, stream_name: Optional[str] = None
    ) -> Generator:
        try:
            _index = None
            for _index, instance in enumerate(stream):
                yield self._process_instance(instance, stream_name)
        except Exception as e:
            if _index is None:
                raise e
            else:
                raise ValueError(
                    f"Error processing instance '{_index}' from stream '{stream_name}' in {self.__class__.__name__} due to: {e}"
                ) from e

    def _process_instance(
        self, instance: Dict[str, Any], stream_name: Optional[str] = None
    ) -> Dict[str, Any]:
        return self.process(instance, stream_name)

    @abstractmethod
    def process(
        self, instance: Dict[str, Any], stream_name: Optional[str] = None
    ) -> Dict[str, Any]:
        pass

    def process_instance(self, instance, stream_name="tmp"):
        return self._process_instance(instance, stream_name)


class StreamInstanceOperatorValidator(StreamInstanceOperator):
    """A class representing a stream instance operator validator in the streaming system.

    A stream instance operator validator is a type of `StreamInstanceOperator` that includes a validation step. It operates on individual instances within a `Stream` and validates the result of processing each instance.
    """

    @abstractmethod
    def validate(self, instance):
        pass

    def _process_stream(
        self, stream: Stream, stream_name: Optional[str] = None
    ) -> Generator:
        iterator = iter(stream)
        try:
            first_instance = next(iterator)
        except StopIteration as e:
            raise StopIteration(f"Stream '{stream_name}' is empty") from e
        result = self._process_instance(first_instance, stream_name)
        self.validate(result)
        yield result
        yield from (
            self._process_instance(instance, stream_name) for instance in iterator
        )


class InstanceOperator(Artifact):
    """A class representing an instance operator in the streaming system.

    An instance operator is a type of `Artifact` that operates on a single instance (represented as a dict) at a time. It takes an instance as input and produces a transformed instance as output.
    """

    def __call__(self, data: dict) -> dict:
        return self.process(data)

    @abstractmethod
    def process(self, data: dict) -> dict:
        pass


class BaseFieldOperator(Artifact):
    """A class representing a field operator in the streaming system.

    A field operator is a type of `Artifact` that operates on a single field within an instance. It takes an instance and a field name as input, processes the field, and updates the field in the instance with the processed value.
    """

    def __call__(self, data: Dict[str, Any], field: str) -> dict:
        value = self.process(data[field])
        data[field] = value
        return data

    @abstractmethod
    def process(self, value: Any) -> Any:
        pass


class InstanceOperatorWithMultiStreamAccess(StreamingOperator):
    """A class representing an instance operator with global access in the streaming system.

    An instance operator with global access is a type of `StreamingOperator` that operates on individual instances within a `Stream` and can also access other streams.
    It uses the `accessible_streams` attribute to determine which other streams it has access to.
    In order to make this efficient and to avoid qudratic complexity, it caches the accessible streams by default.
    """

    def __call__(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
        result = {}

        for stream_name, stream in multi_stream.items():
            stream = Stream(
                self.generator,
                gen_kwargs={"stream": stream, "multi_stream": multi_stream},
            )
            result[stream_name] = stream

        return MultiStream(result)

    def generator(self, stream, multi_stream):
        yield from (self.process(instance, multi_stream) for instance in stream)

    @abstractmethod
    def process(self, instance: dict, multi_stream: MultiStream) -> dict:
        pass


class SequentialOperator(MultiStreamOperator):
    """A class representing a sequential operator in the streaming system.

    A sequential operator is a type of `MultiStreamOperator` that applies a sequence of other operators to a
    `MultiStream`. It maintains a list of `StreamingOperator`s and applies them in order to the `MultiStream`.
    """

    max_steps = None

    steps: List[StreamingOperator] = field(default_factory=list)

    def num_steps(self) -> int:
        return len(self.steps)

    def set_max_steps(self, max_steps):
        assert (
            max_steps <= self.num_steps()
        ), f"Max steps requested ({max_steps}) is larger than defined steps {self.num_steps()}"
        assert max_steps >= 1, f"Max steps requested ({max_steps}) is less than 1"
        self.max_steps = max_steps

    def get_last_step_description(self):
        last_step = (
            self.max_steps - 1 if self.max_steps is not None else len(self.steps) - 1
        )
        description = str(self.steps[last_step])
        return re.sub(r"\w+=None, ", "", description)

    def _get_max_steps(self):
        return self.max_steps if self.max_steps is not None else len(self.steps)

    def process(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
        for operator in self.steps[0 : self._get_max_steps()]:
            multi_stream = operator(multi_stream)
        return multi_stream


class SourceSequentialOperator(SequentialOperator):
    """A class representing a source sequential operator in the streaming system.

    A source sequential operator is a type of `SequentialOperator` that starts with a source operator.
    The first operator in its list of steps is a `SourceOperator`, which generates the initial `MultiStream`
    that the other operators then process.
    """

    def __call__(self) -> MultiStream:
        return super().__call__()

    def process(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
        assert (
            self.num_steps() > 0
        ), "Calling process on a SourceSequentialOperator without any steps"
        multi_stream = self.steps[0]()
        for operator in self.steps[1 : self._get_max_steps()]:
            multi_stream = operator(multi_stream)
        return multi_stream


class SequentialOperatorInitilizer(SequentialOperator):
    """A class representing a sequential operator initializer in the streaming system.

    A sequential operator initializer is a type of `SequntialOperator` that starts with a stream initializer operator. The first operator in its list of steps is a `StreamInitializerOperator`, which generates the initial `MultiStream` based on the provided arguments and keyword arguments.
    """

    def __call__(self, *args, **kwargs) -> MultiStream:
        return self.process(*args, **kwargs)

    def process(self, *args, **kwargs) -> MultiStream:
        assert (
            self.num_steps() > 0
        ), "Calling process on a SequentialOperatorInitilizer without any steps"

        assert isinstance(
            self.steps[0], StreamInitializerOperator
        ), "The first step in a SequentialOperatorInitilizer must be a StreamInitializerOperator"
        multi_stream = self.steps[0](*args, **kwargs)
        for operator in self.steps[1 : self._get_max_steps()]:
            multi_stream = operator(multi_stream)
        return multi_stream