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

mage_ai/command_center/files/factory.py

@@ -23,13 +23,13 @@ class FileFactory(BaseFactory):
         files = []
 
         async def build_and_score(line: str, files=files, factory=self):
-            full_path, file_size, modified_timestamp = line.split(',')
+            full_path, file_size, modified_timestamp = line.split(",")
             filename = os.path.basename(full_path)
             modified_at = datetime.fromtimestamp(float(modified_timestamp)).isoformat()
 
             path_dict = shorten_directory(full_path)
-            extension = path_dict.get('extension')
-            directory = path_dict.get('directory')
+            extension = path_dict.get("extension")
+            directory = path_dict.get("directory")
 
             item_dict = dict(
                 item_type=ItemType.DETAIL,
@@ -55,7 +55,6 @@ class FileFactory(BaseFactory):
                     ),
                 ),
             )
-
             scored = factory.filter_score(item_dict)
             if scored:
                 files.append(scored)
@@ -64,23 +63,26 @@ class FileFactory(BaseFactory):
             now = datetime.utcnow().timestamp()
             cache = FileCache.initialize_cache_with_settings()
             lines = await cache.load() or []
-            print(f'[FileFactory] Load files: {len(lines)} - {datetime.utcnow().timestamp() - now}')
+            print(
+                f"[FileFactory] Load files: {len(lines)} - {datetime.utcnow().timestamp() - now}"
+            )
 
             now = datetime.utcnow().timestamp()
-            await asyncio.gather(
-                *[build_and_score(item_dict) for item_dict in lines]
-            )
+            await asyncio.gather(*[build_and_score(item_dict) for item_dict in lines])
             print(
-                f'[FileFactory] Search {self.search}: '
-                f'{len(files)} - {datetime.utcnow().timestamp() - now}',
+                f"[FileFactory] Search {self.search}: "
+                f"{len(files)} - {datetime.utcnow().timestamp() - now}",
             )
 
             now = datetime.utcnow().timestamp()
             files = await self.rank_items(files)
-            files = [merge_dict(
-                item_dict,
-                add_application_actions(item_dict),
-            ) for item_dict in files]
-            print(f'[FileFactory] Rank items: {datetime.utcnow().timestamp() - now}')
+            files = [
+                merge_dict(
+                    item_dict,
+                    add_application_actions(item_dict),
+                )
+                for item_dict in files
+            ]
+            print(f"[FileFactory] Rank items: {datetime.utcnow().timestamp() - now}")
 
         return files[:10] + items

mage_ai/command_center/utils.py

@@ -4,27 +4,39 @@ from typing import Dict
 
 
 def shorten_directory(full_path: str) -> Dict:
-    parts = Path(os.path.dirname(full_path)).parts
+    # Use Path object for simplicity and cross-platform compatibility
+    path_obj = Path(full_path)
+    parts = path_obj.parts
     parts_count = len(parts)
-    dir_name = ''
-    if parts_count >= 1:
-        os.path.join(*parts[max([0, parts_count - 3]):])
 
+    # Initialize dir_name as an empty string
+    dir_name = ""
+    # Handle directory part; ignore the filename in parts
+    if parts_count > 1:
+        # Join all parts except the last one (filename)
+        dir_name = os.path.join(*parts[: parts_count - 1])
+
+    # Simplifying the directory to show only relevant parts (3 levels up max)
     if parts_count >= 4:
-        if parts_count >= 5:
-            dir_name = os.path.join('.', dir_name)
+        if parts_count > 4:
+            # For deeply nested paths, show as relative path starting with ..
+            relevant_parts = parts[-3:]
+            dir_name = os.path.join("..", *relevant_parts)
         else:
-            dir_name = os.path.join('..', dir_name)
+            # If exactly 4 levels deep, show entire path
+            dir_name = os.path.join(*parts[:-1])
 
-        for _i in range(min(1, parts_count - 4)):
-            dir_name = os.path.join('..', dir_name)
+    # Extracting extension (without the dot)
+    extension = (
+        path_obj.suffix[1:] if path_obj.suffix.startswith(".") else path_obj.suffix
+    )
 
-    extension = Path(full_path or '').suffix
-    if extension:
-        extension = extension.replace('.', '')
+    # Include the filename in the output
+    filename = path_obj.name
 
     return dict(
         directory=dir_name,
         extension=extension,
-        parts=parts,
+        parts=parts[:-1],  # Exclude the filename from 'parts'
+        filename=filename,
     )

mage_ai/data/constants.py

