MindsDB 25.4.1.0__py3-none-any.whl → 25.4.2.1__py3-none-any.whl

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

@@ -14,6 +14,7 @@ from mindsdb_sql_parser.ast import (
 from mindsdb.utilities.exception import EntityNotExistsError
 from mindsdb.api.executor.datahub.datanodes.datanode import DataNode
 from mindsdb.api.executor.datahub.classes.tables_row import TablesRow
+from mindsdb.api.executor.datahub.classes.response import DataHubResponse
 from mindsdb.utilities.partitioning import process_dataframe_in_partitions
 
 
@@ -45,10 +46,6 @@ class ProjectDataNode(DataNode):
         result = [TablesRow.from_dict(row) for row in tables]
         return result
 
-    def has_table(self, table_name):
-        tables = self.project.get_tables()
-        return table_name in tables
-
     def get_table_columns(self, table_name, schema_name=None):
         return [
             {'name': name}
@@ -71,7 +68,7 @@ class ProjectDataNode(DataNode):
 
         return ml_handler.predict(model_name, df, project_name=self.project.name, version=version, params=params)
 
-    def query(self, query=None, native_query=None, session=None):
+    def query(self, query=None, native_query=None, session=None) -> DataHubResponse:
         if query is None and native_query is not None:
             query = parse_sql(native_query)
 
@@ -81,7 +78,7 @@ class ProjectDataNode(DataNode):
             if kb_table:
                 # this is the knowledge db
                 kb_table.update_query(query)
-                return pd.DataFrame(), []
+                return DataHubResponse()
 
             raise NotImplementedError(f"Can't update object: {query_table}")
 
@@ -91,7 +88,7 @@ class ProjectDataNode(DataNode):
             if kb_table:
                 # this is the knowledge db
                 kb_table.delete_query(query)
-                return pd.DataFrame(), []
+                return DataHubResponse()
 
             raise NotImplementedError(f"Can't delete object: {query_table}")
 
@@ -111,8 +108,7 @@ class ProjectDataNode(DataNode):
                         new_query.where,
                         project_filter
                     ])
-                df, columns_info = self.information_schema.query(new_query)
-                return df, columns_info
+                return self.information_schema.query(new_query)
             # endregion
 
             # other table from project
@@ -121,15 +117,15 @@ class ProjectDataNode(DataNode):
                 # this is the view
                 df = self.project.query_view(query, session)
 
-                columns_info = [
-                    {
-                        'name': k,
-                        'type': v
-                    }
-                    for k, v in df.dtypes.items()
-                ]
+                columns_info = [{
+                    'name': k,
+                    'type': v
+                } for k, v in df.dtypes.items()]
 
-                return df, columns_info
+                return DataHubResponse(
+                    data_frame=df,
+                    columns=columns_info
+                )
 
             kb_table = session.kb_controller.get_table(query_table, self.project.id)
             if kb_table:
@@ -143,13 +139,16 @@ class ProjectDataNode(DataNode):
                     for k, v in df.dtypes.items()
                 ]
 
-                return df, columns_info
+                return DataHubResponse(
+                    data_frame=df,
+                    columns=columns_info
+                )
 
             raise EntityNotExistsError(f"Can't select from {query_table} in project")
         else:
             raise NotImplementedError(f"Query not supported {query}")
 
-    def create_table(self, table_name: Identifier, result_set=None, is_replace=False, **kwargs):
+    def create_table(self, table_name: Identifier, result_set=None, is_replace=False, **kwargs) -> DataHubResponse:
         # is_create - create table
         # is_replace - drop table if exists
         # is_create==False and is_replace==False: just insert
@@ -165,5 +164,6 @@ class ProjectDataNode(DataNode):
                 kb_table.clear()
 
             df = result_set.to_df()
-            return kb_table.insert(df)
+            kb_table.insert(df)
+            return DataHubResponse()
         raise NotImplementedError(f"Can't create table {table_name}")

mindsdb/api/executor/planner/plan_join.py

@@ -119,7 +119,7 @@ class PlanJoinTablesQuery:
             query2.from_table = None
             query2.using = None
             query2.cte = None
-            sup_select = QueryStep(query2, from_table=join_step.result)
+            sup_select = QueryStep(query2, from_table=join_step.result, strict_where=False)
             self.planner.plan.add_step(sup_select)
             return sup_select
         return join_step
