MindsDB 25.5.3.0__py3-none-any.whl → 25.5.4.1__py3-none-any.whl

mindsdb/api/executor/datahub/classes/response.py

@@ -1,11 +1,14 @@
 from dataclasses import dataclass, field
-from typing import Optional, List, Dict
+from typing import List, Dict
 
 import pandas as pd
 
+from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
+
 
 @dataclass
 class DataHubResponse:
     data_frame: pd.DataFrame = field(default_factory=pd.DataFrame)
     columns: List[Dict] = field(default_factory=list)
-    affected_rows: Optional[int] = None
+    affected_rows: int | None = None
+    mysql_types: list[MYSQL_DATA_TYPE] | None = None

mindsdb/api/executor/datahub/datanodes/integration_datanode.py

@@ -1,13 +1,12 @@
 import time
 import inspect
 from dataclasses import astuple
+from typing import Iterable
 
 import numpy as np
-from numpy import dtype as np_dtype
 import pandas as pd
-from pandas.api import types as pd_types
 from sqlalchemy.types import (
-    Integer, Float, Text
+    Integer, Float
 )
 
 from mindsdb_sql_parser.ast.base import ASTNode
@@ -121,9 +120,7 @@ class IntegrationDataNode(DataNode):
             tables=[name],
             if_exists=if_exists
         )
-        result = self._query(drop_ast)
-        if result.type == RESPONSE_TYPE.ERROR:
-            raise Exception(result.error_message)
+        self.query(drop_ast)
 
     def create_table(self, table_name: Identifier, result_set: ResultSet = None, columns=None,
                      is_replace=False, is_create=False, **kwargs) -> DataHubResponse:
@@ -138,27 +135,11 @@ class IntegrationDataNode(DataNode):
 
             df = result_set.get_raw_df()
 
-            for idx, col in enumerate(result_set.columns):
-                dtype = col.type
-                # assume this is pandas type
-                column_type = Text
-                if isinstance(dtype, np_dtype):
-                    if pd_types.is_object_dtype(dtype):
-                        # try to infer
-                        dtype = df[idx].infer_objects().dtype
-
-                    if pd_types.is_integer_dtype(dtype):
-                        column_type = Integer
-                    elif pd_types.is_numeric_dtype(dtype):
-                        column_type = Float
-
-                columns.append(
-                    TableColumn(
-                        name=col.alias,
-                        type=column_type
-                    )
-                )
-                table_columns_meta[col.alias] = column_type
+            columns: list[TableColumn] = result_set.get_ast_columns()
+            table_columns_meta = {
+                column.name: column.type
+                for column in columns
+            }
 
         if is_replace:
             # drop
@@ -166,9 +147,7 @@ class IntegrationDataNode(DataNode):
                 tables=[table_name],
                 if_exists=True
             )
-            result = self._query(drop_ast)
-            if result.type == RESPONSE_TYPE.ERROR:
-                raise Exception(result.error_message)
+            self.query(drop_ast)
             is_create = True
 
         if is_create:
@@ -177,9 +156,7 @@ class IntegrationDataNode(DataNode):
                 columns=columns,
                 is_replace=is_replace
             )
-            result = self._query(create_table_ast)
-            if result.type == RESPONSE_TYPE.ERROR:
-                raise Exception(result.error_message)
+            self.query(create_table_ast)
 
         if result_set is None:
             # it is just a 'create table'
@@ -224,56 +201,44 @@ class IntegrationDataNode(DataNode):
         )
 
         try:
-            result = self._query(insert_ast)
+            result: DataHubResponse = self.query(insert_ast)
         except Exception as e:
             msg = f'[{self.ds_type}/{self.integration_name}]: {str(e)}'
             raise DBHandlerException(msg) from e
 
-        if result.type == RESPONSE_TYPE.ERROR:
-            raise Exception(result.error_message)
-
         return DataHubResponse(affected_rows=result.affected_rows)
 
