MindsDB 25.4.5.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/api/executor/sql_query/result_set.py

@@ -1,29 +1,70 @@
 import copy
-from typing import List, Optional
+from dataclasses import dataclass, field, MISSING
 
 import numpy as np
 import pandas as pd
+from pandas.api import types as pd_types
+import sqlalchemy.types as sqlalchemy_types
 
+from mindsdb_sql_parser.ast import TableColumn
+
+from mindsdb.utilities import log
 from mindsdb.api.executor.exceptions import WrongArgumentError
+from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
+
+
+logger = log.getLogger(__name__)
+
 
+def get_mysql_data_type_from_series(series: pd.Series, do_infer: bool = False) -> MYSQL_DATA_TYPE:
+    """Maps pandas Series data type to corresponding MySQL data type.
 
+    This function examines the dtype of a pandas Series and returns the appropriate
+    MySQL data type enum value. For object dtypes, it can optionally attempt to infer
+    a more specific type.
+
+    Args:
+        series (pd.Series): The pandas Series to determine the MySQL type for
+        do_infer (bool): If True and series has object dtype, attempt to infer a more specific type
+
+    Returns:
+        MYSQL_DATA_TYPE: The corresponding MySQL data type enum value
+    """
+    dtype = series.dtype
+    if pd_types.is_object_dtype(dtype) and do_infer is True:
+        dtype = series.infer_objects().dtype
+
+    if pd_types.is_object_dtype(dtype):
+        return MYSQL_DATA_TYPE.TEXT
+    if pd_types.is_datetime64_dtype(dtype):
+        return MYSQL_DATA_TYPE.DATETIME
+    if pd_types.is_string_dtype(dtype):
+        return MYSQL_DATA_TYPE.TEXT
+    if pd_types.is_bool_dtype(dtype):
+        return MYSQL_DATA_TYPE.BOOL
+    if pd_types.is_integer_dtype(dtype):
+        return MYSQL_DATA_TYPE.INT
+    if pd_types.is_numeric_dtype(dtype):
+        return MYSQL_DATA_TYPE.FLOAT
+    return MYSQL_DATA_TYPE.TEXT
+
+
+@dataclass(kw_only=True, slots=True)
 class Column:
-    def __init__(self, name=None, alias=None,
-                 table_name=None, table_alias=None,
-                 type=None, database=None, flags=None,
-                 charset=None):
-        if alias is None:
-            alias = name
-        if table_alias is None:
-            table_alias = table_name
-        self.name = name
-        self.alias = alias
-        self.table_name = table_name
-        self.table_alias = table_alias
-        self.type = type
-        self.database = database
-        self.flags = flags
-        self.charset = charset
+    name: str = field(default=MISSING)
+    alias: str | None = None
+    table_name: str | None = None
+    table_alias: str | None = None
+    type: MYSQL_DATA_TYPE | None = None
+    database: str | None = None
+    flags: dict = None
+    charset: str | None = None
+
+    def __post_init__(self):
+        if self.alias is None:
+            self.alias = self.name
+        if self.table_alias is None:
+            self.table_alias = self.table_name
 
     def get_hash_name(self, prefix):
         table_name = self.table_name if self.table_alias is None else self.table_alias
@@ -32,11 +73,8 @@ class Column:
         name = f'{prefix}_{table_name}_{name}'
         return name
 
-    def __repr__(self):
-        return f'{self.__class__.__name__}({self.__dict__})'
-
 
-def rename_df_columns(df: pd.DataFrame, names: Optional[List] = None) -> None:
+def rename_df_columns(df: pd.DataFrame, names: list | None = None) -> None:
     """Inplace rename of dataframe columns
 
     Args:
@@ -50,7 +88,15 @@ def rename_df_columns(df: pd.DataFrame, names: Optional[List] = None) -> None:
 
 
 class ResultSet:
-    def __init__(self, columns=None, values: List[List] = None, df: pd.DataFrame = None, affected_rows: int = None):
+    def __init__(
+        self,
+        columns: list[Column] | None = None,
+        values: list[list] | None = None,
+        df: pd.DataFrame | None = None,
+        affected_rows: int | None = None,
+        is_prediction: bool = False,
+        mysql_types: list[MYSQL_DATA_TYPE] | None = None
+    ):
         """
         Args:
             columns: list of Columns