@@ -0,0 +1,45 @@
+from typing import AsyncGenerator, Generator, Iterable, List, Union
+
+import polars as pl
+
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.models.pyarrow.record_batch import RecordBatch, TaggedRecordBatch
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.shared.models import BaseEnum
+
+SUPPORTED_VARIABLE_TYPES = [
+    VariableType.DATAFRAME,
+    VariableType.POLARS_DATAFRAME,
+    VariableType.SERIES_PANDAS,
+    VariableType.SERIES_POLARS,
+]
+
+ScanBatchDatasetResult = Union[RecordBatch, TaggedRecordBatch, pl.DataFrame]
+OutputData = Union[
+    DataGenerator,
+    List[pl.DataFrame],
+    Iterable[ScanBatchDatasetResult],
+    List[ScanBatchDatasetResult],
+]
+
+RecordBatchGenerator = Union[Generator[ScanBatchDatasetResult, None, None], DataGenerator]
+
+AsyncRecordBatchGenerator = AsyncGenerator[ScanBatchDatasetResult, None]
+
+
+class InputDataType(BaseEnum):
+    # Batch settings will be used to fetch the data in batches,
+    # and execute the decorated function inside a yield loop.
+    BATCH = 'batch'
+    # @chunker, @chunking, @chunk
+    # CHUNKS = 'chunks' # This will be handled in the block code
+    # Original data will be passed as is.
+    DEFAULT = 'default'
+    # A generator to iterate over
+    GENERATOR = 'generator'
+    # Input variable will be a Callable function that loads the data; e.g. Variable object
+    READER = 'reader'
+
+
+class ReadModeType(BaseEnum):
+    MEMORY = 'memory'

mage_ai/data/models/base.py

@@ -0,0 +1,119 @@
+import os
+from pathlib import Path
+from typing import Any, List, Optional, Union
+
+from mage_ai.data.constants import InputDataType
+from mage_ai.data.models.constants import CHUNKS_DIRECTORY_NAME
+from mage_ai.data.models.utils import variable_type_supported
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.data_preparation.storage.base_storage import BaseStorage
+from mage_ai.settings.repo import get_variables_dir
+from mage_ai.system.models import ResourceUsage
+from mage_ai.system.storage.utils import size_of_path
+
+
+class BaseData:
+    def __init__(
+        self,
+        storage: BaseStorage,
+        variable_dir_path: str,
+        variable_path: str,
+        batch_settings: Optional[BatchSettings] = None,
+        chunks: Optional[List[Any]] = None,
+        input_data_types: Optional[List[InputDataType]] = None,
+        poll_interval: Optional[int] = None,
+        uuid: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
+        variable_types: Optional[List[VariableType]] = None,
+        variables_dir: Optional[str] = None,
+    ):
+        self.storage = storage
+        """
+        ~/.mage_data/[project_name]/pipelines/[pipeline_uuid]/.variables
+            /[pipeline_run.id]/[execution_partition]/[block.uuid]
+        ~/.mage_data/unit_3_observability/pipelines/sklearn_training/.variables
+            /9/20240518T144726_954384/load_models
+        """
+        self.variable_dir_path = variable_dir_path
+        """
+        ~/.mage_data/[project_name]/pipelines/[pipeline_uuid]/.variables
+            /[pipeline_run.id]/[execution_partition]/[block.uuid]
+            /[variable_name]
+
+        ~/.mage_data/unit_3_observability/pipelines/sklearn_training/.variables
+            /9/20240518T144726_954384/load_models
+            /output_0
+        """
+        self.batch_settings = batch_settings
+        self.chunks = chunks
+        self.poll_interval = poll_interval
+        self.variable_path = variable_path
+        self.variable_type = variable_type
+        self.variable_types = variable_types or []
+        self.variables_dir = variables_dir or get_variables_dir(root_project=False)
+        self.uuid = uuid or str(Path(self.variable_dir_path).relative_to(Path(self.variables_dir)))
+
+    @property
+    def data_source(self) -> Union[str, List[str]]:
+        if self.variable_types and len(self.variable_types) >= 1:
+            return self.__data_source_paths
+        return self.data_source_directory_path
+
+    @property
+    def number_of_outputs(self) -> int:
+        return (
+            len(self.variable_types)
+            if self.variable_types and len(self.variable_types) >= 1
+            else 1
+        )
+
+    @property
+    def data_source_directory_path(self) -> str:
+        return self.__build_data_source_path()
+
+    @property
+    def __data_source_paths(self) -> List[str]:
+        if self.variable_types:
+            return [
+                self.__build_data_source_path(str(idx)) for idx in range(len(self.variable_types))
+            ]
+        return [self.data_source_directory_path]
+
+    def __build_data_source_path(self, subdirectory: Optional[str] = None) -> str:
+        return os.path.join(self.variable_path, subdirectory or '', CHUNKS_DIRECTORY_NAME)
+
+    def is_dataframe(self) -> bool:
+        return all(
+            variable_type
+            in [
+                VariableType.DATAFRAME,
+                VariableType.POLARS_DATAFRAME,
+                VariableType.SERIES_PANDAS,
+                VariableType.SERIES_POLARS,
+            ]
+            for variable_type in (self.variable_types or [self.variable_type])
+        )
+
+    @property
+    def resource_usage(self) -> ResourceUsage:
+        return ResourceUsage.combine(self.resource_usages)
+
+    @property
+    def resource_usages(self) -> List[ResourceUsage]:
+        return [
+            ResourceUsage.load(
+                directory=path,
+                size=size_of_path(path, file_extension='parquet'),
+            )
+            for path in self.__data_source_paths
+        ]
+
+    def supported(self, data: Optional[Any] = None) -> bool:
+        if self.variable_types and len(self.variable_types) >= 1:
+            return all(
+                variable_type_supported(variable_type, data)
+                for variable_type in self.variable_types
+            )
+
+        return variable_type_supported(self.variable_type, data)