@@ -423,7 +423,7 @@ class PlanJoinTablesQuery:
             else:
                 query2.where = cond
 
-        step = self.planner.get_integration_select_step(query2)
+        step = self.planner.get_integration_select_step(query2, params=query_in.using)
         self.tables_fetch_step[item.index] = step
 
         self.add_plan_step(step)

mindsdb/api/executor/planner/query_plan.py

@@ -2,6 +2,7 @@
 class QueryPlan:
     def __init__(self, steps=None, **kwargs):
         self.steps = []
+        self.is_resumable = False
 
         if steps:
             for step in steps:

mindsdb/api/executor/planner/query_planner.py

@@ -12,14 +12,13 @@ from mindsdb.api.executor.planner.exceptions import PlanningException
 from mindsdb.api.executor.planner import utils
 from mindsdb.api.executor.planner.query_plan import QueryPlan
 from mindsdb.api.executor.planner.steps import (
-    FetchDataframeStep, ProjectStep, ApplyPredictorStep,
+    PlanStep, FetchDataframeStep, ProjectStep, ApplyPredictorStep,
     ApplyPredictorRowStep, UnionStep, GetPredictorColumns, SaveToTable,
-    InsertToTable, UpdateToTable, SubSelectStep, QueryStep,
-    DeleteStep, DataStep, CreateTableStep
+    InsertToTable, UpdateToTable, SubSelectStep, QueryStep, JoinStep,
+    DeleteStep, DataStep, CreateTableStep, FetchDataframeStepPartition
 )
 from mindsdb.api.executor.planner.utils import (
     disambiguate_predictor_column_identifier,
-    get_deepest_select,
     recursively_extract_column_values,
     query_traversal, filters_to_bin_op
 )
@@ -166,7 +165,11 @@ class QueryPlanner:
 
         query_traversal(query, _prepare_integration_select)
 
-    def get_integration_select_step(self, select):
+    def get_integration_select_step(self, select: Select, params: dict = None) -> PlanStep:
+        """
+        Generate planner step to execute query over integration or over results of previous step (if it is CTE)
+        """
+
         if isinstance(select.from_table, NativeQuery):
             integration_name = select.from_table.integration.parts[-1]
         else:
@@ -188,12 +191,22 @@ class QueryPlanner:
         if fetch_df_select.using is not None:
             fetch_df_select.using = None
 
-        return FetchDataframeStep(integration=integration_name, query=fetch_df_select)
+        if params:
+            fetch_params = params.copy()
+            # remove partition parameters
+            for key in ('batch_size', 'track_column'):
+                if key in params:
+                    del params[key]
+            if 'track_column' in fetch_params and isinstance(fetch_params['track_column'], Identifier):
+                fetch_params['track_column'] = fetch_params['track_column'].parts[-1]
+        else:
+            fetch_params = None
+        return FetchDataframeStep(integration=integration_name, query=fetch_df_select, params=fetch_params)
 
     def plan_integration_select(self, select):
         """Plan for a select query that can be fully executed in an integration"""
 
-        return self.plan.add_step(self.get_integration_select_step(select))
+        return self.plan.add_step(self.get_integration_select_step(select, params=select.using))
 
     def resolve_database_table(self, node: Identifier):
         # resolves integration name and table name
@@ -414,12 +427,6 @@ class QueryPlanner:
 
         return self.plan_mdb_nested_select(select)
 
-    def plan_integration_nested_select(self, select, integration_name):
-        fetch_df_select = copy.deepcopy(select)
-        deepest_select = get_deepest_select(fetch_df_select)
-        self.prepare_integration_select(integration_name, deepest_select)
-        return self.plan.add_step(FetchDataframeStep(integration=integration_name, query=fetch_df_select))
-
     def plan_mdb_nested_select(self, select):
         # plan nested select
 
@@ -818,7 +825,72 @@ class QueryPlanner:
         else:
             raise PlanningException(f'Unsupported query type {type(query)}')
 