@@ -62,15 +108,18 @@ class ResultSet:
             columns = []
         self._columns = columns
 
-        if values is None:
-            df = None
-        elif df is None:
-            df = pd.DataFrame(values)
+        if df is None:
+            if values is None:
+                df = None
+            else:
+                df = pd.DataFrame(values)
         self._df = df
 
         self.affected_rows = affected_rows
 
-        self.is_prediction = False
+        self.is_prediction = is_prediction
+
+        self.mysql_types = mysql_types
 
     def __repr__(self):
         col_names = ', '.join([col.name for col in self._columns])
@@ -89,43 +138,76 @@ class ResultSet:
 
     # --- converters ---
 
-    def from_df(self, df, database=None, table_name=None, table_alias=None):
-        self._columns = [
+    @classmethod
+    def from_df(
+        cls, df: pd.DataFrame, database=None, table_name=None, table_alias=None,
+        is_prediction: bool = False, mysql_types: list[MYSQL_DATA_TYPE] | None = None
+    ):
+        match mysql_types:
+            case None:
+                mysql_types = [None] * len(df.columns)
+            case list() if len(mysql_types) != len(df.columns):
+                raise WrongArgumentError(
+                    f'Mysql types length mismatch: {len(mysql_types)} != {len(df.columns)}'
+                )
+
+        columns = [
             Column(
                 name=column_name,
                 table_name=table_name,
                 table_alias=table_alias,
                 database=database,
-                type=column_dtype
-            ) for column_name, column_dtype
-            in zip(df.columns, df.dtypes)
+                type=mysql_type
+            ) for column_name, mysql_type
+            in zip(df.columns, mysql_types)
         ]
 
         rename_df_columns(df)
-        self._df = df
+        return cls(
+            df=df,
+            columns=columns,
+            is_prediction=is_prediction,
+            mysql_types=mysql_types
+        )
 
-        return self
+    @classmethod
+    def from_df_cols(cls, df: pd.DataFrame, columns_dict: dict[str, Column], strict: bool = True) -> 'ResultSet':
+        """Create ResultSet from dataframe and dictionary of columns
 
-    def from_df_cols(self, df, col_names, strict=True):
-        # find column by alias
-        alias_idx = {}
-        for col in col_names.values():
-            if col.alias is not None:
-                alias_idx[col.alias] = col
+        Args:
+            df (pd.DataFrame): dataframe
+            columns_dict (dict[str, Column]): dictionary of columns
+            strict (bool): if True, raise an error if a column is not found in columns_dict
 
-        for col in df.columns:
-            if col in col_names or strict:
-                column = col_names[col]
-            elif col in alias_idx:
-                column = alias_idx[col]
-            else:
-                column = Column(col)
-            self._columns.append(column)
+        Returns:
+            ResultSet: result set
+
+        Raises:
+            ValueError: if a column is not found in columns_dict and strict is True
+        """
+        alias_idx = {
+            column.alias: column
+            for column in columns_dict.values()
+            if column.alias is not None
+        }
+
+        columns = []
+        for column_name in df.columns:
+            if strict and column_name not in columns_dict:
+                raise ValueError(f'Column {column_name} not found in columns_dict')
+            column = (
+                columns_dict.get(column_name)
+                or alias_idx.get(column_name)
+                or Column(name=column_name)
+            )
+            columns.append(column)
 
         rename_df_columns(df)
-        self._df = df
 
-        return self
+        return cls(
+            columns=columns,
+            df=df
+        )
 
     def to_df(self):
         columns_names = self.get_column_names()
@@ -133,7 +215,7 @@ class ResultSet:
         rename_df_columns(df, columns_names)
         return df
 
-    def to_df_cols(self, prefix=''):
+    def to_df_cols(self, prefix: str = '') -> tuple[pd.DataFrame, dict[str, Column]]:
         # returns dataframe and dict of columns
         #   can be restored to ResultSet by from_df_cols method
 
