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

mage_ai/data_preparation/variable_manager.py

@@ -1,30 +1,41 @@
 import os
+import re
 from datetime import datetime
-from typing import Any, Dict, List
-
-import pandas as pd
-
-from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
-from mage_ai.data_preparation.models.variable import (
-    VARIABLE_DIR,
-    Variable,
-    VariableType,
+from typing import Any, Dict, List, Optional
+
+from mage_ai.data.constants import InputDataType
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.data_preparation.models.utils import (
+    infer_variable_type,
+    warn_for_repo_path,
+)
+from mage_ai.data_preparation.models.variable import VARIABLE_DIR, Variable
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.data_preparation.models.variables.utils import (
+    get_first_data_output_variable_uuid,
+    is_output_variable,
 )
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.data_preparation.storage.local_storage import LocalStorage
+from mage_ai.io.base import ExportWritePolicy
 from mage_ai.settings.platform import project_platform_activated
 from mage_ai.settings.repo import get_repo_path, get_variables_dir
+from mage_ai.settings.server import MEMORY_MANAGER_V2
 from mage_ai.shared.constants import GCS_PREFIX, S3_PREFIX
 from mage_ai.shared.dates import str_to_timedelta
 from mage_ai.shared.environments import is_debug
+from mage_ai.shared.strings import to_ordinal_integers
 from mage_ai.shared.utils import clean_name
+from mage_ai.system.models import ResourceUsage
 
 
 class VariableManager:
     def __init__(self, repo_path=None, variables_dir=None):
+        warn_for_repo_path(repo_path)
         self.repo_path = repo_path or get_repo_path()
         if variables_dir is None:
-            self.variables_dir = self.repo_path
+            self.variables_dir = get_variables_dir(repo_path=self.repo_path)
         else:
             self.variables_dir = variables_dir
         self.storage = LocalStorage()
@@ -32,9 +43,9 @@ class VariableManager:
 
     @classmethod
     def get_manager(
-        self,
-        repo_path: str = None,
-        variables_dir: str = None,
+        cls,
+        repo_path: Optional[str] = None,
+        variables_dir: Optional[str] = None,
     ) -> 'VariableManager':
         manager_args = dict(
             repo_path=repo_path,
@@ -53,45 +64,137 @@ class VariableManager:
         block_uuid: str,
         variable_uuid: str,
         data: Any,
-        partition: str = None,
-        variable_type: VariableType = None,
+        partition: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
         clean_block_uuid: bool = True,
-    ) -> None:
-        if type(data) is pd.DataFrame:
-            variable_type = VariableType.DATAFRAME
-        elif is_spark_dataframe(data):
-            variable_type = VariableType.SPARK_DATAFRAME
-        elif is_geo_dataframe(data):
-            variable_type = VariableType.GEO_DATAFRAME
+        clean_variable_uuid: bool = True,
+        disable_variable_type_inference: bool = False,
+        input_data_types: Optional[List[InputDataType]] = None,
+        resource_usage: Optional[ResourceUsage] = None,
+        read_batch_settings: Optional[BatchSettings] = None,
+        read_chunks: Optional[List] = None,
+        write_batch_settings: Optional[BatchSettings] = None,
+        write_chunks: Optional[List] = None,
+    ) -> Variable:
+        """
+        Used by:
+            block
+            block/data_integration/utils
+        """
+        basic_iterable = False
+        if not disable_variable_type_inference:
+            variable_type, basic_iterable = infer_variable_type(
+                data,
+                repo_path=self.repo_path,
+                variable_type=variable_type,
+            )
+
         variable = Variable(
-            clean_name(variable_uuid),
+            clean_name(variable_uuid) if clean_variable_uuid else variable_uuid,
             self.pipeline_path(pipeline_uuid),
             block_uuid,
+            clean_block_uuid=clean_block_uuid,
+            input_data_types=input_data_types,
             partition=partition,
+            read_batch_settings=read_batch_settings,
+            read_chunks=read_chunks,
+            resource_usage=resource_usage,
             storage=self.storage,
             variable_type=variable_type,
-            clean_block_uuid=clean_block_uuid,
+            variables_dir=self.variables_dir,
+            write_batch_settings=write_batch_settings,
+            write_chunks=write_chunks,
         )
 