-        return self.plan
+        plan = self.handle_partitioning(self.plan)
+
+        return plan
+
+    def handle_partitioning(self, plan: QueryPlan) -> QueryPlan:
+        """
+        If plan has fetching in partitions:
+          try to rebuild plan to send fetched chunk of data through the following steps, if it is possible
+        """
+
+        # handle fetchdataframe partitioning
+        steps_out = []
+
+        partition_step = None
+        for step in plan.steps:
+            if isinstance(step, FetchDataframeStep) and step.params is not None:
+                batch_size = step.params.get('batch_size')
+                if batch_size is not None:
+                    # found batched fetch
+                    partition_step = FetchDataframeStepPartition(
+                        step_num=step.step_num,
+                        integration=step.integration,
+                        query=step.query,
+                        raw_query=step.raw_query,
+                        params=step.params
+                    )
+                    steps_out.append(partition_step)
+                    # mark plan
+                    plan.is_resumable = True
+                    continue
+                else:
+                    step.params = None
+
+            if partition_step is not None:
+                # check and add step into partition
+
+                can_be_partitioned = False
+                if isinstance(step, (JoinStep, ApplyPredictorStep, InsertToTable)):
+                    can_be_partitioned = True
+                elif isinstance(step, QueryStep):
+                    query = step.query
+                    if (
+                        query.group_by is None and query.order_by is None and query.distinct is False
+                        and query.limit is None and query.offset is None
+                    ):
+                        no_identifiers = [
+                            target
+                            for target in step.query.targets
+                            if not isinstance(target, (Star, Identifier))
+                        ]
+                        if len(no_identifiers) == 0:
+                            can_be_partitioned = True
+
+                if not can_be_partitioned:
+                    if len(partition_step.steps) == 0:
+                        # Nothing can be partitioned, failback to old plan
+                        plan.is_resumable = False
+                        return plan
+                    partition_step = None
+                else:
+                    partition_step.steps.append(step)
+                    continue
+
+            steps_out.append(step)
+        plan.steps = steps_out
+        return plan
 
     def prepare_steps(self, query):
         statement_planner = PreparedStatementPlanner(self)

mindsdb/api/executor/planner/steps.py

@@ -104,11 +104,19 @@ class LimitOffsetStep(PlanStep):
 
 class FetchDataframeStep(PlanStep):
     """Fetches a dataframe from external integration"""
-    def __init__(self, integration, query=None, raw_query=None, *args, **kwargs):
+    def __init__(self, integration, query=None, raw_query=None, params=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.integration = integration
         self.query = query
         self.raw_query = raw_query
+        self.params = params
+
+
+class FetchDataframeStepPartition(FetchDataframeStep):
+    """Fetches a dataframe from external integration in partitions"""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.steps = []
 
 
 class ApplyPredictorStep(PlanStep):
@@ -249,11 +257,12 @@ class SubSelectStep(PlanStep):
 
 
 class QueryStep(PlanStep):
-    def __init__(self, query, from_table=None, *args, **kwargs):
+    def __init__(self, query, from_table=None, *args, strict_where=True, **kwargs):
         """Performs query using injected dataframe"""
         super().__init__(*args, **kwargs)
         self.query = query
         self.from_table = from_table
+        self.strict_where = strict_where
 
 
 class DataStep(PlanStep):

mindsdb/api/executor/sql_query/result_set.py

@@ -50,13 +50,14 @@ 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):
-        '''
-
-        :param columns: list of Columns
-        :param values: data of resultSet, have to be list of lists with length equal to column
-        :param df: injected dataframe, have to have enumerated columns and length equal to columns
-        '''
+    def __init__(self, columns=None, values: List[List] = None, df: pd.DataFrame = None, affected_rows: int = None):
+        """
+        Args:
+            columns: list of Columns
+            values (List[List]): data of resultSet, have to be list of lists with length equal to column
+            df (pd.DataFrame): injected dataframe, have to have enumerated columns and length equal to columns
+            affected_rows (int): number of affected rows
+        """
         if columns is None:
             columns = []
         self._columns = columns
@@ -67,6 +68,8 @@ class ResultSet:
             df = pd.DataFrame(values)
         self._df = df
 
+        self.affected_rows = affected_rows
+
         self.is_prediction = False
 
     def __repr__(self):

mindsdb/api/executor/sql_query/sql_query.py