@@ -262,6 +344,57 @@ class ResultSet:
         )
         self.add_raw_df(df)
 
+    def get_ast_columns(self) -> list[TableColumn]:
+        """Converts ResultSet columns to a list of TableColumn objects with SQLAlchemy types.
+
+        This method processes each column in the ResultSet, determines its MySQL data type
+        (inferring it if necessary), and maps it to the appropriate SQLAlchemy type.
+        The resulting TableColumn objects most likely will be used in CREATE TABLE statement.
+
+        Returns:
+            list[TableColumn]: A list of TableColumn objects with properly mapped SQLAlchemy types
+        """
+        columns: list[TableColumn] = []
+
+        type_mapping = {
+            MYSQL_DATA_TYPE.TINYINT: sqlalchemy_types.INTEGER,
+            MYSQL_DATA_TYPE.SMALLINT: sqlalchemy_types.INTEGER,
+            MYSQL_DATA_TYPE.MEDIUMINT: sqlalchemy_types.INTEGER,
+            MYSQL_DATA_TYPE.INT: sqlalchemy_types.INTEGER,
+            MYSQL_DATA_TYPE.BIGINT: sqlalchemy_types.INTEGER,
+            MYSQL_DATA_TYPE.YEAR: sqlalchemy_types.INTEGER,
+            MYSQL_DATA_TYPE.BOOL: sqlalchemy_types.BOOLEAN,
+            MYSQL_DATA_TYPE.BOOLEAN: sqlalchemy_types.BOOLEAN,
+            MYSQL_DATA_TYPE.FLOAT: sqlalchemy_types.FLOAT,
+            MYSQL_DATA_TYPE.DOUBLE: sqlalchemy_types.FLOAT,
+            MYSQL_DATA_TYPE.TIME: sqlalchemy_types.TIME,
+            MYSQL_DATA_TYPE.DATE: sqlalchemy_types.DATE,
+            MYSQL_DATA_TYPE.DATETIME: sqlalchemy_types.DATETIME,
+            MYSQL_DATA_TYPE.TIMESTAMP: sqlalchemy_types.TIMESTAMP,
+        }
+
+        for i, column in enumerate(self._columns):
+            column_type: MYSQL_DATA_TYPE | None = column.type
+
+            # infer MYSQL_DATA_TYPE if not set
+            if isinstance(column_type, MYSQL_DATA_TYPE) is False:
+                if column_type is not None:
+                    logger.warning(f'Unexpected column type: {column_type}')
+                if self._df is None:
+                    column_type = MYSQL_DATA_TYPE.TEXT
+                else:
+                    column_type = get_mysql_data_type_from_series(self._df.iloc[:, i])
+
+            sqlalchemy_type = type_mapping.get(column_type, sqlalchemy_types.TEXT)
+
+            columns.append(
+                TableColumn(
+                    name=column.alias,
+                    type=sqlalchemy_type
+                )
+            )
+        return columns
+
     def to_lists(self, json_types=False):
         """
         :param type_cast: cast numpy types
@@ -281,7 +414,7 @@ class ResultSet:
             return df.to_records(index=False).tolist()
 
         # slower but keep timestamp type
-        df = self._df.replace({np.nan: None})
+        df = self._df.replace({np.nan: None})  # TODO rework
         return df.to_dict('split')['data']
 
     def get_column_values(self, col_idx):

mindsdb/api/executor/sql_query/sql_query.py

@@ -252,7 +252,7 @@ class SQLQuery:
                 # return query info
                 # columns in upper case
                 rec = {k.upper(): v for k, v in self.run_query.get_info().items()}
-                self.fetched_data = ResultSet().from_df(pd.DataFrame([rec]))
+                self.fetched_data = ResultSet.from_df(pd.DataFrame([rec]))
                 self.columns_list = self.fetched_data.columns
                 return
             self.run_query.mark_as_run()

mindsdb/api/executor/sql_query/steps/apply_predictor_step.py

@@ -2,7 +2,6 @@ import datetime as dt
 import re
 
 import pandas as pd
-import dateparser
 
 from mindsdb_sql_parser.ast import (
     BinaryOperation,
@@ -92,21 +91,18 @@ class ApplyPredictorRowStepCall(ApplyPredictorBaseCall):
 
         table_name = get_preditor_alias(step, self.context.get('database'))
 
-        result = ResultSet()
-        result.is_prediction = True
         if len(predictions) == 0:
             columns_names = project_datanode.get_table_columns_names(predictor_name)
             predictions = pd.DataFrame([], columns=columns_names)
 
-        result.from_df(
-            predictions,
+        return ResultSet.from_df(
+            df=predictions,
             database=table_name[0],
             table_name=table_name[1],
-            table_alias=table_name[2]
+            table_alias=table_name[2],
+            is_prediction=True
         )
 
-        return result
-
 
 class ApplyPredictorStepCall(ApplyPredictorBaseCall):
 
@@ -165,15 +161,14 @@ class ApplyPredictorStepCall(ApplyPredictorBaseCall):
                 params['force_ts_infer'] = True
                 _mdb_forecast_offset = None
 
-            data.add_column(Column('__mdb_forecast_offset'), _mdb_forecast_offset)
+            data.add_column(Column(name='__mdb_forecast_offset'), _mdb_forecast_offset)
 
         table_name = get_preditor_alias(step, self.context['database'])
-        result = ResultSet()
-        result.is_prediction = True
 
         project_datanode = self.session.datahub.get(project_name)
         if len(data) == 0:
             columns_names = project_datanode.get_table_columns_names(predictor_name) + ['__mindsdb_row_id']
+            result = ResultSet(is_prediction=True)
             for column_name in columns_names:
                 result.add_column(Column(
                     name=column_name,
@@ -231,11 +226,12 @@ class ApplyPredictorStepCall(ApplyPredictorBaseCall):
                 pred_data = self.apply_ts_filter(pred_data, where_data, step, predictor_metadata)
                 predictions = pd.DataFrame(pred_data)
 
-            result.from_df(
+            result = ResultSet.from_df(
                 predictions,
                 database=table_name[0],
                 table_name=table_name[1],
-                table_alias=table_name[2]
+                table_alias=table_name[2],
+                is_prediction=True
             )
 
         return result
@@ -283,6 +279,8 @@ class ApplyPredictorStepCall(ApplyPredictorBaseCall):
             # Use dateparser as fallback and infer format
             try:
                 # Parse the first sample to get its format
+                # The import is heavy, so we do it here on-demand
+                import dateparser
                 parsed_date = dateparser.parse(samples[0])
                 if parsed_date is None:
                     raise ValueError("Could not parse date")

mindsdb/api/executor/sql_query/steps/fetch_dataframe.py

@@ -7,11 +7,12 @@ from mindsdb_sql_parser.ast import (
     BinaryOperation,
     Tuple,
 )
-from mindsdb.api.executor.planner.steps import FetchDataframeStep
-from mindsdb.integrations.utilities.query_traversal import query_traversal
 
+from mindsdb.api.executor.planner.steps import FetchDataframeStep
+from mindsdb.api.executor.datahub.classes.response import DataHubResponse
 from mindsdb.api.executor.sql_query.result_set import ResultSet
 from mindsdb.api.executor.exceptions import UnknownError
+from mindsdb.integrations.utilities.query_traversal import query_traversal
 from mindsdb.interfaces.query_context.context_controller import query_context_controller
 
 from .base import BaseStepCall
@@ -89,7 +90,7 @@ class FetchDataframeStepCall(BaseStepCall):
             table_alias = (self.context.get('database'), 'result', 'result')
 
             # fetch raw_query
-            response = dn.query(
+            response: DataHubResponse = dn.query(
                 native_query=step.raw_query,
                 session=self.session
             )
@@ -105,7 +106,7 @@ class FetchDataframeStepCall(BaseStepCall):
 
             query, context_callback = query_context_controller.handle_db_context_vars(query, dn, self.session)
 
-            response = dn.query(
+            response: DataHubResponse = dn.query(
                 query=query,
                 session=self.session
             )
@@ -114,13 +115,10 @@ class FetchDataframeStepCall(BaseStepCall):
             if context_callback:
                 context_callback(df, response.columns)
 
-        result = ResultSet()
-
-        result.from_df(
+        return ResultSet.from_df(
             df,
             table_name=table_alias[1],
             table_alias=table_alias[2],
-            database=table_alias[0]
+            database=table_alias[0],
+            mysql_types=response.mysql_types
         )
-
-        return result

mindsdb/api/executor/sql_query/steps/fetch_dataframe_partition.py

@@ -8,7 +8,7 @@ from mindsdb.integrations.utilities.query_traversal import query_traversal
 from mindsdb.interfaces.query_context.context_controller import RunningQuery
 from mindsdb.api.executor.sql_query.result_set import ResultSet
 from mindsdb.utilities import log
-from mindsdb.utilities.config import Config
+from mindsdb.utilities.config import config
 from mindsdb.utilities.partitioning import get_max_thread_count, split_data_frame
 from mindsdb.api.executor.sql_query.steps.fetch_dataframe import get_table_alias, get_fill_param_fnc
 from mindsdb.utilities.context_executor import ContextThreadPoolExecutor
@@ -64,8 +64,6 @@ class FetchDataframePartitionCall(BaseStepCall):
         self.current_step_num = step.step_num
         self.substeps = step.steps
 
-        config = Config()
-
         # ml task queue enabled?
         use_threads, thread_count = False, None
         if config['ml_task_queue']['type'] == 'redis':
@@ -95,21 +93,8 @@ class FetchDataframePartitionCall(BaseStepCall):
         """
 
         results = []