-        # Delete data if it exists
-        variable.delete()
-        variable.variable_type = variable_type
-        variable.write_data(data)
+        if (
+            write_batch_settings
+            and write_batch_settings.mode
+            and ExportWritePolicy.APPEND != write_batch_settings.mode
+        ):
+            # Delete data if it exists
+            variable.delete()
+
+        if (
+            basic_iterable
+            and MEMORY_MANAGER_V2
+            and variable_type
+            in [
+                VariableType.DATAFRAME,
+                VariableType.POLARS_DATAFRAME,
+                VariableType.SERIES_PANDAS,
+                VariableType.SERIES_POLARS,
+            ]
+        ):
+            for idx, data_nested in enumerate(DataGenerator(data)):
+                variable_nested = self.add_variable(
+                    pipeline_uuid,
+                    block_uuid,
+                    os.path.join(str(variable_uuid), str(idx)),
+                    data_nested,
+                    clean_block_uuid=clean_block_uuid,
+                    clean_variable_uuid=False,
+                    input_data_types=input_data_types,
+                    partition=partition,
+                    read_batch_settings=read_batch_settings,
+                    read_chunks=read_chunks,
+                    resource_usage=resource_usage,
+                    variable_type=variable_type,
+                    write_batch_settings=write_batch_settings,
+                    write_chunks=write_chunks,
+                )
+                if variable_nested.variable_type:
+                    variable.variable_types.append(variable_nested.variable_type)
+
+            variable.variable_type = VariableType.ITERABLE
+            variable.write_metadata()
+        else:
+            variable.variable_type = variable_type
+            variable.write_data(data)
+
         if is_debug():
             print(
-                f'Variable {variable_uuid} ({variable_type or "no type"}) for block {block_uuid} '
+                f'Variable {variable_uuid} ({variable_type or "no type"}) '
+                f'for block {block_uuid} '
                 f'in pipeline {pipeline_uuid} '
                 f'stored in {variable.variable_path}'
             )
 
+        return variable
+
+    def add_variable_types(
+        self,
+        pipeline_uuid: str,
+        block_uuid: str,
+        variable_uuid: str,
+        variable_types: List[VariableType],
+        clean_block_uuid: bool = True,
+        clean_variable_uuid: bool = True,
+        partition: Optional[str] = None,
+    ) -> None:
+        variable = self.build_variable(
+            pipeline_uuid,
+            block_uuid,
+            variable_uuid,
+            clean_block_uuid=clean_block_uuid,
+            clean_variable_uuid=clean_variable_uuid,
+            partition=partition,
+            variable_type=VariableType.ITERABLE,
+        )
+        variable.variable_types = variable_types
+        variable.write_metadata()
+
     def build_variable(
         self,
         pipeline_uuid: str,
         block_uuid: str,
         variable_uuid: str,
-        partition: str = None,
-        variable_type: VariableType = None,
+        partition: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
+        clean_block_uuid: bool = True,
         clean_variable_uuid: bool = True,
+        input_data_types: Optional[List[InputDataType]] = None,
+        read_batch_settings: Optional[BatchSettings] = None,
+        read_chunks: Optional[List] = None,
     ) -> Variable:
         return Variable(
             clean_name(variable_uuid) if clean_variable_uuid else variable_uuid,
@@ -100,6 +203,11 @@ class VariableManager:
             partition=partition,
             storage=self.storage,
             variable_type=variable_type,
+            clean_block_uuid=clean_block_uuid,
+            read_batch_settings=read_batch_settings,
+            read_chunks=read_chunks,
+            input_data_types=input_data_types,
+            variables_dir=self.variables_dir,
         )
 
     async def add_variable_async(
@@ -108,16 +216,26 @@ class VariableManager:
         block_uuid: str,
         variable_uuid: str,
         data: Any,
-        partition: str = None,
-        variable_type: VariableType = None,
+        partition: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
         clean_block_uuid: bool = True,
+        input_data_types: Optional[List[InputDataType]] = None,
+        resource_usage: Optional[ResourceUsage] = None,
+        read_batch_settings: Optional[BatchSettings] = None,
+        read_chunks: Optional[List] = None,
+        write_batch_settings: Optional[BatchSettings] = None,
+        write_chunks: Optional[List] = None,
     ) -> None:
-        if type(data) is pd.DataFrame:
-            variable_type = VariableType.DATAFRAME
-        elif is_spark_dataframe(data):
-            variable_type = VariableType.SPARK_DATAFRAME
-        elif is_geo_dataframe(data):
-            variable_type = VariableType.GEO_DATAFRAME
+        """
+        Used by:
+            block
+        """
+        variable_type, _ = infer_variable_type(
+            data,
+            repo_path=self.repo_path,
+            variable_type=variable_type,
+        )
+
         variable = Variable(
             clean_name(variable_uuid),
             self.pipeline_path(pipeline_uuid),
@@ -126,30 +244,43 @@ class VariableManager:
             storage=self.storage,
             variable_type=variable_type,
             clean_block_uuid=clean_block_uuid,
+            input_data_types=input_data_types,
+            resource_usage=resource_usage,
+            read_batch_settings=read_batch_settings,
+            read_chunks=read_chunks,
+            variables_dir=self.variables_dir,
+            write_batch_settings=write_batch_settings,
+            write_chunks=write_chunks,
         )
+
         # Delete data if it exists
         variable.delete()
         variable.variable_type = variable_type
         await variable.write_data_async(data)
 
-    def clean_variables(
-        self,
-        pipeline_uuid: str = None
-    ):
+    def clean_variables(self, pipeline_uuid: Optional[str] = None):
         from mage_ai.data_preparation.models.pipeline import Pipeline
 
         repo_config = get_repo_config()
-        if not repo_config.variables_retention_period:
+        if repo_config.variables_retention_period is None:
             print('Variable retention period is not provided.')
             return
-        min_partition = (datetime.utcnow() -
-                         str_to_timedelta(repo_config.variables_retention_period)).strftime(
-                            format='%Y%m%dT%H%M%S')
+
+        retention_ds = str_to_timedelta(repo_config.variables_retention_period)
+        if retention_ds is None:
+            return
+
+        min_partition = (datetime.utcnow() - retention_ds).strftime(format='%Y%m%dT%H%M%S')
+
         print(f'Clean variables before partition {min_partition}')
         if pipeline_uuid is None:
-            pipeline_uuids = Pipeline.get_all_pipelines(self.repo_path)
+            pipeline_uuids = [
+                d[0] if isinstance(d, tuple) else d
+                for d in Pipeline.get_all_pipelines(self.repo_path)
+            ]
         else:
             pipeline_uuids = [pipeline_uuid]
