mage-ai 0.9.70__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/data_preparation/models/block/data_integration/mixins.py

@@ -29,6 +29,7 @@ from mage_ai.data_preparation.models.constants import (
     PIPELINES_FOLDER,
     BlockLanguage,
     BlockType,
+    PipelineType,
 )
 from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
@@ -215,12 +216,11 @@ class DataIntegrationMixin:
             bool: True if it's a data integration block, False otherwise.
         """
         if not self.pipeline:
-
             return False
 
         actual_project: Project = pipeline_project
         if not actual_project:
-            actual_project = Project(self.pipeline.repo_config)
+            actual_project = self.pipeline.project
 
         if not actual_project.is_feature_enabled(
                         FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
@@ -229,7 +229,8 @@ class DataIntegrationMixin:
             return False
 
         if self.type in [BlockType.DATA_LOADER, BlockType.DATA_EXPORTER] and \
-                BlockLanguage.YAML == self.language:
+                BlockLanguage.YAML == self.language and \
+                self.pipeline.type != PipelineType.STREAMING:
 
             return True
 

mage_ai/data_preparation/models/block/dynamic/__init__.py

@@ -1,14 +1,23 @@
 import os
 from typing import Any, Dict, List, Tuple
 
+from pandas.core.construction import Optional
 
-def all_variable_uuids(block, partition: str = None) -> List[str]:
+
+def all_variable_uuids(
+    block, partition: Optional[str] = None, max_results: Optional[int] = None
+) -> List[str]:
     arr = __all_variable_uuids_and_file_paths_for_reducing_block_output(
         block,
         partition=partition,
     )
 
-    return list(set([tup[0] for tup in arr]))
+    arr = list(set([tup[0] for tup in arr]))
+
+    if max_results is not None:
+        arr = arr[:max_results]
+
+    return arr
 
 
 def reduce_output_from_block(
@@ -99,9 +108,11 @@ def __all_variable_uuids_and_file_paths_for_reducing_block_output(
             # e.g. output_0
             dir_name = os.path.basename(os.path.normpath(dir_path))
 
-            variable_uuid_and_file_paths.append((
-                dir_name,
-                dir_path,
-            ))
+            variable_uuid_and_file_paths.append(
+                (
+                    dir_name,
+                    dir_path,
+                )
+            )
 
     return variable_uuid_and_file_paths

mage_ai/data_preparation/models/block/dynamic/child.py

@@ -1,6 +1,5 @@
 import time
-from logging import Logger
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import pandas as pd
 
@@ -13,63 +12,16 @@ from mage_ai.data_preparation.models.block.dynamic.variables import (
     get_outputs_for_dynamic_block,
     get_outputs_for_dynamic_child,
 )
-from mage_ai.orchestration.db.models.schedules import BlockRun
-from mage_ai.shared.memory import get_memory_usage
-
-
-class DynamicChildController:
-    def __init__(self, block, block_run_id: int):
-        self.block = block
-        self.block_run_id = block_run_id
-        self._block_run = None
-        self._block_runs = None
-        self._pipeline_run = None
-
-    def __getattr__(self, name):
-        val = getattr(self.block, name)
-
-        def _missing(val=val, *args, **kwargs):
-            return val(*args, **kwargs)
-
-        if callable(val):
-            return _missing
-
-        return val
-
-    def block_run(self) -> BlockRun:
-        if self._block_run:
-            return self._block_run
-
-        self._block_run = BlockRun.query.get(self.block_run_id)
-
-        return self._block_run
-
-    def pipeline_run(self):
-        if self._pipeline_run:
-            return self._pipeline_run
-
-        self._pipeline_run = self.block_run().pipeline_run
-
-        return self._pipeline_run
-
-    def block_runs(self) -> List[BlockRun]:
-        if self._block_runs:
-            return self._block_runs
-
-        self._block_runs = BlockRun.query.filter(
-            BlockRun.pipeline_run_id == self.block_run().pipeline_run_id,
-        ).all()
-
-        return self._block_runs
+from mage_ai.data_preparation.models.block.dynamic.wrappers import (
+    DynamicBlockWrapperBase,
+)
 
-    def is_completed(self) -> bool:
-        return True
 
+class DynamicChildController(DynamicBlockWrapperBase):
     def execute_sync(
         self,
-        execution_partition: str = None,
-        logger: Logger = None,
-        logging_tags: Dict = None,
+        execution_partition: Optional[str] = None,
+        logging_tags: Optional[Dict] = None,
         **kwargs,
     ) -> List[Dict]:
         pipeline = self.block.pipeline
@@ -95,27 +47,10 @@ class DynamicChildController:
                 count = 0
                 while tries < 12 and count == 0:
                     # If this block tries to get the data too soon, it’ll return empty.
-                    def __get(
+                    lazy_variable_controller = get_outputs_for_dynamic_child(
+                        upstream_block,
                         execution_partition=execution_partition,
-                        logger=logger,
-                        logging_tags=logging_tags,
-                        upstream_block=upstream_block,
-                    ):
-                        return get_outputs_for_dynamic_child(
-                            upstream_block,
-                            execution_partition=execution_partition,
-                            logger=logger,
-                            logging_tags=logging_tags,
-                        )
-
-                    lazy_variable_controller = get_memory_usage(
-                        logger=logger,
                         logging_tags=logging_tags,
-                        message_prefix=(
-                            f'DynamicChild {self.block.uuid} '
-                            'execute_sync get_outputs_for_dynamic_child'
-                        ),
-                        wrapped_function=__get,
                     )
 
                     if lazy_variable_controller is not None:
@@ -128,8 +63,9 @@ class DynamicChildController:
                 counts_by_upstream_block_uuid[upstream_block.uuid] = count
 
                 if is_dynamic:
-                    metadata_by_upstream_block_uuid[upstream_block.uuid] = \
-                        [lazy_var_set.read_metadata() for lazy_var_set in lazy_variable_controller]
+                    metadata_by_upstream_block_uuid[upstream_block.uuid] = [
+                        lazy_var_set.read_metadata() for lazy_var_set in lazy_variable_controller
+                    ]
             elif is_dynamic:
                 tries = 0
                 count = 0
@@ -164,6 +100,9 @@ class DynamicChildController:
                 time.sleep(10)
                 tries += 1
 
+        if combos is None:
+            return []
+
         block_run_dicts = []
 
         for combo in combos:
@@ -180,31 +119,34 @@ class DynamicChildController:
                 if is_dynamic_child or is_dynamic:
                     count = counts_by_upstream_block_uuid.get(upstream_block.uuid)
 
-                    parent_index = dynamic_block_index % count
-
-                    if is_dynamic_child:
-                        block_runs = block_runs_by_block_uuid[upstream_block.uuid]
-                        block_runs = sorted(
-                            [br for br in block_runs if br.block_uuid != upstream_block.uuid],
-                            key=lambda br: br.id,
-                        )
-                        # This errors list index out of range
-                        if parent_index < len(block_runs):
-                            dynamic_upstream_block_uuids.append(block_runs[parent_index].block_uuid)
-
-                    if metadata_by_upstream_block_uuid.get(upstream_block.uuid):
-                        metadata = metadata_by_upstream_block_uuid[upstream_block.uuid]
-                        if parent_index < len(metadata):
-                            metadata = metadata[parent_index]
-
-                            if metadata.get('block_uuid'):
-                                block_run_block_uuid = metadata.get('block_uuid')
-
-                            if metadata.get('upstream_blocks'):
-                                upstream_blocks_from_metadata.extend(
-                                    metadata.get('upstream_blocks'),
+                    if count is not None and count >= 1:
+                        parent_index = dynamic_block_index % count
+
+                        if is_dynamic_child:
+                            block_runs = block_runs_by_block_uuid[upstream_block.uuid]
+                            block_runs = sorted(
+                                [br for br in block_runs if br.block_uuid != upstream_block.uuid],
+                                key=lambda br: br.id,
+                            )
+                            # This errors list index out of range
+                            if parent_index < len(block_runs):
+                                dynamic_upstream_block_uuids.append(
+                                    block_runs[parent_index].block_uuid
                                 )
 
+                        if metadata_by_upstream_block_uuid.get(upstream_block.uuid):
+                            metadata = metadata_by_upstream_block_uuid[upstream_block.uuid]
+                            if parent_index < len(metadata):
+                                metadata = metadata[parent_index]
+
+                                if metadata.get('block_uuid'):
+                                    block_run_block_uuid = metadata.get('block_uuid')
+
+                                if metadata.get('upstream_blocks'):
+                                    upstream_blocks_from_metadata.extend(
+                                        metadata.get('upstream_blocks'),
+                                    )
+
             block_run_dict = dict(
                 dynamic_block_index=dynamic_block_index,
                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
@@ -229,6 +171,3 @@ class DynamicChildController:
             block_runs.append(block_run)
 
         return block_runs
-
-    def run_tests(self, **kwargs):
-        pass

mage_ai/data_preparation/models/block/dynamic/constants.py

@@ -0,0 +1 @@
+CHILD_DATA_VARIABLE_UUID = 'output_0'

mage_ai/data_preparation/models/block/dynamic/counter.py

@@ -0,0 +1,296 @@
+from typing import Any, Iterable, List, Optional, Sequence
+
+from mage_ai.data.tabular.reader import read_metadata
+from mage_ai.data_preparation.models.block.dynamic.constants import (
+    CHILD_DATA_VARIABLE_UUID,
+)
+from mage_ai.data_preparation.models.block.dynamic.utils import (
+    is_dynamic_block,
+    is_dynamic_block_child,
+    should_reduce_output,
+)
+from mage_ai.data_preparation.models.variables.cache import VariableAggregateCache
+from mage_ai.data_preparation.models.variables.constants import (
+    VariableAggregateDataType,
+    VariableAggregateSummaryGroupType,
+)
+from mage_ai.data_preparation.models.variables.summarizer import (
+    aggregate_summary_info_for_all_variables,
+    dynamic_block_index_paths,
+    get_aggregate_summary_info,
+    get_part_uuids,
+)
+
+
+class DynamicItemCounter:
+    def __init__(
+        self,
+        block: Any,
+        dynamic_block_index: Optional[int] = None,
+        partition: Optional[str] = None,
+        variable_uuid: Optional[str] = None,
+    ):
+        self.block = block
+        self.dynamic_block_index = dynamic_block_index
+        self.variable_uuid = variable_uuid or CHILD_DATA_VARIABLE_UUID
+        self.partition = partition
+
+        self._output = None
+        self._summary_information = None
+
+    @classmethod
+    def build_counter(
+        cls,
+        block: Any,
+        downstream_block: Optional[Any] = None,
+        dynamic_block_index: Optional[int] = None,
+        partition: Optional[str] = None,
+        variable_uuid: Optional[str] = None,
+    ):
+        counter_class = None
+        is_dynamic_parent = (
+            (
+                block.should_dynamically_generate_block(downstream_block)
+                if downstream_block is not None
+                else block.is_dynamic_parent
+            )
+            if block.is_dynamic_v2
+            else is_dynamic_block(block)
+        )
+        is_dynamic_child = (
+            block.is_dynamic_child if block.is_dynamic_v2 else is_dynamic_block_child(block)
+        )
+        if is_dynamic_child:
+            if is_dynamic_parent:
+                counter_class = DynamicDuoItemCounter
+            else:
+                counter_class = DynamicChildItemCounter
+        elif is_dynamic_parent:
+            counter_class = DynamicBlockItemCounter
+
+        if counter_class is not None:
+            return counter_class(
+                block,
+                dynamic_block_index=dynamic_block_index,
+                partition=partition,
+                variable_uuid=variable_uuid,
+            )
+
+    @property
+    def variable_manager(self):
+        return self.block.variable_manager
+
+    @property
+    def pipeline(self):
+        return self.block.pipeline
+
+    @property
+    def output(self):
+        if self._output is not None:
+            return self._output
+
+        self._output = self.variable.read_data()
+        return self._output
+
+    @property
+    def variable(self):
+        return self.variable_manager.get_variable_object(
+            self.pipeline.uuid, self.block.uuid, self.variable_uuid, partition=self.partition
+        )
+
+    def item_count(self, downstream_block: Optional[Any] = None) -> int:
+        return 0
+
+
+class DynamicBlockItemCounter(DynamicItemCounter):
+    @property
+    def summary_information(self) -> Optional[VariableAggregateCache]:
+        if self._summary_information is not None:
+            return self._summary_information
+
+        self._summary_information = get_aggregate_summary_info(
+            self.block.variable_manager,
+            pipeline_uuid=self.block.pipeline.uuid,
+            block_uuid=self.block.uuid,
+            variable_uuid=CHILD_DATA_VARIABLE_UUID,
+            data_type=VariableAggregateDataType.STATISTICS,
+            partition=self.partition,
+        )
+        return self._summary_information
+
+    @property
+    def part_uuids(self) -> Optional[Sequence]:
+        return get_part_uuids(self.variable)
+
+    @property
+    def num_rows_from_parquet_metadata(self) -> Optional[int]:
+        metadata = read_metadata(self.variable.variable_path)
+        if metadata:
+            num_rows = metadata.get('num_rows') or 0
+            if num_rows is not None and str(num_rows).isdigit():
+                return int(str(num_rows))
+
+        return 0
+
+    def item_count(self, downstream_block: Optional[Any] = None) -> int:
+        """
+        Try the following methods to calculate the item count by starting with the
+        lowest resource consumption method first.
+
+        1. Read the original_row_count from:
+            block_uuid/
+                output_0/
+                    statistics.json
+        2. Counting the number of parts within an output variable directory:
+            block_uuid/
+                output_0/
+                    0/
+                        filename.extension
+        3. Parquet metadata: read_metadata('source_directory')['num_rows']
+            block_uuid/
+                output_0/
+                    chunks/
+                        chunk=0/
+                            filename.parquet
+                        filename.parquet
+        4. Read the file into memory (worst case)
+        """
+
+        if (
+            self.summary_information is not None
+            and self.summary_information.statistics is not None
+            and self.summary_information.statistics.original_row_count is not None
+        ):
+            return self.summary_information.statistics.original_row_count
+
+        if self.part_uuids is not None:
+            return len(self.part_uuids)
+
+        if self.num_rows_from_parquet_metadata is not None:
+            return self.num_rows_from_parquet_metadata
+
+        if self.output is not None and isinstance(self.output, Iterable):
+            return sum(1 for _ in self.output)
+
+        return 0
+
+
+class DynamicChildItemCounter(DynamicItemCounter):
+    @property
+    def part_uuids(self) -> Optional[Sequence]:
+        return dynamic_block_index_paths(self.variable)
+
+    def item_count(self, downstream_block: Optional[Any] = None) -> int:
+        """
+        1. Similar to counting the number of parts,
+        count the number of dynamic child block directories
+            block_uuid/
+                0/
+                    output_0/
+                        filename.extension
+        2. If dynamic child block reduces output, the count is 1
+        """
+
+        if (
+            (not self.block.is_dynamic_v2 and should_reduce_output(self.block))
+            or (
+                downstream_block is not None
+                and (
+                    self.block.should_reduce_output_for_downstream_block(downstream_block)
+                    or downstream_block.should_reduce_output_from_upstream_block(self.block)
+                )
+            )
+        ) or self.block.should_reduce_output:
+            return 1
+
+        if self.part_uuids is not None:
+            return len(self.part_uuids)
+
+        if self.output is not None and isinstance(self.output, Iterable):
+            return sum(1 for _ in self.output)
+
+        return 0
+
+
+class DynamicDuoItemCounter(DynamicItemCounter):
+    @property
+    def summary_information(self) -> Optional[VariableAggregateCache]:
+        if self._summary_information is not None:
+            return self._summary_information
+
+        aggregate_summary_info_for_all_variables(
+            self.block.variable_manager,
+            pipeline_uuid=self.block.pipeline.uuid,
+            block_uuid=self.block.uuid,
+            partition=self.partition,
+        )
+        self._summary_information = get_aggregate_summary_info(
+            self.block.variable_manager,
+            self.block.pipeline.uuid,
+            self.block.uuid,
+            self.variable_uuid,
+            VariableAggregateDataType.STATISTICS,
+            group_type=VariableAggregateSummaryGroupType.DYNAMIC,
+            partition=self.partition,
+        )
+
+        return self._summary_information
+
+    @property
+    def indexes(self) -> List[int]:
+        dynamic_child_counter = DynamicChildItemCounter(
+            self.block,
+            partition=self.partition,
+            variable_uuid=self.variable_uuid,
+        )
+        return [i for i in range(dynamic_child_counter.item_count())]
+
+    def calculate_item_count(self) -> int:
+        return sum([
+            DynamicBlockItemCounter(
+                self.block,
+                dynamic_block_index=dynamic_block_index,
+                partition=self.partition,
+                variable_uuid=self.variable_uuid,
+            ).item_count()
+            for dynamic_block_index in self.indexes
+        ])
+
+    def item_count(self, downstream_block: Optional[Any] = None) -> int:
+        """
+        If block is dynamic and a dynamic child,
+        the count is the sum of dynamic_output_item_count() across all dynamic child blocks:
+            0/
+                -> [1, 2, 3]
+            1/
+                -> [4, 5]
+            2/
+                -> [6, 7, 8]
+        Count: 8
+
+        After a block completes, the statistics.json must be created, aggregated, and mapped:
+
+        block_uuid/
+            0/
+                output_0/
+                    0/
+                        data.parquet
+                        statistics.json
+                    1/
+                output_1/
+            1/
+                output_0/
+                    data.json
+                    statistics.json
+        """
+        if (
+            self.summary_information
+            and self.summary_information.dynamic
+            and self.summary_information.dynamic.statistics
+        ):
+            return sum(
+                sum(i.original_row_count for i in arr)
+                for arr in self.summary_information.dynamic.statistics
+            )
+
+        return self.calculate_item_count()

mage_ai/data_preparation/models/block/dynamic/data.py

@@ -0,0 +1,16 @@
+from typing import List
+
+
+def calculate_dynamic_index_data_index(
+    dynamic_block_index: int,
+    upstream_index: int,
+    child_data_count: int,
+    item_counts: List[int],
+):
+    # The step size for each upstream index
+    step_size = 1
+    for i in range(upstream_index + 1, len(item_counts)):
+        step_size *= item_counts[i]
+
+    # Calculate the index with modulo to ensure it fits within child_data_count
+    return (dynamic_block_index // step_size) % child_data_count