@@ -8,11 +8,11 @@
  * permission of MindsDB Inc
  *******************************************************
 """
-import re
 import inspect
 from textwrap import dedent
+from typing import Union, Dict
 
-from mindsdb_sql_parser import parse_sql
+from mindsdb_sql_parser import parse_sql, ASTNode
 from mindsdb.api.executor.planner.steps import (
     ApplyTimeseriesPredictorStep,
     ApplyPredictorRowStep,
@@ -23,7 +23,7 @@ from mindsdb.api.executor.planner.exceptions import PlanningException
 from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
 from mindsdb.api.executor.planner import query_planner
 
-from mindsdb.api.executor.utilities.sql import query_df, get_query_models
+from mindsdb.api.executor.utilities.sql import get_query_models
 from mindsdb.interfaces.model.functions import get_model_record
 from mindsdb.api.executor.exceptions import (
     BadTableError,
@@ -33,19 +33,21 @@ from mindsdb.api.executor.exceptions import (
 import mindsdb.utilities.profiler as profiler
 from mindsdb.utilities.fs import create_process_mark, delete_process_mark
 from mindsdb.utilities.exception import EntityNotExistsError
+from mindsdb.interfaces.query_context.context_controller import query_context_controller
+from mindsdb.utilities.context import context as ctx
+
 
 from . import steps
 from .result_set import ResultSet, Column
 from . steps.base import BaseStepCall
 
-superset_subquery = re.compile(r'from[\s\n]*(\(.*\))[\s\n]*as[\s\n]*virtual_table', flags=re.IGNORECASE | re.MULTILINE | re.S)
-
 
 class SQLQuery:
 
     step_handlers = {}
 
-    def __init__(self, sql, session, execute=True, database=None):
+    def __init__(self, sql: Union[ASTNode, str], session, execute: bool = True,
+                 database: str = None, query_id: int = None):
         self.session = session
 
         if database is not None:
@@ -59,23 +61,22 @@ class SQLQuery:
         }
 
         self.columns_list = None
-        self.steps_data = {}
+        self.steps_data: Dict[int, ResultSet] = {}
 
-        self.planner = None
+        self.planner: query_planner.QueryPlanner = None
         self.parameters = []
-        self.fetched_data = None
+        self.fetched_data: ResultSet = None
 
         self.outer_query = None
+        self.run_query = None
+        self.query_id = query_id
+        if query_id is not None:
+            # resume query
+            run_query = query_context_controller.get_query(self.query_id)
+            run_query.clear_error()
+            sql = run_query.sql
 
         if isinstance(sql, str):
-            # region workaround for subqueries in superset
-            if 'as virtual_table' in sql.lower():
-                subquery = re.findall(superset_subquery, sql)
-                if isinstance(subquery, list) and len(subquery) == 1:
-                    subquery = subquery[0]
-                    self.outer_query = sql.replace(subquery, 'dataframe')
-                    sql = subquery.strip('()')
-            # endregion
             self.query = parse_sql(sql)
             self.context['query_str'] = sql
         else:
@@ -89,7 +90,6 @@ class SQLQuery:
         self.create_planner()
 
         if execute:
-            self.prepare_query(prepare=False)
             self.execute_query()
 
     @classmethod
@@ -190,63 +190,62 @@ class SQLQuery:
             default_namespace=database,
         )
 
-    def fetch(self, view='result_set'):
-        data = self.fetched_data
-
-        if view == 'dataframe':
-            result = data.to_df()
-        else:
-            result = data
-
-        return {
-            'success': True,
-            'result': result
-        }
+    def prepare_query(self):
+        """it is prepared statement call
+        """
+        try:
+            for step in self.planner.prepare_steps(self.query):
+                data = self.execute_step(step)
+                step.set_result(data)
+                self.steps_data[step.step_num] = data
+        except PlanningException as e:
+            raise LogicError(e)
 
-    def prepare_query(self, prepare=True):
-        if prepare:
-            # it is prepared statement call
-            try:
-                for step in self.planner.prepare_steps(self.query):
-                    data = self.execute_step(step)
-                    step.set_result(data)
-                    self.steps_data[step.step_num] = data
-            except PlanningException as e:
-                raise LogicError(e)
-
-            statement_info = self.planner.get_statement_info()
-
-            self.columns_list = []
-            for col in statement_info['columns']:
-                self.columns_list.append(
-                    Column(
-                        database=col['ds'],
-                        table_name=col['table_name'],
-                        table_alias=col['table_alias'],
-                        name=col['name'],
-                        alias=col['alias'],
-                        type=col['type']
-                    )
-                )
+        statement_info = self.planner.get_statement_info()
 
-            self.parameters = [
+        self.columns_list = []
+        for col in statement_info['columns']:
+            self.columns_list.append(
                 Column(
+                    database=col['ds'],
+                    table_name=col['table_name'],
+                    table_alias=col['table_alias'],
                     name=col['name'],
                     alias=col['alias'],
                     type=col['type']
                 )
-                for col in statement_info['parameters']
-            ]
-
-    def execute_query(self, params=None):
+            )
+
+        self.parameters = [
+            Column(
+                name=col['name'],
+                alias=col['alias'],
+                type=col['type']
+            )
+            for col in statement_info['parameters']
+        ]
+
+    def execute_query(self):
         if self.fetched_data is not None:
             # no need to execute
             return
 
+        try:
+            steps = list(self.planner.execute_steps())
+        except PlanningException as e:
+            raise LogicError(e)
+
+        if self.planner.plan.is_resumable:
+            # create query
+            if self.query_id is not None:
+                self.run_query = query_context_controller.get_query(self.query_id)
+            else:
+                self.run_query = query_context_controller.create_query(self.context['query_str'])
+            ctx.run_query_id = self.run_query.record.id
+
         step_result = None
         process_mark = None
         try:
-            steps = list(self.planner.execute_steps(params))
             steps_classes = (x.__class__ for x in steps)
             predict_steps = (ApplyPredictorRowStep, ApplyPredictorStep, ApplyTimeseriesPredictorStep)
             if any(s in predict_steps for s in steps_classes):
@@ -255,10 +254,16 @@ class SQLQuery:
                 with profiler.Context(f'step: {step.__class__.__name__}'):
                     step_result = self.execute_step(step)
                 self.steps_data[step.step_num] = step_result
-        except PlanningException as e:
-            raise LogicError(e)
         except Exception as e:
+            if self.run_query is not None:
+                # set error and place where it stopped
+                self.run_query.on_error(e, step.step_num, self.steps_data)
             raise e
+        else:
+            # mark running query as completed
+            if self.run_query is not None:
+                self.run_query.finish()
+                ctx.run_query_id = None
         finally:
             if process_mark is not None:
                 delete_process_mark('predict', process_mark)
@@ -270,27 +275,7 @@ class SQLQuery:
         if len(self.steps_data) == 0:
             return
 
-        try:
-            if self.outer_query is not None:
-                # workaround for subqueries in superset. remove it?
-                # +++
-                # ???
-
-                result = step_result
-                df = result.to_df()
-
-                df2 = query_df(df, self.outer_query)
-
-                result2 = ResultSet().from_df(df2, database='', table_name='')
-
-                self.columns_list = result2.columns
-                self.fetched_data = result2
-
-            else:
-                result = step_result
-                self.fetched_data = result
-        except Exception as e:
-            raise UnknownError("error in preparing result query step") from e
+        self.fetched_data = step_result
 
         try:
             if hasattr(self, 'columns_list') is False:

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

@@ -1,6 +1,7 @@
 from .apply_predictor_step import ApplyPredictorStepCall, ApplyPredictorRowStepCall, ApplyTimeseriesPredictorStepCall
 from .delete_step import DeleteStepCall
 from .fetch_dataframe import FetchDataframeStepCall
+from .fetch_dataframe_partition import FetchDataframePartitionCall
 from .insert_step import InsertToTableCall, SaveToTableCall, CreateTableCall
 from .join_step import JoinStepCall
 from .map_reduce_step import MapReduceStepCall

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

@@ -44,6 +44,5 @@ class DeleteStepCall(BaseStepCall):
 
         query_traversal(query.where, fill_params)
 
-        dn.query(query=query, session=self.session)
-
-        return ResultSet()
+        response = dn.query(query=query, session=self.session)
+        return ResultSet(affected_rows=response.affected_rows)

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

@@ -89,10 +89,11 @@ class FetchDataframeStepCall(BaseStepCall):
             table_alias = (self.context.get('database'), 'result', 'result')
 
             # fetch raw_query
-            df, columns_info = dn.query(
+            response = dn.query(
                 native_query=step.raw_query,
                 session=self.session
             )
+            df = response.data_frame
         else:
             table_alias = get_table_alias(step.query.from_table, self.context.get('database'))
 
@@ -104,13 +105,14 @@ class FetchDataframeStepCall(BaseStepCall):
 
             query, context_callback = query_context_controller.handle_db_context_vars(query, dn, self.session)
 
-            df, columns_info = dn.query(
+            response = dn.query(
                 query=query,
                 session=self.session
             )
+            df = response.data_frame
 
             if context_callback:
-                context_callback(df, columns_info)
+                context_callback(df, response.columns)
 
         result = ResultSet()