+
         for pipeline_uuid in pipeline_uuids:
             print(f'Removing cached variables from pipeline {pipeline_uuid}')
             pipeline_variable_path = os.path.join(
@@ -177,8 +308,8 @@ class VariableManager:
         pipeline_uuid: str,
         block_uuid: str,
         variable_uuid: str,
-        partition: str = None,
-        variable_type: VariableType = None,
+        partition: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
         clean_block_uuid: bool = True,
     ) -> None:
         Variable(
@@ -189,6 +320,7 @@ class VariableManager:
             storage=self.storage,
             variable_type=variable_type,
             clean_block_uuid=clean_block_uuid,
+            variables_dir=self.variables_dir,
         ).delete()
 
     def get_variable(
@@ -196,14 +328,19 @@ class VariableManager:
         pipeline_uuid: str,
         block_uuid: str,
         variable_uuid: str,
-        dataframe_analysis_keys: List[str] = None,
-        partition: str = None,
-        variable_type: VariableType = None,
+        dataframe_analysis_keys: Optional[List[str]] = None,
+        partition: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
         raise_exception: bool = False,
         sample: bool = False,
-        sample_count: int = None,
-        spark=None,
+        sample_count: Optional[int] = None,
+        spark: Optional[Any] = None,
         clean_block_uuid: bool = True,
+        input_data_types: Optional[List[InputDataType]] = None,
+        read_batch_settings: Optional[BatchSettings] = None,
+        read_chunks: Optional[List] = None,
+        write_batch_settings: Optional[BatchSettings] = None,
+        write_chunks: Optional[List] = None,
     ) -> Any:
         variable = self.get_variable_object(
             pipeline_uuid,
@@ -213,6 +350,11 @@ class VariableManager:
             variable_type=variable_type,
             spark=spark,
             clean_block_uuid=clean_block_uuid,
+            input_data_types=input_data_types,
+            read_batch_settings=read_batch_settings,
+            read_chunks=read_chunks,
+            write_batch_settings=write_batch_settings,
+            write_chunks=write_chunks,
         )
         return variable.read_data(
             dataframe_analysis_keys=dataframe_analysis_keys,
@@ -226,14 +368,29 @@ class VariableManager:
         self,
         pipeline_uuid: str,
         block_uuid: str,
-        variable_uuid: str,
-        partition: str = None,
-        variable_type: VariableType = None,
+        variable_uuid: Optional[str] = None,
+        partition: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
+        variable_types: Optional[List[VariableType]] = None,
         clean_block_uuid: bool = True,
+        skip_check_variable_type: Optional[bool] = None,
         spark=None,
+        input_data_types: Optional[List[InputDataType]] = None,
+        read_batch_settings: Optional[BatchSettings] = None,
+        read_chunks: Optional[List] = None,
+        write_batch_settings: Optional[BatchSettings] = None,
+        write_chunks: Optional[List] = None,
     ) -> Variable:
         if variable_type == VariableType.DATAFRAME and spark is not None:
             variable_type = VariableType.SPARK_DATAFRAME
+
+        if not variable_uuid:
+            variable_uuid = get_first_data_output_variable_uuid(
+                self.get_variables_by_block(pipeline_uuid, block_uuid, partition=partition)
+            )
+        if not variable_uuid:
+            variable_uuid = 'output_0'
+
         return Variable(
             variable_uuid,
             self.pipeline_path(pipeline_uuid),
@@ -242,11 +399,20 @@ class VariableManager:
             spark=spark,
             storage=self.storage,
             variable_type=variable_type,
+            variable_types=variable_types,
             clean_block_uuid=clean_block_uuid,
+            input_data_types=input_data_types,
+            read_batch_settings=read_batch_settings,
+            read_chunks=read_chunks,
+            skip_check_variable_type=skip_check_variable_type,
+            variables_dir=self.variables_dir,
+            write_batch_settings=write_batch_settings,
+            write_chunks=write_chunks,
         )
 
     def get_variables_by_pipeline(self, pipeline_uuid: str) -> Dict[str, List[str]]:
         from mage_ai.data_preparation.models.pipeline import Pipeline
+
         pipeline = Pipeline.get(pipeline_uuid, repo_path=self.repo_path)
         variable_dir_path = os.path.join(self.pipeline_path(pipeline_uuid), VARIABLE_DIR)
         if not self.storage.path_exists(variable_dir_path):
@@ -269,9 +435,10 @@ class VariableManager:
         self,
         pipeline_uuid: str,
         block_uuid: str,
