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

mage_ai/data_preparation/models/variables/cache.py

@@ -0,0 +1,149 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.shared.models import BaseDataClass
+from mage_ai.system.models import ResourceUsage
+
+
+@dataclass
+class VariableTypeInformation(BaseDataClass):
+    type: VariableType
+
+    def __post_init__(self):
+        self.serialize_attribute_enum('type', VariableType)
+
+
+@dataclass
+class StatisticsInformation(BaseDataClass):
+    original_column_count: Optional[int] = None
+    original_row_count: Optional[int] = None
+
+
+InformationData = Optional[
+    Union[
+        Dict[str, Any],
+        ResourceUsage,
+        StatisticsInformation,
+        VariableTypeInformation,
+    ]
+]
+AggregateInformationData = List[InformationData]
+
+
+@dataclass
+class AggregateInformation(BaseDataClass):
+    insights: Optional[
+        Union[
+            List[Dict[str, Any]],
+            List[List[Dict[str, Any]]],
+        ]
+    ] = None
+    metadata: Optional[
+        Union[
+            List[Dict[str, Any]],
+            List[List[Dict[str, Any]]],
+        ]
+    ] = None
+    resource_usage: Optional[
+        Union[
+            List[ResourceUsage],
+            List[List[ResourceUsage]],
+        ]
+    ] = None
+    sample_data: Optional[
+        Union[
+            List[Dict[str, Any]],
+            List[List[Dict[str, Any]]],
+        ]
+    ] = None
+    statistics: Optional[
+        Union[
+            List[StatisticsInformation],
+            List[List[StatisticsInformation]],
+        ]
+    ] = None
+    suggestions: Optional[
+        Union[
+            List[Dict[str, Any]],
+            List[List[Dict[str, Any]]],
+        ]
+    ] = None
+    type: Optional[
+        Union[
+            List[VariableTypeInformation],
+            List[List[VariableTypeInformation]],
+        ]
+    ] = None
+
+    def __post_init__(self):
+        if self.resource_usage and isinstance(self.resource_usage, list):
+            arr = []
+            for item in self.resource_usage:
+                if item and isinstance(item, dict):
+                    arr.append(ResourceUsage.load(**item))
+                elif isinstance(item, ResourceUsage):
+                    arr.append(ResourceUsage)
+                elif isinstance(item, list):
+                    arr_nested = []
+                    for item_nested in item:
+                        if item_nested and isinstance(item_nested, dict):
+                            arr_nested.append(ResourceUsage.load(**item_nested))
+                        elif isinstance(item_nested, ResourceUsage):
+                            arr_nested.append(ResourceUsage)
+                    arr.append(arr_nested)
+            self.resource_usage = arr
+
+        if self.statistics and isinstance(self.statistics, list):
+            arr = []
+            for item in self.statistics:
+                if item and isinstance(item, dict):
+                    arr.append(StatisticsInformation.load(**item))
+                elif isinstance(item, StatisticsInformation):
+                    arr.append(StatisticsInformation)
+                elif isinstance(item, list):
+                    arr_nested = []
+                    for item_nested in item:
+                        if item_nested and isinstance(item_nested, dict):
+                            arr_nested.append(StatisticsInformation.load(**item_nested))
+                        elif isinstance(item_nested, StatisticsInformation):
+                            arr_nested.append(StatisticsInformation)
+                    arr.append(arr_nested)
+            self.statistics = arr
+
+        if self.type and isinstance(self.type, list):
+            arr = []
+            for item in self.type:
+                if item and isinstance(item, dict):
+                    arr.append(VariableTypeInformation.load(**item))
+                elif isinstance(item, VariableTypeInformation):
+                    arr.append(VariableTypeInformation)
+                elif isinstance(item, list):
+                    arr_nested = []
+                    for item_nested in item:
+                        if item_nested and isinstance(item_nested, dict):
+                            arr_nested.append(VariableTypeInformation.load(**item_nested))
+                        elif isinstance(item_nested, VariableTypeInformation):
+                            arr_nested.append(VariableTypeInformation)
+                    arr.append(arr_nested)
+            self.type = arr
+
+
+@dataclass
+class VariableAggregateCache(BaseDataClass):
+    dynamic: Optional[AggregateInformation] = None
+    insights: Optional[Dict[str, Any]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    parts: Optional[AggregateInformation] = None
+    resource_usage: Optional[ResourceUsage] = None
+    sample_data: Optional[Dict[str, Any]] = None
+    statistics: Optional[StatisticsInformation] = None
+    suggestions: Optional[Dict[str, Any]] = None
+    type: Optional[VariableTypeInformation] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('dynamic', AggregateInformation)
+        self.serialize_attribute_class('parts', AggregateInformation)
+        self.serialize_attribute_class('resource_usage', ResourceUsage)
+        self.serialize_attribute_class('statistics', StatisticsInformation)
+        self.serialize_attribute_class('type', VariableTypeInformation)

mage_ai/data_preparation/models/variables/constants.py

@@ -0,0 +1,72 @@
+from enum import Enum
+
+CONFIG_JSON_FILE = 'config.json'
+DATAFRAME_COLUMN_TYPES_FILE = 'data_column_types.json'
+DATAFRAME_CSV_FILE = 'data.csv'
+DATAFRAME_PARQUET_FILE = 'data.parquet'
+DATAFRAME_PARQUET_SAMPLE_FILE = 'sample_data.parquet'
+JOBLIB_FILE = 'model.joblib'
+JOBLIB_OBJECT_FILE = 'object.joblib'
+JSON_FILE = 'data.json'
+MATRIX_NPZ_FILE = 'matrix.npz'
+MATRIX_SAMPLE_NPZ_FILE = 'matrix.npz'
+MEDIA_IMAGE_VISUALIZATION_FILE = 'visualization.png'
+UBJSON_MODEL_FILENAME = 'model.ubj'
+
+
+class VariableType(str, Enum):
+    CUSTOM_OBJECT = 'custom_object'
+    DATAFRAME = 'dataframe'
+    DATAFRAME_ANALYSIS = 'dataframe_analysis'
+    DICTIONARY_COMPLEX = 'dictionary_complex'
+    GEO_DATAFRAME = 'geo_dataframe'
+    ITERABLE = 'iterable'
+    LIST_COMPLEX = 'list_complex'
+    MATRIX_SPARSE = 'matrix_sparse'
+    MODEL_SKLEARN = 'model_sklearn'
+    MODEL_XGBOOST = 'model_xgboost'
+    POLARS_DATAFRAME = 'polars_dataframe'
+    SERIES_PANDAS = 'series_pandas'
+    SERIES_POLARS = 'series_polars'
+    SPARK_DATAFRAME = 'spark_dataframe'
+
+
+class VariableAggregateDataType(str, Enum):
+    INSIGHTS = 'insights'
+    METADATA = 'metadata'
+    RESOURCE_USAGE = 'resource_usage'
+    SAMPLE_DATA = 'sample_data'
+    STATISTICS = 'statistics'
+    SUGGESTIONS = 'suggestions'
+    TYPE = 'type'
+
+
+class VariableAggregateDataTypeFilename(str, Enum):
+    INSIGHTS = f'{VariableAggregateDataType.INSIGHTS}.json'
+    METADATA = f'{VariableAggregateDataType.METADATA}.json'
+    RESOURCE_USAGE = f'{VariableAggregateDataType.RESOURCE_USAGE}.json'
+    SAMPLE_DATA = f'{VariableAggregateDataType.SAMPLE_DATA}.json'
+    STATISTICS = f'{VariableAggregateDataType.STATISTICS}.json'
+    SUGGESTIONS = f'{VariableAggregateDataType.SUGGESTIONS}.json'
+    TYPE = f'{VariableAggregateDataType.TYPE}.json'
+
+
+DATA_TYPE_FILENAME = {
+    VariableAggregateDataType.INSIGHTS: VariableAggregateDataTypeFilename.INSIGHTS,
+    VariableAggregateDataType.METADATA: VariableAggregateDataTypeFilename.METADATA,
+    VariableAggregateDataType.RESOURCE_USAGE: VariableAggregateDataTypeFilename.RESOURCE_USAGE,
+    VariableAggregateDataType.SAMPLE_DATA: VariableAggregateDataTypeFilename.SAMPLE_DATA,
+    VariableAggregateDataType.STATISTICS: VariableAggregateDataTypeFilename.STATISTICS,
+    VariableAggregateDataType.SUGGESTIONS: VariableAggregateDataTypeFilename.SUGGESTIONS,
+    VariableAggregateDataType.TYPE: VariableAggregateDataTypeFilename.TYPE,
+}
+
+
+RESOURCE_USAGE_FILE = VariableAggregateDataTypeFilename.RESOURCE_USAGE.value
+JSON_SAMPLE_FILE = VariableAggregateDataTypeFilename.SAMPLE_DATA.value
+METADATA_FILE = VariableAggregateDataTypeFilename.TYPE.value
+
+
+class VariableAggregateSummaryGroupType(str, Enum):
+    DYNAMIC = 'dynamic'
+    PARTS = 'parts'

mage_ai/data_preparation/models/variables/summarizer.py

@@ -0,0 +1,336 @@
+import os
+from typing import List, Optional, Tuple
+
+from mage_ai.data_preparation.models.variables.cache import (
+    AggregateInformation,
+    VariableAggregateCache,
+    VariableTypeInformation,
+)
+from mage_ai.data_preparation.models.variables.constants import (
+    DATA_TYPE_FILENAME,
+    VariableAggregateDataType,
+    VariableAggregateDataTypeFilename,
+    VariableAggregateSummaryGroupType,
+)
+from mage_ai.data_preparation.models.variables.utils import is_output_variable
+from mage_ai.shared.utils import clean_name
+
+
+def aggregate_summary_info_for_all_variables(
+    variable_manager,
+    pipeline_uuid: str,
+    block_uuid: str,  # Make sure this is the base block_uuid, not with the dynamic block index
+    partition: Optional[str] = None,
+):
+    block_uuid = clean_name(block_uuid)
+    variable_uuids = variable_manager.get_variables_by_block(
+        pipeline_uuid,
+        block_uuid,
+        partition=partition,
+    )
+
+    if variable_uuids and all(not is_output_variable(v) for v in variable_uuids):
+        """
+        Dynamic child blocks have this directory structure:
+            [block_uuid]/
+                [dynamic_block_index]/
+                    [variable_uuid]/
+        """
+        variable_uuids = variable_manager.get_variables_by_block(
+            pipeline_uuid,
+            os.path.join(block_uuid, str(variable_uuids[0])),
+            clean_block_uuid=False,
+            partition=partition,
+        )
+
+    for variable_uuid in variable_uuids:
+        variable = variable_manager.get_variable_object(
+            pipeline_uuid,
+            block_uuid,
+            variable_uuid,
+            partition=partition,
+            skip_check_variable_type=True,
+        )
+        __aggregate_summary_info(variable)
+
+
+def __aggregate_summary_info(variable):
+    """
+    Only consolidate summary files if at least 1 of the following criteria are met:
+        - Dynamic child block
+        - Variable has parts (e.g. output_0/[part_uuid])
+    Consolidate summary files from all subfolders (except data.json):
+        - insights.json
+        - metadata.json
+        - resources.json
+        - sample_data.json
+        - statistics.json
+        - suggestions.json
+        - type.json
+
+    Directory structure
+    [block_uuid]/                  -> variable.variable_dir_path
+        [variable_uuid: output_0]/ -> variable.variable_path
+            0/
+                type.json
+            1/
+                type.json
+            type.json
+        [dynamic_block_index]/
+            [variable_uuid: output_0]/
+                0/
+                    type.json
+                1/
+                    type.json
+                type.json
+    """
+
+    filenames = [key.value for key in VariableAggregateDataTypeFilename]
+
+    """
+    [block_uuid]/                  -> variable.variable_dir_path
+        [variable_uuid: output_0]/ -> variable.variable_path
+            0/
+                type.json
+            1/
+                type.json
+            type.json
+    """
+    mapping_parts = {}
+    if variable.part_uuids:
+        for part_uuid in variable.part_uuids:  # part_uuid: 0/
+            for filename in filenames:  # filename: type.json
+                if not mapping_parts.get(filename):
+                    mapping_parts[filename] = []
+
+                file_path = os.path.join(variable.variable_path, part_uuid, filename)
+                if variable.storage.path_exists(file_path):
+                    mapping_parts[filename].append(
+                        variable.storage.read_json_file(
+                            file_path,
+                            default_value={},
+                            raise_exception=False,
+                        ),
+                    )
+                else:
+                    mapping_parts[filename].append({})
+
+    """
+    [block_uuid]/                      -> variable.variable_dir_path
+        [dynamic_block_index]/         -> [read all of them]
+            [variable_uuid: output_0]/ -> variable.uuid
+                0/
+                    type.json
+                1/
+                    type.json
+                type.json
+    """
+
+    for dirname in VariableAggregateSummaryGroupType:
+        path = os.path.join(variable.variable_dir_path, dirname.value)
+        if variable.storage.path_exists(path):
+            variable.storage.remove_dir(path)
+
+    mapping_dynamic_blocks = {}
+    # dynamic_block_index: 0
+    # index_path:  [block_uuid]/[dynamic_block_index: 0]/
+    # output_path: [block_uuid]/[dynamic_block_index: 0]/[variable_uuid: output_0]/
+    for _dynamic_block_index, _index_path, output_path in dynamic_block_index_paths(variable):
+        mapping_dynamic_children = {}
+        for filename in filenames:  # filename: type.json
+            if not mapping_dynamic_children.get(filename):
+                mapping_dynamic_children[filename] = []
+
+            part_uuids = __get_part_uuids(variable, output_path)
+            if part_uuids and len(part_uuids) >= 1:
+                arr = []
+                for part_uuid in part_uuids:
+                    part_path = os.path.join(output_path, part_uuid, filename)
+                    if variable.storage.path_exists(part_path):
+                        arr.append(
+                            variable.storage.read_json_file(
+                                part_path,
+                                default_value={},
+                                raise_exception=False,
+                            ),
+                        )
+                mapping_dynamic_children[filename].append(arr)
+            else:
+                file_path = os.path.join(output_path, filename)
+                if variable.storage.path_exists(file_path):
+                    mapping_dynamic_children[filename].append(
+                        variable.storage.read_json_file(
+                            file_path,
+                            default_value={},
+                            raise_exception=True,
+                        ),
+                    )
+                else:
+                    mapping_dynamic_children[filename].append({})
+
+        for filename, arr in mapping_dynamic_children.items():
+            if not mapping_dynamic_blocks.get(filename):
+                mapping_dynamic_blocks[filename] = []
+            mapping_dynamic_blocks[filename].append(arr)
+
+    for directory, mapping in [
+        (VariableAggregateSummaryGroupType.DYNAMIC, mapping_dynamic_blocks),
+        (VariableAggregateSummaryGroupType.PARTS, mapping_parts),
+    ]:
+        if len(mapping) == 0:
+            continue
+
+        path = os.path.join(variable.variable_path, directory)
+
+        if not variable.storage.isdir(path):
+            variable.storage.makedirs(path, exist_ok=True)
+
+        for filename, data in mapping.items():
+            variable.storage.write_json_file(os.path.join(path, filename), data)
+
+
+def get_aggregate_summary_info(
+    variable_manager,
+    pipeline_uuid: str,
+    block_uuid: str,
+    variable_uuid: str,
+    data_type: VariableAggregateDataType,
+    default_group_type: Optional[VariableAggregateSummaryGroupType] = None,
+    group_type: Optional[VariableAggregateSummaryGroupType] = None,
+    partition: Optional[str] = None,
+) -> VariableAggregateCache:
+    variable = variable_manager.get_variable_object(
+        pipeline_uuid,
+        block_uuid,
+        variable_uuid,
+        partition=partition,
+        skip_check_variable_type=True,
+    )
+    return __get_aggregate_summary_info(
+        variable, data_type, default_group_type=default_group_type, group_type=group_type
+    )
+
+
+def __get_aggregate_summary_info(
+    variable,
+    data_type: VariableAggregateDataType,
+    default_group_type: Optional[VariableAggregateSummaryGroupType] = None,
+    group_type: Optional[VariableAggregateSummaryGroupType] = None,
+) -> VariableAggregateCache:
+    """
+    Get aggregate summary information for the variable.
+
+    Args:
+        group_type (VariableAggregateSummaryGroupType): Group type of the variable.
+
+    Returns:
+        Optional[AggregateInformationData]: Aggregate summary information for the variable.
+
+    Used
+    """
+    cache = VariableAggregateCache()
+
+    path = os.path.join(
+        variable.variable_path,
+        group_type.value if group_type else '',
+        DATA_TYPE_FILENAME[data_type],
+    )
+
+    if variable.storage.path_exists(path):
+        data = variable.storage.read_json_file(
+            path,
+            default_value=None,
+            raise_exception=False,
+        )
+
+        if group_type:
+            group_info = getattr(cache, group_type.value)
+            group_info = AggregateInformation.load(
+                **(
+                    group_info
+                    if isinstance(group_info, dict)
+                    else group_info.to_dict()
+                    if group_info is not None
+                    else {}
+                )
+            )
+            group_info.update_attributes(**{
+                data_type.value: data,
+            })
+            cache.update_attributes(**{
+                group_type.value: group_info,
+            })
+        else:
+            if (
+                VariableAggregateDataType.TYPE == data_type.value
+                and data is not None
+                and isinstance(data, dict)
+            ):
+                if data.get('type'):
+                    cache.update_attributes(**{
+                        'type': VariableTypeInformation.load(type=data.get('type')),
+                    })
+
+                if default_group_type is not None and data.get('types'):
+                    var_types = data.get('types') or []
+                    group_info = getattr(cache, default_group_type.value)
+                    group_info = AggregateInformation.load(
+                        **(
+                            group_info
+                            if isinstance(group_info, dict)
+                            else group_info.to_dict()
+                            if group_info is not None
+                            else {}
+                        )
+                    )
+                    group_info.update_attributes(**{
+                        data_type.value: [
+                            VariableTypeInformation.load(type=val) for val in var_types
+                        ],
+                    })
+                    cache.update_attributes(**{
+                        default_group_type.value: group_info,
+                    })
+
+            else:
+                cache.update_attributes(**{
+                    data_type.value: data,
+                })
+    return VariableAggregateCache.load(**cache.to_dict())
+
+
+def get_part_uuids(variable) -> Optional[List[str]]:
+    if (
+        not variable.uuid
+        or not is_output_variable(variable.uuid)
+        or not variable.storage.isdir(variable.variable_path)
+    ):
+        return None
+
+    return __get_part_uuids(variable, variable.variable_path)
+
+
+def dynamic_block_index_paths(variable) -> List[Tuple[int, str, str]]:
+    if not variable.storage.isdir(variable.variable_dir_path):
+        return []
+
+    indexes = []
+    for dynamic_block_index in variable.storage.listdir(variable.variable_dir_path):
+        index_path = os.path.join(variable.variable_dir_path, dynamic_block_index)
+        if dynamic_block_index.isdigit() and variable.storage.isdir(index_path):
+            output_path = os.path.join(index_path, variable.uuid)
+            if variable.storage.isdir(output_path):
+                indexes.append((int(dynamic_block_index), index_path, output_path))
+
+    return indexes
+
+
+def __get_part_uuids(variable, path: str) -> Optional[List[str]]:
+    part_uuids = None
+    for chunk_uuid in variable.storage.listdir(path):
+        if chunk_uuid.isdigit() and variable.storage.isdir(os.path.join(path, chunk_uuid)):
+            if part_uuids is None:
+                part_uuids = []
+            part_uuids.append(chunk_uuid)
+
+    return part_uuids

mage_ai/data_preparation/models/variables/utils.py

@@ -0,0 +1,77 @@
+import re
+from typing import Any, List, Optional, Tuple, Union
+
+from mage_ai.data_preparation.models.interfaces import VariableInterface
+from mage_ai.shared.array import find
+from mage_ai.shared.strings import to_ordinal_integers
+
+
+def __sort(text) -> Tuple[int, int]:
+    number = re.findall('\\d+', text)
+    if number:
+        if len(number) == 1:
+            number.append(0)
+        number = number[:2]
+    else:
+        number = to_ordinal_integers(text)
+
+    number = [int(i) for i in number][:2]
+    return tuple(number)
+
+
+def __sort_object(variable: VariableInterface) -> Tuple[int, int]:
+    return __sort(variable.uuid)
+
+
+def sort_variable_uuids(variable_uuids: List[str]) -> List[str]:
+    return sorted(variable_uuids, key=__sort)
+
+
+def sort_variables(variables: List[VariableInterface]) -> List[VariableInterface]:
+    return sorted(variables, key=__sort_object)
+
+
+def is_output_variable(variable_uuid: str, include_df: bool = True) -> bool:
+    """
+    Checks if the given variable UUID represents an output variable.
+
+    `df` variable is only used to save the sample Spark dataframe now.
+
+    Args:
+        variable_uuid (str): The UUID of the variable.
+        include_df (bool): Whether to include the `df` variable. Only set it to True
+            when saving or fetching sample Spark dataframe.
+
+    Returns:
+        bool: True if the variable is an output variable, False otherwise.
+
+    """
+    return (include_df and variable_uuid == 'df') or variable_uuid.startswith('output')
+
+
+def variable_priority(variable_uuid: str) -> Optional[int]:
+    number = re.findall('\\d+', variable_uuid)[:1]
+    if number:
+        number = number[0]
+        if number.isdigit():
+            return int(number)
+
+
+def get_first_data_output_variable(
+    variables: List[Union[VariableInterface, Any]],
+) -> Optional[Union[VariableInterface, Any]]:
+    return find(
+        lambda variable: is_output_variable(variable.uuid)
+        and variable_priority(variable.uuid) == 0,
+        sort_variables(variables),
+    )
+
+
+def get_first_data_output_variable_uuid(
+    variables: List[str],
+) -> Optional[str]:
+    return find(
+        lambda variable_uuid: is_output_variable(variable_uuid)
+        and variable_priority(variable_uuid) == 0,
+        sort_variable_uuids(variables),
+    )