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

mage_ai/data/models/manager.py

@@ -0,0 +1,168 @@
+from typing import Any, Dict, List, Optional, Union
+
+import pandas as pd
+import polars as pl
+
+from mage_ai.data.constants import (
+    InputDataType,
+    OutputData,
+    RecordBatchGenerator,
+    ScanBatchDatasetResult,
+)
+from mage_ai.data.models.base import BaseData
+from mage_ai.data.models.reader import Reader
+from mage_ai.data.models.writer import Writer
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.data_preparation.models.utils import infer_variable_type
+
+
+class DataManager(BaseData):
+    def __init__(
+        self,
+        input_data_types: Optional[List[InputDataType]],
+        read_batch_settings: Optional[BatchSettings] = None,
+        read_chunks: Optional[List] = None,
+        write_batch_settings: Optional[BatchSettings] = None,
+        write_chunks: Optional[List] = None,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.input_data_types = input_data_types or [InputDataType.DEFAULT]
+        self.read_batch_settings = read_batch_settings
+        self.read_chunks = read_chunks
+        self.write_batch_settings = write_batch_settings
+        self.write_chunks = write_chunks
+        self._reader = None
+        self._writer = None
+
+    @property
+    def batch_type(self) -> bool:
+        return InputDataType.BATCH in (self.input_data_types or [])
+
+    @property
+    def default_type(self) -> bool:
+        return not self.batch_type and not self.generator_type and not self.reader_type
+
+    @property
+    def generator_type(self) -> bool:
+        return InputDataType.GENERATOR in (self.input_data_types or [])
+
+    @property
+    def reader_type(self) -> bool:
+        return InputDataType.READER in (self.input_data_types or [])
+
+    @property
+    def reader(self) -> Reader:
+        if not self._reader:
+            self._reader = Reader(
+                batch_settings=self.read_batch_settings,
+                chunks=self.read_chunks,
+                storage=self.storage,
+                variable_dir_path=self.variable_dir_path,
+                variable_path=self.variable_path,
+                variable_type=self.variable_type,
+                variable_types=self.variable_types,
+                variables_dir=self.variables_dir,
+            )
+        return self._reader
+
+    @property
+    def writer(self) -> Writer:
+        if not self._writer:
+            self._writer = Writer(
+                batch_settings=self.write_batch_settings,
+                chunks=self.write_chunks,
+                storage=self.storage,
+                variable_dir_path=self.variable_dir_path,
+                variable_path=self.variable_path,
+                variable_type=self.variable_type,
+                variable_types=self.variable_types,
+                variables_dir=self.variables_dir,
+            )
+        return self._writer
+
+    def read_sync(
+        self,
+        limit_parts: Optional[int] = None,
+        part: Optional[int] = None,
+        sample: bool = False,
+        sample_count: Optional[int] = None,
+    ) -> Optional[Union[OutputData, ScanBatchDatasetResult, RecordBatchGenerator]]:
+        def __process_batch(batch: ScanBatchDatasetResult):
+            if batch is not None:
+                if isinstance(batch, (pd.DataFrame, pd.Series, pl.DataFrame, pl.Series)):
+                    return batch
+                return batch.deserialize()
+
+        def __process_generator(generator_batch):
+            batches = [__process_batch(batch) for batch in generator_batch if batch is not None]
+
+            if self.batch_type:
+                return batches
+
+            if len(batches) >= 1:
+                item = batches[0]
+                if isinstance(item, pl.DataFrame):
+                    return pl.concat(batches)
+                elif isinstance(item, pd.DataFrame):
+                    return pd.concat(batches)
+
+            return batches
+
+        generator = self.reader.read_sync(
+            limit_parts=limit_parts,
+            part=part,
+            sample=sample,
+            sample_count=sample_count,
+        )
+
+        if sample:
+            if self.number_of_outputs == 1:
+                return __process_batch(generator)
+            return [__process_batch(batch) for batch in generator]
+        elif self.reader_type:
+            return self.reader
+        elif self.generator_type:
+            return generator
+
+        if self.number_of_outputs == 1:
+            return __process_generator(generator)
+
+        return [__process_generator(gen) for gen in generator]
+
+    async def read_async(
+        self,
+        limit_parts: Optional[int] = None,
+        sample: bool = False,
+        sample_count: Optional[int] = None,
+    ) -> Optional[Union[Any, str]]:
+        return self.read_sync(limit_parts=limit_parts, sample=sample, sample_count=sample_count)
+
+    async def write_async(
+        self,
+        data: Any,
+        chunk_size: Optional[int] = None,
+    ) -> Optional[Dict[str, int]]:
+        self.__prepare(data, self.writer)
+        return await self.writer.write_async(data)
+
+    def write_sync(
+        self,
+        data: Any,
+        chunk_size: Optional[int] = None,
+    ) -> Optional[Dict[str, int]]:
+        self.__prepare(data, self.writer)
+        return self.writer.write_sync(data)
+
+    def readable(self) -> bool:
+        return self.reader.supported()
+
+    def writeable(self, data: Optional[Any] = None) -> bool:
+        return self.writer.supported(data=data)
+
+    def __prepare(self, data: Any, base_data: BaseData) -> None:
+        if self.variable_type is None:
+            self.variable_type, _ = infer_variable_type(data)
+            if self.variable_type is not None:
+                base_data.variable_type = self.variable_type

mage_ai/data/models/pyarrow/record_batch.py

@@ -0,0 +1,55 @@
+import math
+from typing import Optional
+
+import pyarrow as pa
+import pyarrow.dataset as ds
+from pandas.io.formats.style_render import Union
+
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.models.pyarrow.shared import Base
+from mage_ai.data.models.pyarrow.table import Table
+from mage_ai.data.tabular.constants import DEFAULT_BATCH_ITEMS_VALUE
+from mage_ai.data.tabular.utils import DeserializedBatch
+
+
+class Batch(Base):
+    def __init__(self, target: Union[pa.RecordBatch, ds.TaggedRecordBatch], **kwargs):
+        super().__init__(target, **kwargs)
+
+    def generator(
+        self, batch_size: Optional[int] = None, deserialize_on_consumption: Optional[bool] = False
+    ):
+        batch_size = batch_size or DEFAULT_BATCH_ITEMS_VALUE
+        table = pa.Table.from_batches([self.record_batch])
+        num_rows = table.num_rows
+
+        def __load(
+            index: int,
+            batch_size=batch_size,
+            deserialize_on_consumption=deserialize_on_consumption,
+            object_metadata=self.object_metadata,
+            table=table,
+        ) -> Union[DeserializedBatch, pa.Table]:
+            table_part = Table(
+                table.slice(offset=index * batch_size, length=batch_size),
+            )
+            return table_part.deserialize() if deserialize_on_consumption else table_part
+
+        def __measure(_index: int, batch_size=batch_size, num_rows=num_rows) -> int:
+            return math.ceil(num_rows / batch_size)
+
+        return DataGenerator(load_data=__load, measure_data=__measure)
+
+    @property
+    def record_batch(self) -> pa.RecordBatch:
+        return self.target
+
+
+class RecordBatch(Batch):
+    pass
+
+
+class TaggedRecordBatch(Batch):
+    @property
+    def record_batch(self) -> pa.RecordBatch:
+        return self.target.record_batch

mage_ai/data/models/pyarrow/shared.py

@@ -0,0 +1,21 @@
+from typing import Dict, Optional, Union
+
+import pyarrow as pa
+import pyarrow.dataset as ds
+
+from mage_ai.data.tabular.utils import DeserializedBatch, deserialize_batch
+from mage_ai.shared.models import Delegator
+
+
+class Base(Delegator):
+    def __init__(
+        self,
+        target: Union[pa.RecordBatch, ds.TaggedRecordBatch, pa.Table],
+        object_metadata: Optional[Dict[str, str]] = None,
+    ):
+        self.target = target
+        self.delegate = Delegator(self.target)
+        self.object_metadata = object_metadata
+
+    def deserialize(self) -> DeserializedBatch:
+        return deserialize_batch(self.target, object_metadata=self.object_metadata)

mage_ai/data/models/pyarrow/table.py

@@ -0,0 +1,8 @@
+import pyarrow as pa
+
+from mage_ai.data.models.pyarrow.shared import Base
+
+
+class Table(Base):
+    def __init__(self, target: pa.Table, **kwargs):
+        super().__init__(target, **kwargs)

mage_ai/data/models/reader.py

@@ -0,0 +1,103 @@
+from typing import List, Optional, Union
+
+from mage_ai.data.constants import (
+    OutputData,
+    RecordBatchGenerator,
+    ScanBatchDatasetResult,
+)
+from mage_ai.data.models.base import BaseData
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.tabular.reader import (
+    read_metadata,
+    sample_batch_datasets,
+    scan_batch_datasets_generator,
+)
+
+
+class Reader(BaseData):
+    """
+    - all_upstream_blocks_executed: false
+      color: null
+      configuration:
+        file_source:
+          path: memory_upgrade_v2/transformers/reader.py
+        variables:
+          upstream:
+            someblock:
+              batch_settings:
+                count:
+                  maximum: 1
+                items:
+                  minimum: 100000
+              input_data_types:
+              - generator
+    """
+
+    def read_sync(
+        self,
+        columns: Optional[List[str]] = None,
+        deserialize: Optional[bool] = False,
+        limit: Optional[int] = None,
+        limit_parts: Optional[int] = None,
+        offset: Optional[int] = None,
+        part: Optional[int] = None,
+        sample: bool = False,
+        sample_count: Optional[int] = None,
+    ) -> Optional[
+        Union[
+            OutputData,
+            ScanBatchDatasetResult,
+            RecordBatchGenerator,
+        ]
+    ]:
+        if not self.is_dataframe():
+            return
+
+        def __process_source(
+            source,
+            chunks=self.chunks,
+            columns=columns,
+            deserialize=deserialize,
+            limit=limit,
+            offset=offset,
+            part=part,
+            sample_count=sample_count,
+            settings=self.batch_settings,
+        ):
+            if part is not None and settings is not None and limit is None and offset is None:
+                metadata = read_metadata(source)
+                if metadata:
+                    num_rows = metadata.get('num_rows')
+                    if num_rows:
+                        batch_size = settings.items.maximum or settings.items.minimum
+                        if batch_size:
+                            offset = part * batch_size
+                            limit = batch_size
+
+            if sample and sample_count:
+                return sample_batch_datasets(
+                    source,
+                    chunks=chunks,
+                    columns=columns,
+                    deserialize=deserialize,
+                    limit=limit,
+                    offset=offset,
+                    sample_count=sample_count,
+                    settings=settings,
+                )
+            return scan_batch_datasets_generator(
+                source,
+                chunks=chunks,
+                columns=columns,
+                deserialize=deserialize,
+                limit=limit,
+                offset=offset,
+                settings=settings,
+            )
+
+        if self.number_of_outputs >= 2:
+            sources = self.data_source
+            if limit_parts is not None:
+                sources = sources[:limit_parts]
+            return DataGenerator([__process_source(source) for source in sources])
+        return __process_source(self.data_source)

mage_ai/data/models/utils.py

@@ -0,0 +1,59 @@
+from typing import Any, Optional
+
+from mage_ai.data.constants import SUPPORTED_VARIABLE_TYPES
+from mage_ai.data_preparation.models.utils import infer_variable_type
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.shared.environments import is_debug
+
+
+def variable_type_supported(
+    variable_type: Optional[VariableType] = None, data: Optional[Any] = None
+) -> bool:
+    from mage_ai.data_preparation.models.utils import is_user_defined_complex
+
+    # Need to import here to mock in tests
+    from mage_ai.settings.server import (
+        MEMORY_MANAGER_PANDAS_V2,
+        MEMORY_MANAGER_POLARS_V2,
+        MEMORY_MANAGER_V2,
+    )
+
+    if not MEMORY_MANAGER_V2:
+        if is_debug():
+            print('[WARNING] Memory manager V2 is not enabled.')
+        return False
+
+    basic_iterable = False
+    if data is not None and variable_type is None:
+        variable_type, basic_iterable = infer_variable_type(data)
+
+    if variable_type not in SUPPORTED_VARIABLE_TYPES:
+        if is_debug():
+            print(f'[WARNING] Variable type {variable_type} is not supported.')
+        return False
+
+    if (
+        variable_type
+        in [
+            VariableType.POLARS_DATAFRAME,
+            VariableType.SERIES_POLARS,
+        ]
+        and not MEMORY_MANAGER_POLARS_V2
+    ):
+        if is_debug():
+            print(f'[WARNING] Variable type {variable_type} is not supported.')
+        return False
+
+    if (
+        variable_type
+        in [
+            VariableType.DATAFRAME,
+            VariableType.SERIES_PANDAS,
+        ]
+        and not MEMORY_MANAGER_PANDAS_V2
+    ):
+        if is_debug():
+            print(f'[WARNING] Variable type {variable_type} is not supported.')
+        return False
+
+    return data is None or basic_iterable or is_user_defined_complex(data)

mage_ai/data/models/writer.py

@@ -0,0 +1,91 @@
+import uuid
+from typing import Any, Dict, Optional
+
+from mage_ai.data.models.base import BaseData
+from mage_ai.data.tabular.writer import to_parquet_async, to_parquet_sync
+from mage_ai.io.base import ExportWritePolicy
+from mage_ai.shared.files import (
+    makedirs_async,
+    makedirs_sync,
+    safe_delete_dir_async,
+    safe_delete_dir_sync,
+)
+
+
+class Writer(BaseData):
+    """
+    - all_upstream_blocks_executed: false
+      color: null
+      configuration:
+        file_source:
+          path: memory_upgrade_v2/data_exporters/bold_portal.py
+        variables:
+          write:
+            batch_settings:
+              items:
+                maximum: 170000
+            chunks:
+            - key
+    """
+
+    def write_sync(self, data: Any, replace: Optional[bool] = None) -> Optional[Dict[str, int]]:
+        makedirs_sync(self.data_source_directory_path)
+
+        mode = self.__write_mode
+        replace = replace if replace is not None else ExportWritePolicy.REPLACE == mode
+        if replace:
+            safe_delete_dir_sync(self.data_source_directory_path)
+
+        if self.is_dataframe():
+            return to_parquet_sync(
+                self.data_source_directory_path,
+                df=data,
+                basename_template=self.__basename_template,
+                existing_data_behavior=self.__existing_data_behavior,
+                settings=self.batch_settings,
+                partition_cols=[str(key) for key in (self.chunks or [])],
+            )
+
+    async def write_async(self, data: Any, replace: bool = True) -> Optional[Dict[str, int]]:
+        await makedirs_async(self.data_source_directory_path)
+
+        mode = self.__write_mode
+        replace = replace if replace is not None else ExportWritePolicy.REPLACE == mode
+        if replace:
+            await safe_delete_dir_async(self.data_source_directory_path)
+
+        if self.is_dataframe():
+            return await to_parquet_async(
+                self.data_source_directory_path,
+                df=data,
+                basename_template=self.__basename_template,
+                existing_data_behavior=self.__existing_data_behavior,
+                settings=self.batch_settings,
+                partition_cols=[str(key) for key in (self.chunks or [])],
+            )
+
+    @property
+    def __basename_template(self) -> str:
+        mode = self.__write_mode
+        filename = uuid.uuid5(uuid.NAMESPACE_OID, self.uuid)
+        if ExportWritePolicy.APPEND == mode:
+            filename = str(uuid.uuid4())
+        return '{i}-' + f'{filename}.parquet'
+
+    @property
+    def __existing_data_behavior(self) -> str:
+        mode = self.__write_mode
+        return (
+            'error'
+            if ExportWritePolicy.FAIL == mode
+            else 'overwrite_or_ignore'
+            if ExportWritePolicy.APPEND == mode
+            else 'delete_matching'
+        )
+
+    @property
+    def __write_mode(self) -> ExportWritePolicy:
+        if self.batch_settings and self.batch_settings.mode is not None:
+            return ExportWritePolicy.from_value(self.batch_settings.mode)
+
+        return ExportWritePolicy.REPLACE

mage_ai/data/tabular/constants.py

@@ -0,0 +1,23 @@
+from enum import Enum
+
+# Cannot use any names with __ in the beginning or end or else PyArrow won’t read those files
+# when loading data into a PyArrow Table.
+COLUMN_CHUNK = 'mage_chunk'
+DEFAULT_BATCH_ITEMS_VALUE = 1_000_000
+DEFAULT_BATCH_BYTE_VALUE = 1024 * 1024 * 100  # 100 MB
+DEFAULT_BATCH_COUNT_VALUE = 1_000
+
+
+class FilterComparison(Enum):
+    EQUAL = '=='
+    GREATER_THAN = '>'
+    GREATER_THAN_OR_EQUAL = '>='
+    LESS_THAN = '<'
+    LESS_THAN_OR_EQUAL = '<='
+    NOT_EQUAL = '!='
+
+
+class BatchStrategy(Enum):
+    BYTES = 'bytes'  # Byte size per batch
+    COUNT = 'count'  # Number of batches
+    ITEMS = 'items'  # Rows per batch

mage_ai/data/tabular/mocks.py

@@ -0,0 +1,19 @@
+from typing import Union
+
+import numpy as np
+import pandas as pd
+import polars as pl
+
+DEFAULT_ROWS = 1_000_000
+DEFAULT_COLUMNS = 10
+
+
+def create_dataframe(
+    n_rows: int = DEFAULT_ROWS, n_cols: int = DEFAULT_COLUMNS, use_pandas: bool = False
+) -> Union[pd.DataFrame, pl.DataFrame]:
+    data = {f'col_{i}': np.random.rand(n_rows) for i in range(n_cols)}
+    if use_pandas:
+        df = pd.DataFrame(data)
+    else:
+        df = pl.DataFrame(data)
+    return df

mage_ai/data/tabular/models.py

@@ -0,0 +1,126 @@
+from dataclasses import dataclass, field
+from functools import reduce
+from math import ceil
+from typing import Any, Dict, List, Optional, Union
+
+from mage_ai.data.constants import ReadModeType
+from mage_ai.data.tabular.constants import (
+    DEFAULT_BATCH_BYTE_VALUE,
+    DEFAULT_BATCH_COUNT_VALUE,
+    DEFAULT_BATCH_ITEMS_VALUE,
+    BatchStrategy,
+)
+from mage_ai.io.base import ExportWritePolicy
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class Settings(BaseDataClass):
+    maximum: int = 0
+    minimum: int = 0
+
+
+@dataclass
+class BatchSettings(BaseDataClass):
+    """
+    Batches with less items uses less memory
+        Batch by 10K rows:
+            Time elapsed: 64.632 ms
+            Memory start: 0.192 gb
+            Memory final: 0.795 gb
+            Memory churn: 16.318 gb
+            Memory peak : 13.953 gb
+        Batch by 100K rows:
+            Time elapsed: 64.961 ms
+            Memory start: 0.194 gb
+            Memory final: 1.362 gb
+            Memory churn: 17.684 gb
+            Memory peak : 14.831 gb
+        Batch by 1M rows:
+            Time elapsed: 72.813 ms
+            Memory start: 0.198 gb
+            Memory final: 0.935 gb
+            Memory churn: 18.075 gb
+            Memory peak : 14.756 gb
+
+    Batches with less bytes uses less memory
+        Batch by 100MB:
+            Time elapsed: 73.180 ms
+            Memory start: 0.874 gb
+            Memory final: 1.028 gb
+            Memory churn: 16.988 gb
+            Memory peak : 14.823 gb
+        Batch by 1GB:
+            Time elapsed: 63.710 ms
+            Memory start: 0.890 gb
+            Memory final: 0.948 gb
+            Memory churn: 17.181 gb
+            Memory peak : 14.559 gb
+        Batch by 2GB:
+            Time elapsed: 67.282 ms
+            Memory start: 0.755 gb
+            Memory final: 1.059 gb
+            Memory churn: 17.256 gb
+            Memory peak : 15.057 gb
+    """
+
+    count: Settings = field(default_factory=Settings)
+    items: Settings = field(default_factory=Settings)
+    mode: Optional[Union[ReadModeType, ExportWritePolicy]] = None
+    size: Settings = field(default_factory=Settings)
+
+    def __post_init__(self):
+        self.serialize_attribute_class('count', Settings)
+        self.serialize_attribute_class('items', Settings)
+        self.serialize_attribute_class('size', Settings)
+
+        if self.mode:
+            for mode in (ReadModeType, ExportWritePolicy):
+                if mode.has_value(self.mode):
+                    self.mode = mode.from_value(self.mode)
+                    break
+
+    @property
+    def strategy(self) -> BatchStrategy:
+        if self.count is not None:
+            return BatchStrategy.COUNT
+        elif self.items is not None:
+            return BatchStrategy.ITEMS
+        elif self.size is not None:
+            return BatchStrategy.BYTES
+        else:
+            return BatchStrategy.ITEMS
+
+    @property
+    def batch_value(self) -> int:
+        if BatchStrategy.ITEMS == self.strategy and self.items is not None:
+            return self.items.maximum or DEFAULT_BATCH_ITEMS_VALUE
+        elif BatchStrategy.BYTES == self.strategy and self.size is not None:
+            return self.size.maximum or DEFAULT_BATCH_BYTE_VALUE
+        elif BatchStrategy.COUNT == self.strategy and self.count is not None:
+            return self.count.maximum or DEFAULT_BATCH_COUNT_VALUE
+        else:
+            return DEFAULT_BATCH_ITEMS_VALUE
+
+    def batch_size(self, metadatas: List[Dict[str, Any]]) -> int:
+        total_row_count, total_byte_size = reduce(
+            lambda acc, metadata: (
+                acc[0] + metadata['num_rows'],
+                acc[1] + metadata['total_byte_size'],
+            ),
+            metadatas,
+            (0, 0),
+        )
+        if BatchStrategy.ITEMS == self.strategy:
+            batch_size = self.batch_value or DEFAULT_BATCH_ITEMS_VALUE
+        elif BatchStrategy.BYTES == self.strategy:
+            batch_value = self.batch_value or DEFAULT_BATCH_BYTE_VALUE
+            estimated_bytes_per_row = total_byte_size / total_row_count
+            batch_size = max(1, int(batch_value / estimated_bytes_per_row))
+        elif BatchStrategy.COUNT == self.strategy:
+            batch_value = self.batch_value or DEFAULT_BATCH_COUNT_VALUE
+            batch_size = ceil(total_row_count / batch_value)
+        else:
+            batch_size = DEFAULT_BATCH_ITEMS_VALUE
+
+        return batch_size