-        partition: str = None,
+        partition: Optional[str] = None,
         clean_block_uuid: bool = True,
-        max_results: int = None,
+        max_results: Optional[int] = None,
+        output_variable_only: Optional[bool] = None,
     ) -> List[str]:
         variable_dir_path = os.path.join(
             self.pipeline_path(pipeline_uuid),
@@ -281,8 +448,29 @@ class VariableManager:
         )
         if not self.storage.path_exists(variable_dir_path):
             return []
-        variables = self.storage.listdir(variable_dir_path, max_results=max_results)
-        return sorted([v.split('.')[0] for v in variables])
+
+        opts = {}
+        if max_results is not None:
+            opts['max_results'] = max_results
+
+        variable_uuids = self.storage.listdir(variable_dir_path, **opts)
+        variable_uuids = [
+            v
+            for v in variable_uuids
+            if v.split('.')[0] and (not output_variable_only or is_output_variable(v))
+        ]
+
+        def __sort_variables(text):
+            number = re.findall('\\d+', text)
+            if number:
+                if len(number) == 1:
+                    number.append(0)
+                number = number[:2]
+            else:
+                number = to_ordinal_integers(text)[:2]
+            return [int(i) for i in number]
+
+        return sorted(variable_uuids, key=__sort_variables)
 
     def pipeline_path(self, pipeline_uuid: str) -> str:
         path = os.path.join(self.variables_dir, 'pipelines', pipeline_uuid)
@@ -308,9 +496,7 @@ class GCSVariableManager(VariableManager):
         self.storage = GCSStorage(dirpath=variables_dir)
 
 
-def clean_variables(
-    pipeline_uuid: str = None
-):
+def clean_variables(pipeline_uuid: str = None):
     variables_dir = get_variables_dir()
     VariableManager(variables_dir=variables_dir).clean_variables(pipeline_uuid=pipeline_uuid)
 
