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

mage_ai/data/tabular/writer.py

@@ -0,0 +1,266 @@
+import asyncio
+import glob
+import json
+from typing import Dict, List, Optional, Union
+
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+from pyarrow import parquet as pq
+
+from mage_ai.data.tabular.constants import COLUMN_CHUNK
+from mage_ai.data.tabular.models import DEFAULT_BATCH_ITEMS_VALUE, BatchSettings
+from mage_ai.data.tabular.utils import multi_series_to_frame
+from mage_ai.shared.environments import is_debug
+
+
+def append_chunk_column(
+    df: pl.DataFrame,
+    chunk_size: Optional[int] = None,
+    total_rows: Optional[int] = None,
+    use_index_as_chunk: Optional[bool] = False,
+    index: int = 0,
+) -> pl.DataFrame:
+    if total_rows is None:
+        total_rows = df.height
+
+    if use_index_as_chunk:
+        chunks = [index] * total_rows
+    elif chunk_size is not None:
+        chunks = [i // chunk_size for i in range(total_rows)]
+    else:
+        chunks = [
+            0
+        ] * total_rows  # Default to all rows in the same chunk if no parameters are provided
+
+    return df.with_columns([pl.Series(COLUMN_CHUNK, chunks).cast(pl.Int32)])
+
+
+def add_chunk_column(
+    df: pl.DataFrame,
+    chunk_size: Optional[int] = None,
+    num_buckets: Optional[int] = None,
+    use_index_as_chunk: Optional[bool] = False,
+    index: int = 0,
+) -> pl.DataFrame:
+    """
+    Adds a 'chunk' column to the DataFrame based on the specified number
+    of buckets or chunk size or uses index as chunk number when asked.
+    :param df: Input Polars DataFrame.
+    :param num_buckets: The desired number of chunks (buckets).
+    :param chunk_size: The number of rows each chunk should contain.
+    :param use_index_as_chunk: Use index as the chunk number.
+    :param index: Index number when using index as the chunk number.
+    :return: DataFrame with an added 'chunk' column.
+    """
+    total_rows = df.height
+
+    if num_buckets is not None:
+        # Calculate chunk size based on the total number of rows and the desired number of buckets
+        chunk_size = max(total_rows // num_buckets, 1)
+        num_buckets = (total_rows + chunk_size - 1) // chunk_size
+
+    if is_debug():
+        print(f'Chunk size: {chunk_size}, Total rows: {total_rows}, Num buckets: {num_buckets}')
+
+    return append_chunk_column(df, chunk_size, total_rows, use_index_as_chunk, index)
+
+
+def add_custom_metadata_to_table(table: pa.Table, metadata: Dict) -> pa.Table:
+    # This is a placeholder function. Adapt it based on your actual metadata handling.
+    # Access existing metadata (if any) and merge with custom metadata
+    existing_metadata = table.schema.metadata if table.schema.metadata is not None else {}
+    updated_metadata = {
+        **existing_metadata,
+        **{key.encode(): str(value).encode() for key, value in metadata.items()},
+    }
+
+    # Update PyArrow Table schema with merged metadata
+    new_schema = table.schema.with_metadata(updated_metadata)
+    return table.replace_schema_metadata(new_schema.metadata)
+
+
+def to_parquet_sync(
+    output_dir: str,
+    df: Optional[Union[pl.DataFrame, pl.Series, pd.Series]] = None,
+    dfs: Optional[Union[List[pl.DataFrame], pl.Series, pd.Series]] = None,
+    basename_template: Optional[str] = None,
+    existing_data_behavior: Optional[str] = None,
+    metadata: Optional[Dict] = None,
+    partition_cols: Optional[List[str]] = None,
+    settings: Optional[BatchSettings] = None,
+) -> Dict[str, int]:
+    total_rows = 0
+    total_columns = 0
+
+    for table, partition_columns in __prepare_data(
+        output_dir=output_dir,
+        df=df,
+        dfs=dfs,
+        metadata=metadata,
+        partition_cols=partition_cols,
+        settings=settings,
+    ):
+        pq.write_to_dataset(
+            table,
+            basename_template=basename_template,
+            compression='snappy',
+            existing_data_behavior=existing_data_behavior,
+            partition_cols=partition_columns,
+            root_path=output_dir,
+            use_dictionary=True,
+        )
+        total_rows += table.num_rows
+        total_columns = max(len(table.schema.names), total_columns)
+
+    return dict(
+        columns=total_columns,
+        rows=total_rows,
+    )
+
+
+async def to_parquet_async(
+    output_dir: str,
+    df: Optional[Union[pl.DataFrame, pl.Series, pd.Series]] = None,
+    dfs: Optional[Union[List[pl.DataFrame], pl.Series, pd.Series]] = None,
+    basename_template: Optional[str] = None,
+    existing_data_behavior: Optional[str] = None,
+    metadata: Optional[Dict] = None,
+    partition_cols: Optional[List[str]] = None,
+    settings: Optional[BatchSettings] = None,
+) -> Dict[str, int]:
+    total_rows = 0
+    total_columns = 0
+
+    for table, partition_columns in __prepare_data(
+        output_dir=output_dir,
+        df=df,
+        dfs=dfs,
+        metadata=metadata,
+        partition_cols=partition_cols,
+        settings=settings,
+    ):
+        await __write_to_dataset_async(
+            table,
+            root_path=output_dir,
+            existing_data_behavior=existing_data_behavior,
+            partition_cols=partition_columns,
+        )
+        total_rows += table.num_rows
+        total_columns = max(len(table.schema.names), total_columns)
+
+    return dict(
+        columns=total_columns,
+        rows=total_rows,
+    )
+
+
+async def __write_to_dataset_async(
+    table: pa.Table,
+    root_path: str,
+    basename_template: Optional[str] = None,
+    existing_data_behavior: Optional[str] = None,
+    partition_cols: Optional[List[str]] = None,
+):
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(
+        None,
+        pq.write_to_dataset,
+        table,
+        basename_template=basename_template,
+        compression='snappy',
+        existing_data_behavior=existing_data_behavior,
+        partition_cols=partition_cols,
+        root_path=root_path,
+        use_dictionary=True,
+    )
+
+
+def __prepare_data(
+    output_dir: str,
+    df: Optional[Union[pd.DataFrame, pl.DataFrame, pl.Series, pd.Series]] = None,
+    dfs: Optional[Union[List[pd.DataFrame], List[pl.DataFrame], pl.Series, pd.Series]] = None,
+    metadata: Optional[Dict] = None,
+    partition_cols: Optional[List[str]] = None,
+    settings: Optional[BatchSettings] = None,
+):
+    """
+    Writes a Polars DataFrame to partitioned Parquet files directly using PyArrow,
+    while adding custom metadata and managing partitioning automatically.
+    :param df: Polars DataFrame to write.
+    :param output_dir: Base directory for writing partitioned Parquet files.
+    :param partition_cols: List of column names on which to partition the data.
+    :param chunk_size: Number of rows per chunk.
+    :param metadata: Dictionary with custom metadata to add to the Parquet files.
+    """
+    if not settings:
+        settings = BatchSettings()
+
+    chunk_size = None
+    num_buckets = None
+    if settings.items and settings.items.maximum:
+        chunk_size = settings.items.maximum
+    elif settings.count and settings.count.maximum:
+        num_buckets = settings.count.maximum
+
+    if chunk_size is not None and num_buckets is not None:
+        chunk_size = DEFAULT_BATCH_ITEMS_VALUE
+
+    df, dfs, series_sample, object_metadata = multi_series_to_frame(df, dfs)
+
+    chunk_sizes = [] + ([chunk_size] if chunk_size else [])
+    if dfs is not None:
+        num_buckets = len(dfs)
+        df = pl.DataFrame([])
+
+        # Process a list of DataFrames with index as chunk if dfs is provided
+        for index, dataframe in enumerate(dfs):
+            chunk_size_for_index = dataframe.height
+
+            # Use index as chunk number
+            dataframe_with_chunk = append_chunk_column(
+                dataframe,
+                use_index_as_chunk=True,
+                index=index,
+                total_rows=chunk_size_for_index,
+            )
+
+            # Convert the Polars DataFrame to a PyArrow Table
+            df = df.vstack(dataframe_with_chunk)
+
+            chunk_sizes.append(chunk_size_for_index)
+            if is_debug():
+                print(f'Chunk size for index {index}: {chunk_size_for_index}')
+    elif df is not None and (chunk_size or num_buckets) and isinstance(df, pl.DataFrame):
+        df = add_chunk_column(df, chunk_size=chunk_size, num_buckets=num_buckets)
+
+    if df is None:
+        raise ValueError('No DataFrame provided to write.')
+
+    metadata = metadata or {}
+    if chunk_sizes:
+        metadata['chunk_sizes'] = chunk_sizes
+    if num_buckets:
+        metadata['num_buckets'] = num_buckets
+    if object_metadata:
+        metadata['object'] = json.dumps(object_metadata)
+
+    # Convert the Polars DataFrame to a PyArrow Table first
+    table = add_custom_metadata_to_table(df.to_arrow(), metadata)
+
+    if chunk_size or num_buckets:
+        partition_cols = (partition_cols or []) + [COLUMN_CHUNK]
+
+    # Utilize PyArrow's write_to_dataset function to handle partitioning
+    yield table, partition_cols
+
+    partition_directories = glob.glob(f'{output_dir}/*')
+    num_partitions = len(partition_directories)
+    total_rows = table.num_rows
+    total_columns = len(table.schema.names)
+
+    if is_debug():
+        print(
+            f'{total_rows} rows with {total_columns} columns '
+            f"written to '{output_dir}' across {num_partitions} partitions."
+        )

mage_ai/data/variables/wrapper.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import Any, List, Optional
+
+from mage_ai.data.constants import InputDataType
+
+
+class VariableWrapper:
+    def __init__(self, variable, input_data_type: InputDataType):
+        self.variable = variable
+        self.input_data_type = input_data_type
+
+
+def wrap_input_variables(
+    block,
+    input_data: List[Any],
+    block_uuids: List[str],
+    variable_uuids: List[List[str]],
+) -> List[Any]:
+    return [
+        wrap_outputs(outputs, variables, block.input_types(block_uuid))
+        for outputs, block_uuid, variables in zip(input_data, block_uuids, variable_uuids)
+    ]
+
+
+def wrap_outputs(
+    outputs: Any,
+    variable_uuids: List[str],
+    input_types: Optional[List[InputDataType]],
+) -> Any:
+    n_variables = len(variable_uuids)
+    n_outputs = len(outputs)
+
+    if n_outputs > n_variables:
+        raise Exception(
+            f'[WRAP_DATA] Too many outputs for the number of variables: {n_outputs} > '
+            f'{n_variables}'
+        )
+
+    if not input_types:
+        input_types = [InputDataType.DEFAULT]
+    if n_outputs < n_outputs:
+        input_types += [InputDataType.DEFAULT] * (n_variables - n_outputs)
+
+    return [
+        wrap_data(data, input_data_type)
+        for data, input_data_type in zip(outputs, input_types[:n_outputs])
+    ]
+
+
+def wrap_data(data: Any, input_data_type: InputDataType) -> Any:
+    if InputDataType.GENERATOR == input_data_type:
+        pass
+    return data

mage_ai/data_cleaner/analysis/charts.py

@@ -1,7 +1,13 @@
+import math
+from typing import Dict
+
+import numpy as np
+import pandas as pd
+
 from mage_ai.data_cleaner.analysis.constants import (
     CHART_TYPE_BAR_HORIZONTAL,
-    CHART_TYPE_LINE_CHART,
     CHART_TYPE_HISTOGRAM,
+    CHART_TYPE_LINE_CHART,
     DATA_KEY_SCATTER_PLOT,
     DATA_KEY_SCATTER_PLOT_LABELS,
     DATA_KEY_TIME_SERIES,
@@ -9,18 +15,16 @@ from mage_ai.data_cleaner.analysis.constants import (
 )
 from mage_ai.data_cleaner.column_types.constants import ColumnType
 from mage_ai.data_cleaner.estimators.encoders import MultipleColumnLabelEncoder
-import math
-import numpy as np
-import pandas as pd
+from mage_ai.shared.parsers import convert_matrix_to_dataframe
 
-DD_KEY = 'lambda.analysis_charts'
+DD_KEY = "lambda.analysis_charts"
 BUCKETS = 40
 SCATTER_PLOT_SAMPLE_COUNT = 200
 SCATTER_PLOT_CATEGORY_LIMIT = 10
 TIME_SERIES_BUCKETS = 40
 
 
-def increment(metric, tags={}):
+def increment(metric, tags: Dict = None):
     pass
 
 
@@ -34,13 +38,17 @@ def build_buckets(min_value, max_value, max_buckets, column_type):
         diff = 0
 
     is_integer = False
-    parts = str(diff).split('.')
+    parts = str(diff).split(".")
     if len(parts) == 1:
         is_integer = True
     else:
         is_integer = int(parts[1]) == 0
 
-    if ColumnType.NUMBER == column_type and total_interval <= max_buckets and is_integer:
+    if (
+        ColumnType.NUMBER == column_type
+        and total_interval <= max_buckets
+        and is_integer
+    ):
         number_of_buckets = int(total_interval)
         bucket_interval = 1
     elif bucket_interval > 1:
@@ -64,7 +72,10 @@ def build_buckets(min_value, max_value, max_buckets, column_type):
 
 
 def build_histogram_data(col1, series, column_type):
-    increment(f'{DD_KEY}.build_histogram_data.start', dict(feature_uuid=col1))
+    increment(f"{DD_KEY}.build_histogram_data.start", dict(feature_uuid=col1))
+
+    if not isinstance(series, pd.Series):
+        series = convert_matrix_to_dataframe(series)
 
     max_value = series.max() if len(series) >= 1 else None
     min_value = series.min() if len(series) >= 1 else None
@@ -74,7 +85,7 @@ def build_histogram_data(col1, series, column_type):
     if bucket_interval == 0:
         return
 
-    bins = [b['min_value'] for b in buckets] + [buckets[-1]['max_value']]
+    bins = [b["min_value"] for b in buckets] + [buckets[-1]["max_value"]]
     count, _ = np.histogram(series, bins=bins)
 
     x = []
@@ -83,13 +94,13 @@ def build_histogram_data(col1, series, column_type):
     for idx, bucket in enumerate(buckets):
         x.append(
             dict(
-                max=bucket['max_value'],
-                min=bucket['min_value'],
+                max=bucket["max_value"],
+                min=bucket["min_value"],
             )
         )
         y.append(dict(value=count[idx]))
 
-    increment(f'{DD_KEY}.build_histogram_data.succeeded', dict(feature_uuid=col1))
+    increment(f"{DD_KEY}.build_histogram_data.succeeded", dict(feature_uuid=col1))
 
     return dict(
         type=CHART_TYPE_HISTOGRAM,
@@ -114,11 +125,13 @@ def build_correlation_data(df):
                 if value is not None:
                     x.append(dict(label=col2))
                     y.append(dict(value=value))
-        charts[col1] = [dict(
-            type=CHART_TYPE_BAR_HORIZONTAL,
-            x=x,
-            y=y,
-        )]
+        charts[col1] = [
+            dict(
+                type=CHART_TYPE_BAR_HORIZONTAL,
+                x=x,
+                y=y,
+            )
+        ]
     return charts
 
 
@@ -148,11 +161,12 @@ def build_time_series_data(df, features, datetime_column):
     df_copy[datetime_column] = datetimes.view(int) / 10**9
 
     for bucket in buckets:
-        max_value = bucket['max_value']
-        min_value = bucket['min_value']
+        max_value = bucket["max_value"]
+        min_value = bucket["min_value"]
 
         df_filtered = df_copy[
-            (df_copy[datetime_column] >= min_value) & (df_copy[datetime_column] < max_value)
+            (df_copy[datetime_column] >= min_value)
+            & (df_copy[datetime_column] < max_value)
         ]
 
         x.append(
@@ -164,8 +178,8 @@ def build_time_series_data(df, features, datetime_column):
 
         series_count = df_filtered.shape[0]
         for f in features:
-            col = f['uuid']
-            column_type = f['column_type']
+            col = f["uuid"]
+            column_type = f["column_type"]
             if col not in y_dict:
                 y_dict[col] = []
 
@@ -233,22 +247,22 @@ def build_overview_data(
     datetime_features,
     numeric_features,
 ):
-    increment(f'{DD_KEY}.build_overview_data.start')
+    increment(f"{DD_KEY}.build_overview_data.start")
 
     time_series = []
     df_copy = df.copy()
 
     for feature in datetime_features:
-        column_type = feature['column_type']
-        datetime_column = feature['uuid']
+        column_type = feature["column_type"]
+        datetime_column = feature["uuid"]
         tags = dict(datetime_column=datetime_column)
-        increment(f'{DD_KEY}.build_overview_time_series.start', tags)
+        increment(f"{DD_KEY}.build_overview_time_series.start", tags)
 
         if df_copy[datetime_column].count() <= 1:
             continue
 
         df_copy[datetime_column] = pd.to_datetime(
-            df[datetime_column], infer_datetime_format=True, errors='coerce'
+            df[datetime_column], infer_datetime_format=True, errors="coerce"
         )
         df_copy[datetime_column] = df_copy[datetime_column].view(int) / 10**9
 
@@ -262,11 +276,12 @@ def build_overview_data(
         y = []
 
         for bucket in buckets:
-            max_value = bucket['max_value']
-            min_value = bucket['min_value']
+            max_value = bucket["max_value"]
+            min_value = bucket["min_value"]
 
             df_filtered = df_copy[
-                (df_copy[datetime_column] >= min_value) & (df_copy[datetime_column] < max_value)
+                (df_copy[datetime_column] >= min_value)
+                & (df_copy[datetime_column] < max_value)
             ]
 
             x.append(
@@ -294,7 +309,7 @@ def build_overview_data(
             )
         )
 
-        increment(f'{DD_KEY}.build_overview_time_series.succeeded', tags)
+        increment(f"{DD_KEY}.build_overview_time_series.succeeded", tags)
 
     """
     Build sample data for scatter plot. Sample data consits of two parts:
@@ -306,7 +321,7 @@ def build_overview_data(
     else:
         df_sample = df.copy()
 
-    df_sample_numeric = df_sample[numeric_features].dropna(axis=1, how='all')
+    df_sample_numeric = df_sample[numeric_features].dropna(axis=1, how="all")
     """
     Calculate low cardinality categorical features:
     1. unique count <= SCATTER_PLOT_CATEGORY_LIMIT and unique count > 1
@@ -315,27 +330,34 @@ def build_overview_data(
     non_numeric_features = list(set(df.columns) - set(numeric_features))
     non_numeric_nuniques = df_sample[non_numeric_features].nunique()
     non_numeric_nuniques_filtered = non_numeric_nuniques[
-        (non_numeric_nuniques <= SCATTER_PLOT_CATEGORY_LIMIT) & (non_numeric_nuniques > 1)
+        (non_numeric_nuniques <= SCATTER_PLOT_CATEGORY_LIMIT)
+        & (non_numeric_nuniques > 1)
     ]
     non_numeric_counts = df_sample[non_numeric_features].count()
     sample_count = df_sample.shape[0]
-    non_numeric_counts_filtered = non_numeric_counts[non_numeric_counts > sample_count / 2]
+    non_numeric_counts_filtered = non_numeric_counts[
+        non_numeric_counts > sample_count / 2
+    ]
     eligible_category_features = set(non_numeric_nuniques_filtered.index) & set(
         non_numeric_counts_filtered.index
     )
     if len(eligible_category_features) > 0:
         encoder = MultipleColumnLabelEncoder(input_type=str)
-        df_sample_category = encoder.fit_transform(df_sample[eligible_category_features])
-        class_mappings = {k: list(e.label_classes()) for k, e in encoder.encoders.items()}
+        df_sample_category = encoder.fit_transform(
+            df_sample[eligible_category_features]
+        )
+        class_mappings = {
+            k: list(e.label_classes()) for k, e in encoder.encoders.items()
+        }
         df_sample_filtered = pd.concat([df_sample_numeric, df_sample_category], axis=1)
     else:
         class_mappings = dict()
         df_sample_filtered = df_sample_numeric
 
-    increment(f'{DD_KEY}.build_overview_data.succeeded')
+    increment(f"{DD_KEY}.build_overview_data.succeeded")
 
     return {
         DATA_KEY_TIME_SERIES: time_series,
-        DATA_KEY_SCATTER_PLOT: df_sample_filtered.to_dict('list'),
+        DATA_KEY_SCATTER_PLOT: df_sample_filtered.to_dict("list"),
         DATA_KEY_SCATTER_PLOT_LABELS: class_mappings,
     }