mage-ai 0.9.57__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/data_preparation/models/file.py

@@ -124,7 +124,11 @@ class File:
         repo_path_alt = repo_path
         if repo_path_alt is None:
             repo_path_alt = get_repo_path(file_path=file_path)
-        return File(os.path.basename(file_path), os.path.dirname(file_path), repo_path_alt)
+        return File(
+            os.path.basename(file_path or ''),
+            os.path.dirname(file_path or ''),
+            repo_path_alt or '',
+        )
 
     @classmethod
     def get_all_files(
@@ -274,6 +278,8 @@ class File:
 
         update_caches(repo_path, dir_path, filename)
 
+        update_file_cache()
+
     @classmethod
     async def write_async(
         self,
@@ -304,6 +310,8 @@ class File:
 
         await update_caches_async(repo_path, dir_path, filename)
 
+        update_file_cache()
+
     def exists(self) -> bool:
         return self.file_exists(self.file_path)
 
@@ -322,7 +330,7 @@ class File:
             async with aiofiles.open(self.file_path, mode='r', encoding='utf-8') as fp:
                 file_content = await fp.read()
             return file_content
-        except FileNotFoundError as err:
+        except Exception as err:
             print(err)
         return ''
 
@@ -369,6 +377,8 @@ class File:
     def delete(self):
         os.remove(self.file_path)
 
+        update_file_cache()
+
     def rename(self, dir_path: str, filename):
         full_path = os.path.join(self.repo_path, dir_path, filename)
 
@@ -395,6 +405,8 @@ class File:
         self.dir_path = dir_path
         self.filename = filename
 
+        update_file_cache()
+
     def to_dict(self, include_content=False):
         data = dict(name=self.filename, path=os.path.join(self.dir_path, self.filename))
         if include_content:
@@ -511,7 +523,7 @@ async def update_caches_async(repo_path: str, dir_path: str, filename: str) -> N
             dbt_cache = await DBTCache.initialize_cache_async(root_project=True)
             await dbt_cache.update_async(file_path=os.path.join(repo_path, dir_path, filename))
         except Exception as err:
-            print(f'[ERROR] File.update_caches DBTCache: {err}.')
+            print(f'[ERROR] File update DBTCache for {repo_path}, {dir_path}, {filename}: {err}.')
             if is_debug():
                 raise err
 
@@ -526,6 +538,13 @@ def update_caches(repo_path: str, dir_path: str, filename: str) -> None:
             dbt_cache = DBTCache.initialize_cache(root_project=True)
             dbt_cache.update(file_path=os.path.join(repo_path, dir_path, filename))
         except Exception as err:
-            print(f'[ERROR] File.update_caches DBTCache: {err}.')
+            print(f'[ERROR] File update DBTCache for {repo_path}, {dir_path}, {filename}: {err}.')
             if is_debug():
                 raise err
+
+
+def update_file_cache() -> None:
+    project_model = Project(root_project=True)
+    if project_model and project_model.is_feature_enabled(FeatureUUID.COMMAND_CENTER):
+        from mage_ai.cache.file import FileCache
+        FileCache().invalidate()

mage_ai/data_preparation/models/pipeline.py

@@ -3,7 +3,7 @@ import datetime
 import json
 import os
 import shutil
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Tuple, Union
 
 import aiofiles
 import dateutil.parser
@@ -310,6 +310,7 @@ class Pipeline:
             raise_exception (bool, optional): Whether to raise Exception.
                 If raise_exception = False, return None as the config when exception happens.
         """
+        include_repo_path = repo_path is not None
         repo_path = repo_path or get_repo_path()
         config_path = os.path.join(
             repo_path,
@@ -330,6 +331,9 @@ class Pipeline:
                 raise e
             config = None
 
+        if include_repo_path and config:
+            config['repo_path'] = repo_path
+
         return config
 
     @classmethod
@@ -390,7 +394,11 @@ class Pipeline:
         return pipeline
 
     @classmethod
-    def get_all_pipelines_all_projects(self, *args, **kwargs):
+    def get_all_pipelines_all_projects(
+        self,
+        *args,
+        **kwargs,
+    ) -> Union[List[str], List[Tuple[str, str]]]:
         if project_platform_activated():
             repo_paths = [d.get(
                 'full_path',
@@ -409,7 +417,8 @@ class Pipeline:
         repo_path: str = None,
         repo_paths: List[str] = None,
         disable_pipelines_folder_creation: bool = False,
-    ) -> List[str]:
+        include_repo_path: bool = False,
+    ) -> Union[List[str], List[Tuple[str, str]]]:
         arr = []
 
         paths = []
@@ -422,12 +431,13 @@ class Pipeline:
             pipelines_folder = os.path.join(path, PIPELINES_FOLDER)
             pipelines_folder_exists = os.path.exists(pipelines_folder)
             if not pipelines_folder_exists and not disable_pipelines_folder_creation:
-                os.mkdir(pipelines_folder)
-                pipelines_folder_exists = True
+                if os.path.exists(os.path.dirname(pipelines_folder)):
+                    os.mkdir(pipelines_folder)
+                    pipelines_folder_exists = True
 
             if pipelines_folder_exists:
                 arr.extend([
-                    d
+                    (d, path) if include_repo_path else d
                     for d in os.listdir(pipelines_folder)
                     if self.is_valid_pipeline(os.path.join(pipelines_folder, d))
                 ])
@@ -1157,7 +1167,7 @@ class Pipeline:
 
             for block in self.blocks_by_uuid.values():
                 if old_uuid:
-                    cache.remove_pipeline(block, old_uuid)
+                    cache.remove_pipeline(block, old_uuid, self.repo_path)
                 cache.update_pipeline(block, self)
 
         if should_update_tag_cache:
@@ -1167,7 +1177,7 @@ class Pipeline:
 
             for tag_uuid in self.tags:
                 if old_uuid:
-                    cache.remove_pipeline(tag_uuid, old_uuid)
+                    cache.remove_pipeline(tag_uuid, old_uuid, self.repo_path)
                 cache.add_pipeline(tag_uuid, self)
 
     def __update_block_order(self, blocks: List[Dict]) -> bool:
@@ -1357,6 +1367,7 @@ class Pipeline:
         spark=None,
         index: int = None,
         sample_count: int = None,
+        dynamic_block_index: int = None,
     ):
         block = self.get_block(block_uuid)
 
@@ -1385,6 +1396,7 @@ class Pipeline:
             raise_exception=raise_exception,
             spark=spark,
             variable_uuid=variable_name,
+            dynamic_block_index=dynamic_block_index,
         )
 
         return variable

mage_ai/data_preparation/models/project/__init__.py

@@ -23,8 +23,7 @@ class Project():
         self.root_project = root_project
         self.repo_path = repo_path or get_repo_path(root_project=self.root_project)
 
-        parts = self.repo_path.split('/')
-        self.name = parts[-1]
+        self.name = os.path.basename(self.repo_path)
         self.settings = None
 
         if not root_project and project_platform_activated():
@@ -32,12 +31,9 @@ class Project():
             if self.settings and self.settings.get('uuid'):
                 self.name = self.settings.get('uuid')
 
-        self.repo_config = repo_config or get_repo_config(
-            repo_path=self.repo_path,
-            root_project=self.root_project,
-        )
         self.version = VERSION
         self._features = None
+        self._repo_config = repo_config
 
         self.__repo_config_root_project = None
 
@@ -47,6 +43,15 @@ class Project():
                 root_project=True,
             )
 
+    @property
+    def repo_config(self):
+        if not self._repo_config:
+            self._repo_config = get_repo_config(
+                repo_path=self.repo_path,
+                root_project=self.root_project,
+            )
+        return self._repo_config
+
     @property
     def workspace_config_defaults(self) -> Dict:
         config = self.repo_config.workspace_config_defaults or {}
@@ -174,7 +179,7 @@ class Project():
     def is_feature_enabled(self, feature_name: FeatureUUID) -> bool:
         feature_enabled = self.features.get(feature_name.value, False)
 
-        if is_debug():
+        if is_debug() and not os.getenv('DISABLE_DATABASE_TERMINAL_OUTPUT'):
             print(f'[Project.is_feature_enabled]: {feature_name} | {feature_enabled}')
 
         return feature_enabled

mage_ai/data_preparation/models/project/constants.py

@@ -3,6 +3,8 @@ from enum import Enum
 
 class FeatureUUID(str, Enum):
     ADD_NEW_BLOCK_V2 = 'add_new_block_v2'
+    CODE_BLOCK_V2 = 'code_block_v2'
+    COMMAND_CENTER = 'command_center'
     COMPUTE_MANAGEMENT = 'compute_management'
     CUSTOM_DESIGN = 'custom_design'
     DATA_INTEGRATION_IN_BATCH_PIPELINE = 'data_integration_in_batch_pipeline'

mage_ai/data_preparation/models/triggers/__init__.py

@@ -29,6 +29,13 @@ class ScheduleType(str, enum.Enum):
     TIME = 'time'
 
 
+SCHEDULE_TYPE_TO_LABEL = {
+  ScheduleType.API: 'API',
+  ScheduleType.EVENT: 'Event',
+  ScheduleType.TIME: 'Schedule',
+}
+
+
 class ScheduleInterval(str, enum.Enum):
     ONCE = '@once'
     HOURLY = '@hourly'

mage_ai/data_preparation/models/utils.py

@@ -4,6 +4,7 @@ from typing import Any, Dict
 import dask.dataframe as dd
 import numpy
 import pandas as pd
+import polars as pl
 import simplejson
 import yaml
 
@@ -24,6 +25,11 @@ CAST_TYPE_COLUMN_TYPES = set([
     'float64',
 ])
 
+POLARS_CAST_TYPE_COLUMN_TYPES = {
+    'Float64': pl.Float64,
+    'Int64': pl.Int64,
+}
+
 
 def serialize_columns(row: pd.Series, column_types: Dict) -> pd.Series:
     for column, column_type in column_types.items():
@@ -54,6 +60,16 @@ def cast_column_types(df: pd.DataFrame, column_types: Dict):
     return df
 
 
+def cast_column_types_polars(df: pl.DataFrame, column_types: Dict):
+    for column, column_type in column_types.items():
+        if column_type in POLARS_CAST_TYPE_COLUMN_TYPES:
+            try:
+                df = df.cast({column: POLARS_CAST_TYPE_COLUMN_TYPES.get(column_type)})
+            except Exception:
+                traceback.print_exc()
+    return df
+
+
 def deserialize_columns(row: pd.Series, column_types: Dict) -> pd.Series:
     for column, column_type in column_types.items():
         if column_type not in JSON_SERIALIZABLE_COLUMN_TYPES:
@@ -85,6 +101,10 @@ def apply_transform_pandas(df: pd.DataFrame, apply_function) -> pd.DataFrame:
     return df.apply(apply_function, axis=1)
 
 
+def apply_transform_polars(df: pl.DataFrame, apply_function) -> pl.DataFrame:
+    return df.apply(apply_function, axis=1)
+
+
 def should_serialize_pandas(column_types: Dict) -> bool:
     if not column_types:
         return False

mage_ai/data_preparation/models/variable.py

@@ -21,6 +21,7 @@ from mage_ai.data_preparation.models.utils import (  # dask_from_pandas,
     STRING_SERIALIZABLE_COLUMN_TYPES,
     apply_transform_pandas,
     cast_column_types,
+    cast_column_types_polars,
     deserialize_columns,
     serialize_columns,
     should_deserialize_pandas,
@@ -36,6 +37,8 @@ DATAFRAME_PARQUET_FILE = 'data.parquet'
 DATAFRAME_PARQUET_SAMPLE_FILE = 'sample_data.parquet'
 DATAFRAME_CSV_FILE = 'data.csv'
 
+METADATA_FILE = 'type.json'
+
 JSON_FILE = 'data.json'
 JSON_SAMPLE_FILE = 'sample_data.json'
 
@@ -87,26 +90,43 @@ class Variable:
     def variable_path(self):
         return os.path.join(self.variable_dir_path, self.uuid or '')
 
+    @property
+    def metadata_path(self):
+        return os.path.join(self.variable_path, METADATA_FILE)
+
     @classmethod
     def dir_path(self, pipeline_path, block_uuid):
         return os.path.join(pipeline_path, VARIABLE_DIR, clean_name(block_uuid))
 
     def check_variable_type(self, spark=None):
         """
-        Infer variable type based on data in the storage.
+        If the variable has a metadata file, read the variable type from the metadata file.
+        Fallback to inferring variable type based on data in the storage.
         """
+        if self.variable_type is None:
+            try:
+                if self.storage.path_exists(self.metadata_path):
+                    metadata = self.storage.read_json_file(self.metadata_path, raise_exception=True)
+                    self.variable_type = metadata.get('type')
+            except Exception:
+                traceback.print_exc()
+
         if self.variable_type is None and self.storage.path_exists(
             os.path.join(self.variable_path, DATAFRAME_PARQUET_FILE)
         ):
             # If parquet file exists for given variable, set the variable type to DATAFRAME
             self.variable_type = VariableType.DATAFRAME
-        elif ((self.variable_type == VariableType.DATAFRAME or self.variable_type is None)
-                and os.path.exists(
-                os.path.join(self.variable_dir_path, f'{self.uuid}', 'data.sh'))):
+        elif (
+            self.variable_type == VariableType.DATAFRAME or self.variable_type is None
+        ) and os.path.exists(
+            os.path.join(self.variable_path, f'{self.uuid}', 'data.sh')
+        ):
             self.variable_type = VariableType.GEO_DATAFRAME
-        elif self.variable_type is None and \
-                len(self.storage.listdir(self.variable_path, suffix='.parquet')) > 0 and \
-                spark is not None:
+        elif (
+            self.variable_type is None
+            and len(self.storage.listdir(self.variable_path, suffix='.parquet')) > 0
+            and spark is not None
+        ):
             self.variable_type = VariableType.SPARK_DATAFRAME
 
     def convert_parquet_to_csv(self):
@@ -164,6 +184,12 @@ class Variable:
                 sample=sample,
                 sample_count=sample_count,
             )
+        elif self.variable_type == VariableType.POLARS_DATAFRAME:
+            return self.__read_polars_parquet(
+                raise_exception=raise_exception,
+                sample=sample,
+                sample_count=sample_count,
+            )
         elif self.variable_type == VariableType.SPARK_DATAFRAME:
             return self.__read_spark_parquet(sample=sample, sample_count=sample_count, spark=spark)
         elif self.variable_type == VariableType.GEO_DATAFRAME:
@@ -193,6 +219,11 @@ class Variable:
         """
         if self.variable_type == VariableType.DATAFRAME:
             return self.__read_parquet(sample=sample, sample_count=sample_count)
+        elif self.variable_type == VariableType.POLARS_DATAFRAME:
+            return self.__read_polars_parquet(
+                sample=sample,
+                sample_count=sample_count,
+            )
         elif self.variable_type == VariableType.SPARK_DATAFRAME:
             return self.__read_spark_parquet(sample=sample, sample_count=sample_count, spark=spark)
         elif self.variable_type == VariableType.DATAFRAME_ANALYSIS:
@@ -231,6 +262,12 @@ class Variable:
         elif is_geo_dataframe(data):
             self.variable_type = VariableType.GEO_DATAFRAME
 
+        # Dataframe analysis variables share the same uuid as the original dataframe variable
+        # so we won't write the metadata file for them
+        if self.variable_type == VariableType.DATAFRAME_ANALYSIS:
+            self.__write_dataframe_analysis(data)
+            return
+
         if self.variable_type == VariableType.DATAFRAME:
             self.__write_parquet(data)
         elif self.variable_type == VariableType.POLARS_DATAFRAME:
@@ -239,11 +276,11 @@ class Variable:
             self.__write_spark_parquet(data)
         elif self.variable_type == VariableType.GEO_DATAFRAME:
             self.__write_geo_dataframe(data)
-        elif self.variable_type == VariableType.DATAFRAME_ANALYSIS:
-            self.__write_dataframe_analysis(data)
         else:
             self.__write_json(data)
 
+        self.write_metadata()
+
     async def write_data_async(self, data: Any) -> None:
         """
         Write variable data to the persistent storage.
@@ -260,6 +297,10 @@ class Variable:
         elif is_geo_dataframe(data):
             self.variable_type = VariableType.GEO_DATAFRAME
 
+        if self.variable_type == VariableType.DATAFRAME_ANALYSIS:
+            self.__write_dataframe_analysis(data)
+            return
+
         if self.variable_type == VariableType.DATAFRAME:
             self.__write_parquet(data)
         elif self.variable_type == VariableType.POLARS_DATAFRAME:
@@ -268,11 +309,22 @@ class Variable:
             self.__write_spark_parquet(data)
         elif self.variable_type == VariableType.GEO_DATAFRAME:
             self.__write_geo_dataframe(data)
-        elif self.variable_type == VariableType.DATAFRAME_ANALYSIS:
-            self.__write_dataframe_analysis(data)
         else:
             await self.__write_json_async(data)
 
+        self.write_metadata()
+
+    def write_metadata(self) -> None:
+        """
+        Write metadata to the persistent storage.
+        """
+        metadata = dict(
+            type=self.variable_type.value
+            if isinstance(self.variable_type, VariableType)
+            else self.variable_type,
+        )
+        self.storage.write_json_file(self.metadata_path, metadata)
+
     def __delete_dataframe_analysis(self) -> None:
         for k in DATAFRAME_ANALYSIS_KEYS:
             file_path = os.path.join(self.variable_path, f'{k}.json')
@@ -300,7 +352,7 @@ class Variable:
         self,
         default_value: Dict = None,
         raise_exception: bool = False,
-        sample: bool = False
+        sample: bool = False,
     ) -> Dict:
         if default_value is None:
             default_value = {}
@@ -443,6 +495,51 @@ class Variable:
             df = cast_column_types(df, column_types)
         return df
 
+    def __read_polars_parquet(
+        self,
+        sample: bool = False,
+        sample_count: int = None,
+        raise_exception: bool = False,
+    ) -> pl.DataFrame:
+        file_path = os.path.join(self.variable_path, DATAFRAME_PARQUET_FILE)
+        sample_file_path = os.path.join(
+            self.variable_path, DATAFRAME_PARQUET_SAMPLE_FILE
+        )
+
+        read_sample_success = False
+        if sample:
+            try:
+                df = self.storage.read_polars_parquet(sample_file_path, use_pyarrow=True)
+                read_sample_success = True
+            except Exception as ex:
+                if raise_exception:
+                    raise Exception(
+                        f'Failed to read parquet file: {sample_file_path}'
+                    ) from ex
+                else:
+                    traceback.print_exc()
+        if not read_sample_success:
+            try:
+                df = self.storage.read_polars_parquet(file_path, use_pyarrow=True)
+            except Exception as ex:
+                if raise_exception:
+                    raise Exception(f'Failed to read parquet file: {file_path}') from ex
+                else:
+                    traceback.print_exc()
+                df = pl.DataFrame()
+        if sample:
+            sample_count = sample_count or DATAFRAME_SAMPLE_COUNT
+            if df.shape[0] > sample_count:
+                df = df.head(sample_count)
+
+        column_types_filename = os.path.join(self.variable_path, DATAFRAME_COLUMN_TYPES_FILE)
+        if self.storage.path_exists(column_types_filename):
+            column_types = self.storage.read_json_file(column_types_filename)
+            # No Mage specific code to serialize columns for polars when writing a variable,
+            # so no need to deserialize columns here
+            df = cast_column_types_polars(df, column_types)
+        return df
+
     def __read_spark_parquet(self, sample: bool = False, sample_count: int = None, spark=None):
         if spark is None:
             return None
@@ -514,7 +611,7 @@ class Variable:
                     pl_df = pl.from_pandas(df_output)
                     self.__write_polars_dataframe(pl_df)
                     # Test read dataframe from parquet
-                    self. __read_parquet(sample=True, raise_exception=True)
+                    self.__read_parquet(sample=True, raise_exception=True)
 
                     return
                 except Exception:
@@ -536,7 +633,7 @@ class Variable:
         try:
             df_sample_output = df_output_serialized.iloc[
                 :DATAFRAME_SAMPLE_COUNT,
-                :DATAFRAME_SAMPLE_MAX_COLUMNS
+                :DATAFRAME_SAMPLE_MAX_COLUMNS,
             ]
 
             self.storage.write_parquet(
@@ -616,7 +713,8 @@ class Variable:
             if dataframe_analysis_keys is not None and k not in dataframe_analysis_keys:
                 continue
             result[k] = await self.storage.read_json_file_async(
-                os.path.join(self.variable_path, f'{k}.json'))
+                os.path.join(self.variable_path, f'{k}.json')
+            )
         return result
 
     def __write_dataframe_analysis(self, data: Dict[str, Dict]) -> None:

mage_ai/data_preparation/shared/utils.py

@@ -4,21 +4,24 @@ from typing import Callable, Dict
 from mage_ai.shared.hash import get_json_value, merge_dict
 
 
-def get_template_vars() -> Dict[str, Callable]:
+def get_template_vars(include_python_libraries: Dict = None) -> Dict[str, Callable]:
     from mage_ai.data_preparation.shared.secrets import get_secret_value
-    no_db_kwargs = get_template_vars_no_db()
+    no_db_kwargs = get_template_vars_no_db(include_python_libraries=include_python_libraries)
 
     kwargs = dict(mage_secret_var=get_secret_value)
 
     return merge_dict(no_db_kwargs, kwargs)
 
 
-def get_template_vars_no_db() -> Dict[str, Callable]:
+def get_template_vars_no_db(include_python_libraries: Dict = None) -> Dict[str, Callable]:
     kwargs = dict(
         env_var=os.getenv,
         json_value=get_json_value,
     )
 
+    if include_python_libraries:
+        kwargs.update(include_python_libraries)
+
     try:
         from mage_ai.services.aws.secrets_manager.secrets_manager import get_secret
         kwargs['aws_secret_var'] = get_secret

mage_ai/data_preparation/storage/base_storage.py

@@ -2,6 +2,9 @@ from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import Dict, List
 
+import pandas as pd
+import polars as pl
+
 
 class BaseStorage(ABC):
     @abstractmethod
@@ -46,6 +49,20 @@ class BaseStorage(ABC):
         """
         pass
 
+    @abstractmethod
+    def read_parquet(self, file_path: str, **kwargs) -> pd.DataFrame:
+        """
+        Read parquet from a file with file path and return a pandas DataFrame.
+        """
+        pass
+
+    @abstractmethod
+    def read_polars_parquet(self, file_path: str, **kwargs) -> pl.DataFrame:
+        """
+        Read parquet from a file with file path and return a polars DataFrame.
+        """
+        pass
+
     @abstractmethod
     def read_json_file(
         self,

mage_ai/data_preparation/storage/gcs_storage.py

@@ -107,6 +107,10 @@ class GCSStorage(BaseStorage):
         buffer = io.BytesIO(self.bucket.blob(gcs_url_path(file_path)).download_as_bytes())
         return pd.read_parquet(buffer, **kwargs)
 
+    def read_polars_parquet(self, file_path: str, **kwargs) -> pl.DataFrame:
+        buffer = io.BytesIO(self.bucket.blob(gcs_url_path(file_path)).download_as_bytes())
+        return pl.read_parquet(buffer, **kwargs)
+
     def write_parquet(self, df: pd.DataFrame, file_path: str) -> None:
         buffer = io.BytesIO()
         df.to_parquet(buffer)

mage_ai/data_preparation/storage/local_storage.py

@@ -96,6 +96,9 @@ class LocalStorage(BaseStorage):
     def read_parquet(self, file_path: str, **kwargs) -> pd.DataFrame:
         return pd.read_parquet(file_path, engine='pyarrow')
 
+    def read_polars_parquet(self, file_path: str, **kwargs) -> pl.DataFrame:
+        return pl.read_parquet(file_path, use_pyarrow=True)
+
     def write_csv(self, df: pd.DataFrame, file_path: str) -> None:
         File.create_parent_directories(file_path)
         df.to_csv(file_path, index=False)
@@ -135,3 +138,15 @@ class LocalStorage(BaseStorage):
             except Exception as err:
                 if is_debug():
                     print(f'[ERROR] LocalStorage.read_async: {err}')
+
+    def read(self, file_path: str) -> str:
+        dirname = os.path.dirname(file_path)
+        if not os.path.isdir(dirname):
+            os.mkdir(dirname)
+
+        with open(file_path, mode='r') as file:
+            try:
+                return file.read()
+            except Exception as err:
+                if is_debug():
+                    print(f'[ERROR] LocalStorage.read: {err}')

mage_ai/data_preparation/storage/s3_storage.py

@@ -97,6 +97,10 @@ class S3Storage(BaseStorage):
         buffer = io.BytesIO(self.client.get_object(s3_url_path(file_path)).read())
         return pd.read_parquet(buffer, **kwargs)
 
+    def read_polars_parquet(self, file_path: str, **kwargs) -> pl.DataFrame:
+        buffer = io.BytesIO(self.client.get_object(s3_url_path(file_path)).read())
+        return pl.read_parquet(buffer, **kwargs)
+
     def write_parquet(self, df: pd.DataFrame, file_path: str) -> None:
         buffer = io.BytesIO()
         df.to_parquet(buffer)

mage_ai/data_preparation/templates/data_exporters/qdrant.py

@@ -5,7 +5,6 @@ from pandas import DataFrame
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.io.config import ConfigFileLoader
 from mage_ai.io.qdrant import Qdrant
-from sentence_transformers import SentenceTransformer
 
 if 'data_exporter' not in globals():
     from mage_ai.data_preparation.decorators import data_exporter

mage_ai/data_preparation/templates/data_loaders/qdrant.py

@@ -8,9 +8,10 @@ from os import path
 {{ super() -}}
 {% endblock %}
 
-DEFAULT_MODEL = 'all-MiniLM-L6-v2'
 
 {% block content %}
+DEFAULT_MODEL = 'all-MiniLM-L6-v2'
+
 @data_loader
 def load_data_from_qdrant(*args, **kwargs):
     """

mage_ai/data_preparation/variable_manager.py

@@ -71,6 +71,7 @@ class VariableManager:
             variable_type=variable_type,
             clean_block_uuid=clean_block_uuid,
         )
+
         # Delete data if it exists
         variable.delete()
         variable.variable_type = variable_type