@@ -318,24 +504,39 @@ def clean_variables(
 def get_global_variables(
     pipeline_uuid: str,
     pipeline=None,
+    context_data: Dict = None,
+    repo_path: str = None,
+    variables_dir: str = None,
 ) -> Dict[str, Any]:
     """
     Get all global variables. Global variables are stored together with project's code.
     """
     from mage_ai.data_preparation.models.pipeline import Pipeline
 
-    pipeline = pipeline or Pipeline.get(pipeline_uuid, all_projects=project_platform_activated())
+    pipeline = pipeline or Pipeline.get(
+        pipeline_uuid,
+        all_projects=project_platform_activated(),
+        context_data=context_data,
+        repo_path=repo_path,
+    )
     if pipeline.variables is not None:
         global_variables = pipeline.variables
     else:
-        variables_dir = get_variables_dir()
-        variables = VariableManager(variables_dir=variables_dir).get_variables_by_block(
+        variables_dir = variables_dir or get_variables_dir(repo_path=repo_path)
+        variable_manager = VariableManager(
+            repo_path=repo_path, variables_dir=variables_dir
+        )
+        variables = variable_manager.get_variables_by_block(
             pipeline_uuid,
             'global',
         )
         global_variables = dict()
         for variable in variables:
-            global_variables[variable] = get_global_variable(pipeline_uuid, variable)
+            global_variables[variable] = variable_manager.get_variable(
+                pipeline_uuid,
+                'global',
+                variable,
+            )
 
     return global_variables
 
@@ -349,7 +550,10 @@ def get_global_variable(
     """
     from mage_ai.data_preparation.models.pipeline import Pipeline
 
-    pipeline = Pipeline.get(pipeline_uuid, all_projects=project_platform_activated())
+    pipeline = Pipeline.get(
+        pipeline_uuid,
+        all_projects=project_platform_activated(),
+    )
     if pipeline.variables is not None:
         return pipeline.variables.get(key)
     else:
@@ -360,12 +564,7 @@ def get_global_variable(
         )
 
 
-def get_variable(
-    pipeline_uuid: str,
-    block_uuid: str,
-    key: str,
-    **kwargs
-) -> Any:
+def get_variable(pipeline_uuid: str, block_uuid: str, key: str, **kwargs) -> Any:
     """
     Set block intermediate variable by key.
     Block intermediate variables are stored in variables dir.
@@ -382,12 +581,15 @@ def set_global_variable(
     pipeline_uuid: str,
     key: str,
     value: Any,
+    repo_path: str = None,
 ) -> None:
     """
     Set global variable by key. Global variables are stored together with project's code.
     """
     from mage_ai.data_preparation.models.pipeline import Pipeline
+
     pipeline = Pipeline.get(pipeline_uuid)
+    pipeline = Pipeline.get(pipeline_uuid, repo_path=repo_path)
 
     if pipeline.variables is None:
         pipeline.variables = get_global_variables(pipeline_uuid)
@@ -397,12 +599,15 @@ def set_global_variable(
 def delete_global_variable(
     pipeline_uuid: str,
     key: str,
+    repo_path: str = None,
 ) -> None:
     """
     Delete global variable by key. Global variables are stored together with project's code.
     """
     from mage_ai.data_preparation.models.pipeline import Pipeline
+
     pipeline = Pipeline.get(pipeline_uuid)
+    pipeline = Pipeline.get(pipeline_uuid, repo_path=repo_path)
     if pipeline.variables is not None:
         pipeline.delete_global_variable(key)
     else:

mage_ai/io/base.py

@@ -8,6 +8,7 @@ from pandas import DataFrame
 
 from mage_ai.io.constants import SQL_RESERVED_WORDS
 from mage_ai.shared.logger import VerbosePrintHandler
+from mage_ai.shared.models import BaseEnum
 from mage_ai.shared.utils import clean_name
 
 QUERY_ROW_LIMIT = 10_000_000
@@ -47,7 +48,7 @@ class FileFormat(str, Enum):
     XML = 'xml'
 
 
-class ExportWritePolicy(str, Enum):
+class ExportWritePolicy(BaseEnum):
     APPEND = 'append'
     FAIL = 'fail'
     REPLACE = 'replace'
@@ -150,7 +151,7 @@ class BaseFile(BaseIO):
         elif format == FileFormat.XML:
             return pd.read_xml
         else:
-            raise ValueError(f'Invalid format \'{format}\' specified.')
+            raise ValueError(f"Invalid format '{format}' specified.")
 
     def _read(
         self,

mage_ai/io/bigquery.py

@@ -242,6 +242,7 @@ WHERE TABLE_NAME = '{table_name}'
             elif len(parts) == 3:
                 database, schema, table_name = parts
 
+            database = database or self.default_database()
             df_existing = self.client.query(f"""
 SELECT 1
 FROM `{database}.{schema}.__TABLES_SUMMARY__`

mage_ai/io/redshift.py

@@ -239,12 +239,14 @@ class Redshift(BaseSQL):
                         col_with_types = ', '.join(col_with_types)
                         query = f'CREATE TABLE IF NOT EXISTS {full_table_name} ({col_with_types})'
                         cur.execute(query)
+                    if df.shape[0] > 0:
+                        columns = ', '.join([t[0] for t in columns_with_type])
+                        values = [f"""({', '.join([format_value(x) for x in v])})"""
+                                  for v in df.values]
+                        values = ', '.join(values)
+                        query = f'INSERT INTO {full_table_name} ({columns})\nVALUES {values}'
 
-                    columns = ', '.join([t[0] for t in columns_with_type])
-                    values = [f"""({', '.join([format_value(x) for x in v])})""" for v in df.values]
-                    values = ', '.join(values)
-                    query = f'INSERT INTO {full_table_name} ({columns})\nVALUES {values}'
-                    cur.execute(query)
+                        cur.execute(query)
 
                 self.conn.commit()