mage_ai/data/models/constants.py

@@ -0,0 +1 @@
+CHUNKS_DIRECTORY_NAME = 'chunks'

mage_ai/data/models/generator.py

@@ -0,0 +1,115 @@
+from collections.abc import Generator
+from typing import Any, Callable, Dict, List, Optional
+
+
+class DataGenerator:
+    def __init__(
+        self,
+        data: Optional[Any] = None,
+        load_data: Optional[Callable[[int], Any]] = None,
+        measure_data: Optional[Callable[[int], int]] = None,
+    ):
+        self.data = data
+        self.index = 0
+        self.load_data = load_data
+        self.measure_data = measure_data
+        self._length = None
+
+    def __len__(self):
+        # Check if the total length has been determined already
+        if self._length is None:
+            # Logic to determine the length from the datasource
+            # This could involve querying the database,
+            # reading file info, etc., without loading all data
+            self._length = self._data_length(self.index)
+        return self._length
+
+    def __iter__(self):
+        if self.data is not None and hasattr(self.data, '__iter__'):
+            for item in self.data:
+                yield item
+        elif self.data is not None and hasattr(self.data, '__getitem__'):
+            for item in self.__sequence_iterator():
+                yield item
+        else:
+            while self.index < self._data_length(self.index):
+                item = None
+                if self.load_data:
+                    item = self.load_data(self.index)
+                elif self.data is not None:
+                    item = self.data[self.index]
+                self.index += 1
+                yield item
+        self.data = None  # Explicitly set data to None after iteration is finished
+
+    def __next__(self):
+        if self.index < self._data_length(self.index):
+            # Load the next item
+            item = None
+            if self.load_data:
+                item = self.load_data(self.index)
+            elif self.data is not None:
+                item = self.data[self.index]
+            self.index += 1
+            return item
+        else:
+            self.data = None  # Set data to None when the generator is exhausted
+            # No more data, stop iteration
+            raise StopIteration
+
+    def iterable(self) -> bool:
+        return hasattr(self.data, '__iter__') or hasattr(self.data, '__getitem__')
+
+    def _data_length(self, index: Optional[int] = None) -> int:
+        if self.measure_data:
+            return self.measure_data(index or 0)
+        return len(self.data) if self.data is not None else 0
+
+    def __sequence_iterator(self) -> Any:
+        # Custom iterator for objects supporting the sequence protocol without __iter__
+        class SeqIterator:
+            def __init__(self, seq):
+                self.seq = seq
+                self.index = 0
+
+            def __next__(self):
+                try:
+                    result = self.seq[self.index]
+                except IndexError:
+                    raise StopIteration
+                self.index += 1
+                return result
+
+        return SeqIterator(self.data)
+
+
+class GeneratorWithMetadata:
+    def __init__(self, generator: Generator, metadata: Optional[List[Dict[str, Any]]]):
+        """
+        Initializes the custom generator wrapper with a generator and associated metadata.
+
+        :param generator: The generator to wrap.
+        :param metadata: A dictionary containing metadata related to the generator.
+        """
+        self._generator = generator
+        self.metadata = metadata
+
+    def __iter__(self):
+        """
+        Returns the iterator object itself.
+        """
+        return self
+
+    def __next__(self):
+        """
+        Returns the next item from the generator.
+        """
+        return next(self._generator)
+
+    def get_metadata(self) -> Dict[str, Any]:
+        """
+        Returns the stored metadata.
+
+        :return: A dictionary containing metadata.
+        """
+        return self.metadata

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)