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

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

@@ -0,0 +1,288 @@
+import pandas as pd
+import threading
+import queue
+from typing import List
+
+from mindsdb_sql_parser import ASTNode
+from mindsdb.api.executor.planner.steps import FetchDataframeStepPartition
+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.context import Context, context as ctx
+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 .base import BaseStepCall
+
+
+logger = log.getLogger(__name__)
+
+
+class FetchDataframePartitionCall(BaseStepCall):
+    """
+    Alternative to FetchDataframeCall but fetch data by batches wrapping user's query to:
+
+     select * from ({user query})
+      where {track_column} > {previous value}
+      order by track_column
+      limit size {batch_size} `
+
+    """
+
+    bind = FetchDataframeStepPartition
+
+    def call(self, step: FetchDataframeStepPartition) -> ResultSet:
+        """
+        Parameters:
+        - batch_size - count of rows to fetch from database per iteration, optional default 1000
+        - threads - run partitioning in threads, bool or int, optinal, if set:
+           - int value: use this as count of threads
+           - true: table threads, autodetect count of thread
+           - false: disable threads even if ml task queue is enabled
+        - track_column - column used for creating partitions
+          - query will be sorted by this column and select will be limited by batch_size
+        - error (default raise)
+          - when `error='skip'`, errors in partition will be skipped and execution will be continued
+        """
+
+        self.dn = self.session.datahub.get(step.integration)
+        query = step.query
+
+        # fill params
+        fill_params = get_fill_param_fnc(self.steps_data)
+        query_traversal(query, fill_params)
+
+        # get query record
+        run_query = self.sql_query.run_query
+        if run_query is None:
+            raise RuntimeError('Error with partitioning of the query')
+        run_query.set_params(step.params)
+
+        self.table_alias = get_table_alias(step.query.from_table, self.context.get('database'))
+        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':
+            use_threads = True
+
+        # use threads?
+        if 'threads' in step.params:
+            threads = step.params['threads']
+            if isinstance(threads, int):
+                thread_count = threads
+                use_threads = True
+            if threads is True:
+                use_threads = True
+            if threads is False:
+                # disable even with ml task queue
+                use_threads = False
+
+        on_error = step.params.get('error', 'raise')
+        if use_threads:
+            return self.fetch_threads(run_query, query, thread_count=thread_count, on_error=on_error)
+        else:
+            return self.fetch_iterate(run_query, query, on_error=on_error)
+
+    def fetch_iterate(self, run_query: RunningQuery, query: ASTNode, on_error: str = None) -> ResultSet:
+        """
+         Process batches one by one in circle
+        """
+
+        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)
+            try:
+                sub_data = self.exec_sub_steps(df)
+                results.append(sub_data)
+            except Exception as e:
+                if on_error == 'skip':
+                    logger.error(e)
+                else:
+                    raise e
+
+            run_query.set_progress(df, max_track_value)
+
+        return self.concat_results(results)
+
+    def concat_results(self, results: List[ResultSet]) -> ResultSet:
+        """
+        Concatenate list of result sets to single result set
+        """
+        df_list = []
+        for res in results:
+            df, col_names = res.to_df_cols()
+            if len(df) > 0:
+                df_list.append(df)
+
+        data = ResultSet()
+        if len(df_list) > 0:
+            data.from_df_cols(pd.concat(df_list), col_names)
+
+        return data
+
+    def exec_sub_steps(self, df: pd.DataFrame) -> ResultSet:
+        """
+        FetchDataframeStepPartition has substeps defined
+        Every batch of data have to be used to execute these substeps
+        - batch of data is put as result of FetchDataframeStepPartition
+        - 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(
+            df,
+            table_name=self.table_alias[1],
+            table_alias=self.table_alias[2],
+            database=self.table_alias[0]
+        )
+
+        # execute with modified previous results
+        steps_data2 = self.steps_data.copy()
+        steps_data2[self.current_step_num] = input_data
+
+        sub_data = None
+        for substep in self.substeps:
+            sub_data = self.sql_query.execute_step(substep, steps_data=steps_data2)
+            steps_data2[substep.step_num] = sub_data
+        return sub_data
+
+    def fetch_threads(self, run_query: RunningQuery, query: ASTNode,
+                      thread_count: int = None, on_error: str = None) -> ResultSet:
+        """
+        Process batches in threads
+        - spawn required count of threads
+        - create in/out queue to communicate with threads
+        - send task to threads and receive results
+        """
+
+        # create communication queues
+        queue_in = queue.Queue()
+        queue_out = queue.Queue()
+        self.stop_event = threading.Event()
+
+        if thread_count is None:
+            thread_count = get_max_thread_count()
+
+        # 3 tasks per worker during 1 batch
+        partition_size = int(run_query.batch_size / thread_count / 3)
+        # min partition size
+        if partition_size < 10:
+            partition_size = 10
+
+        # create N workers pool
+        workers = []
+        results = []
+
+        try:
+            for i in range(thread_count):
+                worker = threading.Thread(target=self._worker, daemon=True, args=(ctx.dump(), queue_in,
+                                                                                  queue_out, self.stop_event))
+                worker.start()
+                workers.append(worker)
+
+            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)
+
+                # split into chunks and send to workers
+                sent_chunks = 0
+                for df2 in split_data_frame(df, partition_size):
+                    queue_in.put([sent_chunks, df2])
+                    sent_chunks += 1
+
+                batch_results = []
+                for i in range(sent_chunks):
+                    res = queue_out.get()
+                    if 'error' in res:
+                        if on_error == 'skip':
+                            logger.error(res['error'])
+                        else:
+                            raise RuntimeError(res['error'])
+
+                    if res['data']:
+                        batch_results.append(res)
+
+                # sort results
+                batch_results.sort(key=lambda x: x['num'])
+
+                results.append(self.concat_results(
+                    [item['data'] for item in batch_results]
+                ))
+
+                # 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)
+        finally:
+            self.close_workers(workers)
+
+        return self.concat_results(results)
+
+    def close_workers(self, workers: List[threading.Thread]):
+        """
+        Sent signal to workers to stop
+        """
+
+        self.stop_event.set()
+        for worker in workers:
+            if worker.is_alive():
+                worker.join()
+
+    def _worker(self, context: Context, queue_in: queue.Queue, queue_out: queue.Queue, stop_event: threading.Event):
+        """
+        Worker function. Execute incoming tasks unless stop_event is set
+        """
+        ctx.load(context)
+        while True:
+            if stop_event.is_set():
+                break
+
+            try:
+                chunk_num, df = queue_in.get(timeout=1)
+                if df is None:
+                    continue
+
+                sub_data = self.exec_sub_steps(df)
+
+                queue_out.put({'data': sub_data, 'num': chunk_num})
+            except queue.Empty:
+                continue
+
+            except Exception as e:
+                queue_out.put({'error': str(e)})
+                stop_event.set()

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

@@ -91,13 +91,13 @@ class InsertToTableCall(BaseStepCall):
             else:
                 col_names.add(col.alias)
 
-        dn.create_table(
+        response = dn.create_table(
             table_name=table_name,
             result_set=data,
             is_replace=is_replace,
             is_create=is_create
         )
-        return ResultSet()
+        return ResultSet(affected_rows=response.affected_rows)
 
 
 class SaveToTableCall(InsertToTableCall):

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

@@ -47,10 +47,10 @@ class GetTableColumnsCall(BaseStepCall):
         dn = self.session.datahub.get(step.namespace)
         ds_query = Select(from_table=Identifier(table), targets=[Star()], limit=Constant(0))
 
-        data, columns_info = dn.query(ds_query, session=self.session)
+        response = dn.query(ds_query, session=self.session)
 
         data = ResultSet()
-        for column in columns_info:
+        for column in response.columns:
             data.add_column(Column(
                 name=column['name'],
                 type=column.get('type'),

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

@@ -3,13 +3,7 @@ from collections import defaultdict
 import pandas as pd
 
 from mindsdb_sql_parser.ast import (
-    Identifier,
-    Select,
-    Star,
-    Constant,
-    Parameter,
-    Function,
-    Variable
+    Identifier, Select, Star, Constant, Parameter, Function, Variable, BinaryOperation
 )
 
 from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import SERVER_VARIABLES
@@ -87,7 +81,7 @@ class QueryStepCall(BaseStepCall):
 
     bind = QueryStep
 
-    def call(self, step):
+    def call(self, step: QueryStep):
         query = step.query
 
         if step.from_table is not None:
@@ -190,6 +184,24 @@ class QueryStepCall(BaseStepCall):
         fill_params = get_fill_param_fnc(self.steps_data)
         query_traversal(query, fill_params)
 
+        if not step.strict_where:
+            # remove conditions with not-existed columns.
+            #   these conditions can be already used as input to model or knowledge base
+            #   but can be absent in their output
+
+            def remove_not_used_conditions(node, **kwargs):
+                # find last in where
+                if isinstance(node, BinaryOperation):
+                    for arg in node.args:
+                        if isinstance(arg, Identifier) and len(arg.parts) > 1:
+                            key = tuple(arg.parts[-2:])
+                            if key not in col_idx:
+                                # exclude
+                                node.args = [Constant(0), Constant(0)]
+                                node.op = '='
+
+            query_traversal(query.where, remove_not_used_conditions)
+
         query_traversal(query, check_fields)
         query.where = query_context_controller.remove_lasts(query.where)
 

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

@@ -18,8 +18,6 @@ class UpdateToTableCall(BaseStepCall):
     bind = UpdateToTable
 
     def call(self, step):
-        data = ResultSet()
-
         if len(step.table.parts) > 1:
             integration_name = step.table.parts[0]
             table_name_parts = step.table.parts[1:]
@@ -85,8 +83,8 @@ class UpdateToTableCall(BaseStepCall):
 
             if result_step is None:
                 # run as is
-                dn.query(query=update_query, session=self.session)
-                return data
+                response = dn.query(query=update_query, session=self.session)
+                return ResultSet(affected_rows=response.affected_rows)
             result_data = self.steps_data[result_step.result.step_num]
 
             # link nodes with parameters for fast replacing with values
@@ -125,5 +123,5 @@ class UpdateToTableCall(BaseStepCall):
             for param_name, param in params_map_index:
                 param.value = row[param_name]
 
-            dn.query(query=update_query, session=self.session)
-        return data
+            response = dn.query(query=update_query, session=self.session)
+        return ResultSet(affected_rows=response.affected_rows)

mindsdb/api/http/namespaces/sql.py

@@ -59,7 +59,10 @@ class Query(Resource):
                 result = mysql_proxy.process_query(query)
 
                 if result.type == SQL_RESPONSE_TYPE.OK:
-                    query_response = {"type": SQL_RESPONSE_TYPE.OK}
+                    query_response = {
+                        "type": SQL_RESPONSE_TYPE.OK,
+                        "affected_rows": result.affected_rows
+                    }
                 elif result.type == SQL_RESPONSE_TYPE.TABLE:
                     data = result.data.to_lists(json_types=True)
                     query_response = {

mindsdb/api/mysql/mysql_proxy/data_types/mysql_packets/ok_packet.py

@@ -40,7 +40,7 @@ class OkPacket(Packet):
     def setup(self):
         eof = self._kwargs.get('eof', False)
         self.ok_header = Datum('int<1>', 0xFE if eof is True else 0)
-        self.affected_rows = Datum('int<lenenc>', self._kwargs.get('affected_rows', 0))
+        self.affected_rows = Datum('int<lenenc>', self._kwargs.get('affected_rows') or 0)
         self.last_insert_id = Datum('int<lenenc>', 0)
         status = self._kwargs.get('status', 0x0002)
         self.server_status = Datum('int<2>', status)

mindsdb/api/mysql/mysql_proxy/executor/mysql_executor.py

@@ -4,6 +4,7 @@ from mindsdb.api.executor.planner import utils as planner_utils
 import mindsdb.utilities.profiler as profiler
 from mindsdb.api.executor.sql_query.result_set import Column
 from mindsdb.api.executor.sql_query import SQLQuery
+from mindsdb.api.executor.data_types.answer import ExecuteAnswer
 from mindsdb.api.executor.command_executor import ExecuteCommands
 from mindsdb.api.mysql.mysql_proxy.utilities import ErSqlSyntaxError
 from mindsdb.utilities import log
@@ -12,37 +13,20 @@ logger = log.getLogger(__name__)
 
 
 class Executor:
-    """This class stores initial and intermediate params
-    between different steps of query execution. And it is also
-    creates a separate instance of ExecuteCommands to execute the current
-    query step.
-
-    IMPORTANT: A public API of this class is a contract.
-    And there are at least 2 classes strongly depend on it:
-        ExecuctorClient
-        ExecutorService.
-    These classes do the same work as Executor when
-    MindsDB works in 'modularity' mode.
-    Thus please make sure that IF you change the API,
-    you must update the API of these two classes as well!"""
-
     def __init__(self, session, sqlserver):
         self.session = session
         self.sqlserver = sqlserver
 
         self.query = None
 
-        # returned values
-        # all this attributes needs to be added in
-        # self.json() method
         self.columns = []
         self.params = []
         self.data = None
-        self.state_track = None
         self.server_status = None
         self.is_executed = False
         self.error_message = None
         self.error_code = None
+        self.executor_answer: ExecuteAnswer = None
 
         self.sql = ""
         self.sql_lower = ""
@@ -126,14 +110,7 @@ class Executor:
         if self.is_executed:
             return
 
-        ret = self.command_executor.execute_command(self.query)
-        self.error_code = ret.error_code
-        self.error_message = ret.error_message
+        executor_answer: ExecuteAnswer = self.command_executor.execute_command(self.query)
+        self.executor_answer = executor_answer
 
         self.is_executed = True
-
-        if ret.data is not None:
-            self.data = ret.data
-            self.columns = ret.data.columns
-
-        self.state_track = ret.state_track

mindsdb/api/mysql/mysql_proxy/libs/constants/mysql.py

@@ -94,6 +94,7 @@ class COMMANDS(object):
     COM_STMT_PREPARE = int('0x16', 0)
     COM_STMT_EXECUTE = int('0x17', 0)
     COM_STMT_FETCH = int('0x1c', 0)
+    COM_STMT_RESET = int('0x1a', 0)
     COM_STMT_CLOSE = int('0x19', 0)
     COM_FIELD_LIST = int('0x04', 0)  # deprecated
 

mindsdb/api/mysql/mysql_proxy/mysql_proxy.py

@@ -21,7 +21,8 @@ import sys
 import tempfile
 import traceback
 from functools import partial
-from typing import Dict, List
+from typing import Dict, List, Optional
+from dataclasses import dataclass
 
 from numpy import dtype as np_dtype
 from pandas.api import types as pd_types
@@ -71,6 +72,7 @@ from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import (
     TYPES,
     getConstName,
 )
+from mindsdb.api.executor.data_types.answer import ExecuteAnswer
 from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
 from mindsdb.api.mysql.mysql_proxy.utilities import (
     ErWrongCharset,
@@ -93,24 +95,16 @@ def empty_fn():
     pass
 
 
+@dataclass
 class SQLAnswer:
-    def __init__(
-        self,
-        resp_type: RESPONSE_TYPE,
-        columns: List[Dict] = None,
-        data: List[Dict] = None,
-        status: int = None,
-        state_track: List[List] = None,
-        error_code: int = None,
-        error_message: str = None,
-    ):
-        self.resp_type = resp_type
-        self.columns = columns
-        self.data = data
-        self.status = status
-        self.state_track = state_track
-        self.error_code = error_code
-        self.error_message = error_message
+    resp_type: RESPONSE_TYPE = RESPONSE_TYPE.OK
+    columns: Optional[List[Dict]] = None
+    data: Optional[List[Dict]] = None   # resultSet ?
+    status: Optional[int] = None
+    state_track: Optional[List[List]] = None
+    error_code: Optional[int] = None
+    error_message: Optional[str] = None
+    affected_rows: Optional[int] = None
 
     @property
     def type(self):
@@ -333,7 +327,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 packages.append(self.last_packet())
             self.send_package_group(packages)
         elif answer.type == RESPONSE_TYPE.OK:
-            self.packet(OkPacket, state_track=answer.state_track).send()
+            self.packet(OkPacket, state_track=answer.state_track, affected_rows=answer.affected_rows).send()
         elif answer.type == RESPONSE_TYPE.ERROR:
             self.packet(
                 ErrPacket, err_code=answer.error_code, msg=answer.error_message
@@ -546,21 +540,23 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
     @profiler.profile()
     def process_query(self, sql):
         executor = Executor(session=self.session, sqlserver=self)
-
         executor.query_execute(sql)
+        executor_answer = executor.executor_answer
 
-        if executor.data is None:
+        if executor_answer.data is None:
             resp = SQLAnswer(
                 resp_type=RESPONSE_TYPE.OK,
-                state_track=executor.state_track,
+                state_track=executor_answer.state_track,
+                affected_rows=executor_answer.affected_rows
             )
         else:
             resp = SQLAnswer(
                 resp_type=RESPONSE_TYPE.TABLE,
-                state_track=executor.state_track,
-                columns=self.to_mysql_columns(executor.columns),
-                data=executor.data,
+                state_track=executor_answer.state_track,
+                columns=self.to_mysql_columns(executor_answer.data.columns),
+                data=executor_answer.data,
                 status=executor.server_status,
+                affected_rows=executor_answer.affected_rows
             )
 
         # Increment the counter and include metadata in attributes
@@ -604,18 +600,20 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
 
     def answer_stmt_execute(self, stmt_id, parameters):
         prepared_stmt = self.session.prepared_stmts[stmt_id]
-        executor = prepared_stmt["statement"]
+        executor: Executor = prepared_stmt["statement"]
 
         executor.stmt_execute(parameters)
 
-        if executor.data is None:
+        executor_answer: ExecuteAnswer = executor.executor_answer
+
+        if executor_answer.data is None:
             resp = SQLAnswer(
-                resp_type=RESPONSE_TYPE.OK, state_track=executor.state_track
+                resp_type=RESPONSE_TYPE.OK, state_track=executor_answer.state_track
             )
             return self.send_query_answer(resp)
 
         # TODO prepared_stmt['type'] == 'lock' is not used but it works
-        columns_def = self.to_mysql_columns(executor.columns)
+        columns_def = self.to_mysql_columns(executor_answer.data.columns)
         packages = [self.packet(ColumnCountPacket, count=len(columns_def))]
 
         packages.extend(self._get_column_defenition_packets(columns_def))
@@ -624,14 +622,14 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             packages.append(self.packet(EofPacket, status=0x0062))
 
         # send all
-        for row in executor.data.to_lists():
+        for row in executor_answer.data.to_lists():
             packages.append(
                 self.packet(BinaryResultsetRowPacket, data=row, columns=columns_def)
             )
 
         server_status = executor.server_status or 0x0002
         packages.append(self.last_packet(status=server_status))
-        prepared_stmt["fetched"] += len(executor.data)
+        prepared_stmt["fetched"] += len(executor_answer.data)
 
         return self.send_package_group(packages)
 
@@ -639,23 +637,24 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         prepared_stmt = self.session.prepared_stmts[stmt_id]
         executor = prepared_stmt["statement"]
         fetched = prepared_stmt["fetched"]
+        executor_answer: ExecuteAnswer = executor.executor_answer
 
-        if executor.data is None:
+        if executor_answer.data is None:
             resp = SQLAnswer(
-                resp_type=RESPONSE_TYPE.OK, state_track=executor.state_track
+                resp_type=RESPONSE_TYPE.OK, state_track=executor_answer.state_track
             )
             return self.send_query_answer(resp)
 
         packages = []
-        columns = self.to_mysql_columns(executor.columns)
-        for row in executor.data[fetched:limit].to_lists():
+        columns = self.to_mysql_columns(executor_answer.data.columns)
+        for row in executor_answer.data[fetched:limit].to_lists():
             packages.append(
                 self.packet(BinaryResultsetRowPacket, data=row, columns=columns)
             )
 
-        prepared_stmt["fetched"] += len(executor.data[fetched:limit])
+        prepared_stmt["fetched"] += len(executor_answer.data[fetched:limit])
 
-        if len(executor.data) <= limit + fetched:
+        if len(executor_answer.data) <= limit + fetched:
             status = sum(
                 [
                     SERVER_STATUS.SERVER_STATUS_AUTOCOMMIT,
@@ -772,6 +771,8 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 elif p.type.value == COMMANDS.COM_FIELD_LIST:
                     # this command is deprecated, but console client still use it.
                     response = SQLAnswer(RESPONSE_TYPE.OK)
+                elif p.type.value == COMMANDS.COM_STMT_RESET:
+                    response = SQLAnswer(RESPONSE_TYPE.OK)
                 else:
                     logger.warning("Command has no specific handler, return OK msg")
                     logger.debug(str(p))