-        while True:
-
-            # fetch batch
-            query2 = run_query.get_partition_query(self.current_step_num, query)
-            response = self.dn.query(
-                query=query2,
-                session=self.session
-            )
-            df = response.data_frame
-
-            if df is None or len(df) == 0:
-                break
 
-            # executing of sub steps can modify dataframe columns, lets memorise max tracking value
-            max_track_value = run_query.get_max_track_value(df)
+        for df in run_query.get_partitions(self.dn, self, query):
             try:
                 sub_data = self.exec_sub_steps(df)
                 results.append(sub_data)
@@ -119,8 +104,6 @@ class FetchDataframePartitionCall(BaseStepCall):
                 else:
                     raise e
 
-            run_query.set_progress(df, max_track_value)
-
         return self.concat_results(results)
 
     def concat_results(self, results: List[ResultSet]) -> ResultSet:
@@ -135,7 +118,7 @@ class FetchDataframePartitionCall(BaseStepCall):
 
         data = ResultSet()
         if len(df_list) > 0:
-            data.from_df_cols(pd.concat(df_list), col_names)
+            data = ResultSet.from_df_cols(pd.concat(df_list), col_names)
 
         return data
 
@@ -147,10 +130,7 @@ class FetchDataframePartitionCall(BaseStepCall):
         - substep are executed using result of previos step (like it is all fetched data is available)
         - the final result is returned and used outside to concatenate with results of other's batches
         """
-
-        input_data = ResultSet()
-
-        input_data.from_df(
+        input_data = ResultSet.from_df(
             df,
             table_name=self.table_alias[1],
             table_alias=self.table_alias[2],
@@ -191,22 +171,7 @@ class FetchDataframePartitionCall(BaseStepCall):
 
         with ContextThreadPoolExecutor(max_workers=thread_count) as executor:
 
-            while True:
-                # fetch batch
-                query2 = run_query.get_partition_query(self.current_step_num, query)
-                response = self.dn.query(
-                    query=query2,
-                    session=self.session
-                )
-                df = response.data_frame
-
-                if df is None or len(df) == 0:
-                    # TODO detect circles: data handler ignores condition and output is repeated
-
-                    # exit & stop workers
-                    break
-
-                max_track_value = run_query.get_max_track_value(df)
+            for df in run_query.get_partitions(self.dn, self, query):
 
                 # split into chunks and send to workers
                 futures = []
@@ -225,9 +190,5 @@ class FetchDataframePartitionCall(BaseStepCall):
                 if self.sql_query.stop_event is not None and self.sql_query.stop_event.is_set():
                     executor.shutdown()
                     raise RuntimeError('Query is interrupted')
-                # TODO
-                #  1. get next batch without updating track_value:
-                #    it allows to keep queue_in filled with data between fetching batches
-                run_query.set_progress(df, max_track_value)
 
         return self.concat_results(results)

mindsdb/api/executor/sql_query/steps/insert_step.py

@@ -1,17 +1,18 @@
 from mindsdb_sql_parser.ast import (
     Identifier,
 )
+
 from mindsdb.api.executor.planner.steps import (
     SaveToTable,
     InsertToTable,
     CreateTableStep
 )
-
 from mindsdb.api.executor.sql_query.result_set import ResultSet, Column
 from mindsdb.api.executor.exceptions import (
     NotSupportedYet,
     LogicError
 )
+from mindsdb.integrations.libs.response import INF_SCHEMA_COLUMNS_NAMES
 
 from .base import BaseStepCall
 
@@ -30,12 +31,32 @@ class InsertToTableCall(BaseStepCall):
             if step.is_replace:
                 is_replace = True
 
+        if len(step.table.parts) > 1:
+            integration_name = step.table.parts[0]
+            table_name = Identifier(parts=step.table.parts[1:])
+        else:
+            integration_name = self.context['database']
+            table_name = step.table
+
+        dn = self.session.datahub.get(integration_name)
+
+        if hasattr(dn, 'create_table') is False:
+            raise NotSupportedYet(f"Creating table in '{integration_name}' is not supported")
+
         if step.dataframe is not None:
             data = self.steps_data[step.dataframe.result.step_num]
         elif step.query is not None:
             data = ResultSet()
-            for col in step.query.columns:
-                data.add_column(Column(col.name))
+            if step.query.columns is None:
+                # Is query like: INSERT INTO table VALUES (...)
+                table_columns_df = dn.get_table_columns_df(str(table_name))
+                columns_names = table_columns_df[INF_SCHEMA_COLUMNS_NAMES.COLUMN_NAME].to_list()
+                for column_name in columns_names:
+                    data.add_column(Column(name=column_name))
+            else:
+                # Is query like: INSERT INTO table (column_name, ...) VALUES (...)
+                for col in step.query.columns:
+                    data.add_column(Column(name=col.name))
 
             records = []
             for row in step.query.values:
@@ -53,18 +74,6 @@ class InsertToTableCall(BaseStepCall):
         else:
             raise LogicError(f'Data not found for insert: {step}')
 
-        if len(step.table.parts) > 1:
-            integration_name = step.table.parts[0]
-            table_name = Identifier(parts=step.table.parts[1:])
-        else:
-            integration_name = self.context['database']
-            table_name = step.table
-
-        dn = self.session.datahub.get(integration_name)
-
-        if hasattr(dn, 'create_table') is False:
-            raise NotSupportedYet(f"Creating table in '{integration_name}' is not supported")
-
         #  del 'service' columns
         for col in data.find_columns('__mindsdb_row_id'):
             data.del_column(col)

mindsdb/api/executor/sql_query/steps/join_step.py

@@ -101,7 +101,7 @@ class JoinStepCall(BaseStepCall):
         resp_df.replace({np.nan: None}, inplace=True)
 
         names_a.update(names_b)
-        data = ResultSet().from_df_cols(resp_df, col_names=names_a)
+        data = ResultSet.from_df_cols(df=resp_df, columns_dict=names_a)
 
         for col in data.find_columns('__mindsdb_row_id'):
             data.del_column(col)

mindsdb/api/executor/sql_query/steps/project_step.py

@@ -83,4 +83,4 @@ class ProjectStepCall(BaseStepCall):
 
         res = query_df(df, query, session=self.session)
 
-        return ResultSet().from_df_cols(res, col_names, strict=False)
+        return ResultSet.from_df_cols(df=res, columns_dict=col_names, strict=False)

mindsdb/api/executor/sql_query/steps/sql_steps.py

@@ -38,4 +38,4 @@ class DataStepCall(BaseStepCall):
     def call(self, step):
         # create resultset
         df = pd.DataFrame(step.data)
-        return ResultSet().from_df(df, database='', table_name='')
+        return ResultSet.from_df(df, database='', table_name='')