-    def _query(self, query) -> HandlerResponse:
-        time_before_query = time.perf_counter()
-        result = self.integration_handler.query(query)
-        elapsed_seconds = time.perf_counter() - time_before_query
-        query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
-            get_class_name(self.integration_handler), result.type)
-        query_time_with_labels.observe(elapsed_seconds)
-
-        num_rows = 0
-        if result.data_frame is not None:
-            num_rows = len(result.data_frame.index)
-        response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
-            get_class_name(self.integration_handler), result.type)
-        response_size_with_labels.observe(num_rows)
-        return result
-
-    def _native_query(self, native_query) -> HandlerResponse:
-        time_before_query = time.perf_counter()
-        result = self.integration_handler.native_query(native_query)
-        elapsed_seconds = time.perf_counter() - time_before_query
-        query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
-            get_class_name(self.integration_handler), result.type)
-        query_time_with_labels.observe(elapsed_seconds)
-
-        num_rows = 0
-        if result.data_frame is not None:
-            num_rows = len(result.data_frame.index)
-        response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
-            get_class_name(self.integration_handler), result.type)
-        response_size_with_labels.observe(num_rows)
-        return result
+    def has_support_stream(self) -> bool:
+        # checks if data handler has query_stream method
+        return hasattr(self.integration_handler, 'query_stream') and callable(self.integration_handler.query_stream)
+
+    @profiler.profile()
+    def query_stream(self, query: ASTNode, fetch_size: int = None) -> Iterable:
+        # returns generator of results from handler (split by chunks)
+        return self.integration_handler.query_stream(query, fetch_size=fetch_size)
 
     @profiler.profile()
     def query(self, query: ASTNode | None = None, native_query: str | None = None, session=None) -> DataHubResponse:
         try:
+            time_before_query = time.perf_counter()
             if query is not None:
-                result: HandlerResponse = self._query(query)
+                result: HandlerResponse = self.integration_handler.query(query)
             else:
                 # try to fetch native query
-                result: HandlerResponse = self._native_query(native_query)
+                result: HandlerResponse = self.integration_handler.native_query(native_query)
+
+            # metrics
+            elapsed_seconds = time.perf_counter() - time_before_query
+            query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
+                get_class_name(self.integration_handler), result.type)
+            query_time_with_labels.observe(elapsed_seconds)
+
+            num_rows = 0
+            if result.data_frame is not None:
+                num_rows = len(result.data_frame.index)
+            response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
+                get_class_name(self.integration_handler), result.type)
+            response_size_with_labels.observe(num_rows)
         except Exception as e:
             msg = str(e).strip()
             if msg == '':
@@ -294,7 +259,8 @@ class IntegrationDataNode(DataNode):
 
         try:
             # replace python's Nan, np.NaN, np.nan and pd.NA to None
-            df.replace([np.NaN, pd.NA], None, inplace=True)
+            # TODO keep all NAN to the end of processing, bacause replacing also changes dtypes
+            df.replace([np.NaN, pd.NA, pd.NaT], None, inplace=True)
         except Exception as e:
             logger.error(f"Issue with clearing DF from NaN values: {e}")
         # endregion
@@ -310,5 +276,6 @@ class IntegrationDataNode(DataNode):
         return DataHubResponse(
             data_frame=df,
             columns=columns_info,
-            affected_rows=result.affected_rows
+            affected_rows=result.affected_rows,
+            mysql_types=result.mysql_types
         )

mindsdb/api/executor/planner/query_planner.py

@@ -848,11 +848,12 @@ class QueryPlanner:
         """
 
         # handle fetchdataframe partitioning
+        steps_in = plan.steps
         steps_out = []
 
         step = None
         partition_step = None
-        for step in plan.steps:
+        for step in steps_in:
             if isinstance(step, FetchDataframeStep) and step.params is not None:
                 batch_size = step.params.get('batch_size')
                 if batch_size is not None:
@@ -905,6 +906,14 @@ class QueryPlanner:
 
         if plan.is_resumable and isinstance(step, InsertToTable):
             plan.is_async = True
+        else:
+            # special case: register insert from select (it is the same as mark resumable)
+            if (
+                len(steps_in) == 2
+                and isinstance(steps_in[0], FetchDataframeStep)
+                and isinstance(steps_in[1], InsertToTable)
+            ):
+                plan.is_resumable = True
 
         plan.steps = steps_out
         return plan

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

@@ -91,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):
 
@@ -164,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,
@@ -230,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

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