Repository: mindsdb/mindsdb
Pull request: Run prediction in partitions
## Description

Depends on: https://github.com/mindsdb/mindsdb_sql/pull/382

Fixes https://linear.app/mindsdb/issue/SQL-20/research-partitioning-approach

Execute prediction in partitions. If partition_size param is used:
```sql
select * from int.tab1 a
    join proj.pred1 p1          
   using partition_size=1000
```
Planner will split fetched data into partitions by size using MapReduce step. Inside this step prediction and join steps will be executed over chunk of data. Execution will be split between sub-threads using [execute_in_threads](https://github.com/mindsdb/mindsdb/pull/9359/files#diff-b5ca614df036c24e57c3442aca7435c9ad1bc3d75d60eebe14274fbc0071af86R19) function. It keeps task queue in optimal size for workers (to have enough tasks available and not generate all tasks in memory at once)
On cloud ApplyPredictorStep will be executed by mindsdb-ml pod using task queue (implemented in [ML task queue PR](https://github.com/mindsdb/mindsdb/pull/7437)).
 By increasing count of these pods we will able to scale prediction horizontally 
Max count of parallel chunk executions is:
- for cloud: defined in MAX_QUERY_PARTITIONS env variable. 10 if not set
- for local: count of CPUs - 2 (keep 2 CPU cores for mindsdb server)


Minor changes:
- [changed approach](https://github.com/mindsdb/mindsdb/pull/9359/commits/5c835cf953690f1a0141085c469f9b8ec59ea077) to fix #8842)
- added 'stop after first failure' flag for some tests in CI

Demo: 
https://www.loom.com/share/c8269745c225492c85a25141e925fd59

## Type of change

- [x] ⚡ New feature (non-breaking change which adds functionality)

## Verification Process

To ensure the changes are working as expected:

 - [ ]   Test Location: Specify the URL or path for testing.
 - [ ]   Verification Steps: Outline the steps or queries needed to validate the change. Include any data, configurations, or actions required to reproduce or see the new functionality.

## Additional Media:

- [x] I have attached a brief loom video or screenshots showcasing the new functionality or change.

## Checklist:

- [x] My code follows the style guidelines(PEP 8) of MindsDB.
- [x] I have appropriately commented on my code, especially in complex areas.
- [ ] Necessary documentation updates are either made or tracked in issues.
- [x] Relevant unit and integration tests are updated or added.





Problem statement:
[Bug]: Error if `JOIN` model with empty data
### Short description of current behavior

If run batch prediction query (`join` model with data), where data query returns empty list of records, then error appear.

### Video or screenshots

_No response_

### Expected behavior

_No response_

### How to reproduce the error

1. Follow `Forecast Quarterly House Sales`
2. instead of query:
```sql
SELECT m.saledate as date, m.ma as forecast
FROM mindsdb.house_sales_model as m
JOIN example_db.house_sales as t
WHERE t.saledate > LATEST
AND t.type = 'house'
AND t.bedrooms=2
LIMIT 4;
```
run
```sql
SELECT m.saledate as date, m.ma as forecast
FROM mindsdb.house_sales_model as m
JOIN example_db.house_sales as t
WHERE t.saledate > LATEST AND t.type = 'unit'
AND t.bedrooms=4
LIMIT 4;
```

### Anything else?

_No response_

Hints from issue discussion:

Commit message: mindsdb sql dep

Patch to review:
diff --git a/.github/workflows/test_on_push.yml b/.github/workflows/test_on_push.yml
index ff060a77e68..45079ecbe74 100644
--- a/.github/workflows/test_on_push.yml
+++ b/.github/workflows/test_on_push.yml
@@ -125,12 +125,12 @@ jobs:
       - name: Run unit tests
         run: |
           if [ "$RUNNER_OS" == "Linux" ]; then
-            env PYTHONPATH=./ pytest tests/unit/test_executor.py
-            env PYTHONPATH=./ pytest tests/unit/test_project_structure.py
-            env PYTHONPATH=./ pytest tests/unit/test_predictor_params.py
-            env PYTHONPATH=./ pytest tests/unit/test_mongodb_handler.py
-            env PYTHONPATH=./ pytest tests/unit/test_mongodb_server.py
-            env PYTHONPATH=./ pytest tests/unit/test_cache.py
+            env PYTHONPATH=./ pytest tests/unit/test_executor.py -x
+            env PYTHONPATH=./ pytest tests/unit/test_project_structure.py -x
+            env PYTHONPATH=./ pytest tests/unit/test_predictor_params.py -x
+            env PYTHONPATH=./ pytest tests/unit/test_mongodb_handler.py -x
+            env PYTHONPATH=./ pytest tests/unit/test_mongodb_server.py -x
+            env PYTHONPATH=./ pytest tests/unit/test_cache.py -x
             env PYTHONPATH=./ pytest tests/unit/test_llm_utils.py
             env PYTHONPATH=./ pytest tests/unit/ml_handlers/test_minds_endpoint.py
             env PYTHONPATH=./ pytest tests/unit/ml_handlers/test_openai.py
diff --git a/mindsdb/api/executor/sql_query/sql_query.py b/mindsdb/api/executor/sql_query/sql_query.py
index 10675ea32b2..769fc9c0943 100644
--- a/mindsdb/api/executor/sql_query/sql_query.py
+++ b/mindsdb/api/executor/sql_query/sql_query.py
@@ -58,7 +58,7 @@ def __init__(self, sql, session, execute=True, database=None):
         }
 
         self.columns_list = None
-        self.steps_data = []
+        self.steps_data = {}
 
         self.planner = None
         self.parameters = []
@@ -214,7 +214,7 @@ def prepare_query(self, prepare=True):
                 for step in self.planner.prepare_steps(self.query):
                     data = self.execute_step(step)
                     step.set_result(data)
-                    self.steps_data.append(data)
+                    self.steps_data[step.step_num] = data
             except PlanningException as e:
                 raise LogicError(e)
 
@@ -247,6 +247,7 @@ def execute_query(self, params=None):
             # no need to execute
             return
 
+        step_result = None
         process_mark = None
         try:
             steps = list(self.planner.execute_steps(params))
@@ -256,9 +257,8 @@ def execute_query(self, params=None):
                 process_mark = create_process_mark('predict')
             for step in steps:
                 with profiler.Context(f'step: {step.__class__.__name__}'):
-                    data = self.execute_step(step)
-                step.set_result(data)
-                self.steps_data.append(data)
+                    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:
@@ -280,7 +280,7 @@ def execute_query(self, params=None):
                 # +++
                 # ???
 
-                result = self.steps_data[-1]
+                result = step_result
                 df = result.to_df()
 
                 df2 = query_df(df, self.outer_query)
@@ -291,7 +291,7 @@ def execute_query(self, params=None):
                 self.fetched_data = result2
 
             else:
-                result = self.steps_data[-1]
+                result = step_result
                 self.fetched_data = result
         except Exception as e:
             raise UnknownError("error in preparing result query step") from e
@@ -310,13 +310,13 @@ def execute_query(self, params=None):
         except Exception as e:
             raise UnknownError("error in column list step") from e
 
-    def execute_step(self, step):
+    def execute_step(self, step, steps_data=None):
         cls_name = step.__class__.__name__
         handler = self.step_handlers.get(cls_name)
         if handler is None:
             raise UnknownError(f"Unknown step: {cls_name}")
 
-        return handler(self).call(step)
+        return handler(self, steps_data=steps_data).call(step)
 
 
 SQLQuery.register_steps()
diff --git a/mindsdb/api/executor/sql_query/steps/apply_predictor_step.py b/mindsdb/api/executor/sql_query/steps/apply_predictor_step.py
index b8b990f83dd..19cecce4913 100644
--- a/mindsdb/api/executor/sql_query/steps/apply_predictor_step.py
+++ b/mindsdb/api/executor/sql_query/steps/apply_predictor_step.py
@@ -116,12 +116,15 @@ def call(self, step):
 
         params = step.params or {}
 
-        for table in data.get_tables()[:1]:  # add  __mindsdb_row_id only for first table
+        # adding __mindsdb_row_id, use first table if exists
+        if len(data.find_columns('__mindsdb_row_id')) == 0:
+            table = data.get_tables()[0] if len(data.get_tables()) > 0 else None
+
             row_id_col = Column(
                 name='__mindsdb_row_id',
-                database=table['database'],
-                table_name=table['table_name'],
-                table_alias=table['table_alias']
+                database=table['database'] if table is not None else None,
+                table_name=table['table_name'] if table is not None else None,
+                table_alias=table['table_alias'] if table is not None else None
             )
 
             row_id = self.context.get('row_id')
diff --git a/mindsdb/api/executor/sql_query/steps/base.py b/mindsdb/api/executor/sql_query/steps/base.py
index 3cbe944ebfe..8ae929bb483 100644
--- a/mindsdb/api/executor/sql_query/steps/base.py
+++ b/mindsdb/api/executor/sql_query/steps/base.py
@@ -2,8 +2,10 @@
 class BaseStepCall:
     bind = None
 
-    def __init__(self, sql_query):
-        self.steps_data = sql_query.steps_data
+    def __init__(self, sql_query, steps_data=None):
+        if steps_data is None:
+            steps_data = sql_query.steps_data
+        self.steps_data = steps_data
 
         self.sql_query = sql_query
         self.context = sql_query.context
diff --git a/mindsdb/api/executor/sql_query/steps/map_reduce_step.py b/mindsdb/api/executor/sql_query/steps/map_reduce_step.py
index e99722900d7..691721b7d15 100644
--- a/mindsdb/api/executor/sql_query/steps/map_reduce_step.py
+++ b/mindsdb/api/executor/sql_query/steps/map_reduce_step.py
@@ -1,3 +1,4 @@
+import os
 import copy
 
 from mindsdb_sql.parser.ast import (
@@ -13,9 +14,10 @@
 
 from mindsdb.api.executor.sql_query.result_set import ResultSet
 from mindsdb.api.executor.exceptions import LogicError
+from mindsdb.utilities.config import Config
+from mindsdb.utilities.context_executor import execute_in_threads
 
 from .base import BaseStepCall
-from .fetch_dataframe import FetchDataframeStepCall
 
 
 def markQueryVar(where):
@@ -30,17 +32,6 @@ def markQueryVar(where):
             where.var_name = where.value
 
 
-def unmarkQueryVar(where):
-    if isinstance(where, BinaryOperation):
-        unmarkQueryVar(where.args[0])
-        unmarkQueryVar(where.args[1])
-    elif isinstance(where, UnaryOperation):
-        unmarkQueryVar(where.args[0])
-    elif isinstance(where, Constant):
-        if hasattr(where, 'is_var') and where.is_var is True:
-            where.value = where.var_name
-
-
 def replaceQueryVar(where, var_value, var_name):
     if isinstance(where, BinaryOperation):
         replaceQueryVar(where.args[0], var_value, var_name)
@@ -64,10 +55,98 @@ class MapReduceStepCall(BaseStepCall):
 
     bind = MapReduceStep
 
-    def call(self, step):
+    def call(self, step: MultipleSteps):
         if step.reduce != 'union':
             raise LogicError(f'Unknown MapReduceStep type: {step.reduce}')
 
+        partition = getattr(step, 'partition', None)
+
+        if partition is not None:
+            data = self._reduce_partition(step, partition)
+
+        else:
+            data = self._reduce_vars(step)
+
+        return data
+
+    def _reduce_partition(self, step, partition):
+        if not isinstance(partition, int):
+            raise NotImplementedError('Only integers are supported in partition definition.')
+
+        input_idx = step.values.step_num
+        input_data = self.steps_data[input_idx]
+        input_columns = list(input_data.columns)
+
+        substeps = step.step
+        if not isinstance(substeps, list):
+            substeps = [substeps]
+
+        data = ResultSet()
+
+        df = input_data.get_raw_df()
+
+        # tasks
+        def split_data_f(df):
+            chunk = 0
+            while chunk * partition < len(df):
+                # create results with partition
+                df1 = df[chunk * partition: (chunk + 1) * partition]
+                chunk += 1
+                yield df1, substeps, input_idx, input_columns
+
+        tasks = split_data_f(df)
+
+        # workers count
+        is_cloud = Config().get('cloud', False)
+        if is_cloud:
+            max_threads = os.getenv('MAX_QUERY_PARTITIONS')
+            if max_threads is not None:
+                max_threads = int(max_threads)
+            else:
+                max_threads = 10
+        else:
+            max_threads = os.cpu_count() - 2
+
+        # don't exceed chunk_count
+        chunk_count = int(len(df) / partition)
+        max_threads = min(max_threads, chunk_count)
+
+        if max_threads < 1:
+            max_threads = 1
+
+        if max_threads == 1:
+            # don't spawn threads
+
+            for task in tasks:
+                sub_data = self._exec_partition(*task)
+                if sub_data:
+                    data = join_query_data(data, sub_data)
+
+        else:
+            for sub_data in execute_in_threads(self._exec_partition, tasks, thread_count=max_threads):
+                if sub_data:
+                    data = join_query_data(data, sub_data)
+
+        return data
+
+    def _exec_partition(self, df, substeps, input_idx, input_columns):
+
+        input_data2 = ResultSet(columns=input_columns.copy())
+        input_data2.add_raw_df(df)
+
+        # execute with modified previous results
+        steps_data2 = self.steps_data.copy()
+        steps_data2[input_idx] = input_data2
+
+        sub_data = None
+        for substep in substeps:
+            sub_data = self.sql_query.execute_step(substep, steps_data=steps_data2)
+            steps_data2[substep.step_num] = sub_data
+
+        return sub_data
+
+    def _reduce_vars(self, step):
+        # extract vars
         step_data = self.steps_data[step.values.step_num]
         vars = []
         for row in step_data.get_records():
@@ -77,70 +156,25 @@ def call(self, step):
                 if name != '__mindsdb_row_id':
                     var_group[name] = value
 
-        data = ResultSet()
-
         substep = step.step
-        if type(substep) is FetchDataframeStep:
-            query = substep.query
-            if len(vars) == 0:
-                substep.query.limit = Constant(0)
-                substep.query.where = None
-
-                sub_data = self._fetch_dataframe_step(substep)
-
-                for column in sub_data.columns:
-                    data.add_column(column)
-
-                data.add_from_result_set(sub_data)
-
-            for var_group in vars:
-                markQueryVar(query.where)
-                for name, value in var_group.items():
-                    replaceQueryVar(query.where, value, name)
-                sub_data = self._fetch_dataframe_step(substep)
-                if len(data.columns) == 0:
-                    data = sub_data
-                else:
-                    data.add_from_result_set(sub_data)
-
-                unmarkQueryVar(query.where)
-        elif type(substep) is MultipleSteps:
-            data = self._multiple_steps_reduce(substep, vars)
-        else:
-            raise LogicError(f'Unknown step type: {step.step}')
-        return data
-
-    def _multiple_steps_reduce(self, step, vars):
-        if step.reduce != 'union':
-            raise LogicError(f'Unknown MultipleSteps type: {step.reduce}')
 
         data = ResultSet()
 
-        # mark vars
-        steps = []
-        for substep in step.steps:
-            if isinstance(substep, FetchDataframeStep) is False:
-                raise LogicError(f'Wrong step type for MultipleSteps: {step}')
-            substep = copy.deepcopy(substep)
-            markQueryVar(substep.query.where)
-            steps.append(substep)
-
         for var_group in vars:
-            steps2 = copy.deepcopy(steps)
-            for name, value in var_group.items():
-                for substep in steps2:
-                    replaceQueryVar(substep.query.where, value, name)
-            sub_data = self._multiple_steps(steps2)
-            data = join_query_data(data, sub_data)
+            steps2 = copy.deepcopy(substep)
 
-        return data
+            self._fill_vars(steps2, var_group)
 
-    def _multiple_steps(self, steps):
-        data = ResultSet()
-        for substep in steps:
-            sub_data = self._fetch_dataframe_step(substep)
+            sub_data = self.sql_query.execute_step(steps2)
             data = join_query_data(data, sub_data)
+
         return data
 
-    def _fetch_dataframe_step(self, step):
-        return FetchDataframeStepCall(self.sql_query).call(step)
+    def _fill_vars(self, step, var_group):
+        if isinstance(step, MultipleSteps):
+            for substep in step.steps:
+                self._fill_vars(substep, var_group)
+        if isinstance(step, FetchDataframeStep):
+            markQueryVar(step.query.where)
+            for name, value in var_group.items():
+                replaceQueryVar(step.query.where, value, name)
diff --git a/mindsdb/integrations/handlers/openai_handler/constants.py b/mindsdb/integrations/handlers/openai_handler/constants.py
index 0d8878eebb5..1c288f1933e 100644
--- a/mindsdb/integrations/handlers/openai_handler/constants.py
+++ b/mindsdb/integrations/handlers/openai_handler/constants.py
@@ -11,7 +11,7 @@
     'gpt-4o'
 )
 COMPLETION_MODELS = ('babbage-002', 'davinci-002')
-FINETUNING_MODELS = ('gpt-3.5-turbo', 'babbage-002', 'davinci-002','gpt-4')
+FINETUNING_MODELS = ('gpt-3.5-turbo', 'babbage-002', 'davinci-002', 'gpt-4')
 COMPLETION_LEGACY_BASE_MODELS = ('davinci', 'curie', 'babbage', 'ada')
 
 FINETUNING_LEGACY_MODELS = FINETUNING_MODELS
diff --git a/mindsdb/integrations/handlers/redshift_handler/__about__.py b/mindsdb/integrations/handlers/redshift_handler/__about__.py
index 16c99739b5c..437cfbd6440 100644
--- a/mindsdb/integrations/handlers/redshift_handler/__about__.py
+++ b/mindsdb/integrations/handlers/redshift_handler/__about__.py
@@ -6,4 +6,4 @@
 __github__ = 'https://github.com/mindsdb/mindsdb'
 __pypi__ = 'https://pypi.org/project/mindsdb/'
 __license__ = 'MIT'
-__copyright__ = 'Copyright 2022- mindsdb'
\ No newline at end of file
+__copyright__ = 'Copyright 2022- mindsdb'
diff --git a/mindsdb/utilities/context_executor.py b/mindsdb/utilities/context_executor.py
index d6f0aed6cc8..7328336fa74 100644
--- a/mindsdb/utilities/context_executor.py
+++ b/mindsdb/utilities/context_executor.py
@@ -1,3 +1,5 @@
+import time
+import types
 from concurrent.futures import ThreadPoolExecutor
 import contextvars
 
@@ -12,3 +14,46 @@ def __init__(self, max_workers=None):
     def _set_child_context(self):
         for var, value in self.context.items():
             var.set(value)
+
+
+def execute_in_threads(func, tasks, thread_count=3, queue_size_k=1.5):
+    """
+    Should be used as generator.
+    Can accept input tasks as generator and keep queue size the same to not overflow the RAM
+
+    :param func: callable, function to execute in threads
+    :param tasks: generator or iterable, list of input for function
+    :param thread_count: number of threads
+    :param queue_size_k: how a queue for workers is bigger than count of threads
+    :return: yield results
+    """
+    executor = ContextThreadPoolExecutor(max_workers=thread_count)
+
+    queue_size = int(thread_count * queue_size_k)
+
+    if not isinstance(tasks, types.GeneratorType):
+        tasks = iter(tasks)
+
+    futures = None
+    while futures is None or len(futures) > 0:
+        if futures is None:
+            futures = []
+
+        # add new portion
+        for i in range(queue_size):
+            try:
+                args = next(tasks)
+                futures.append(executor.submit(func, *args))
+            except StopIteration:
+                break
+
+        # save results
+        for task in futures:
+            if task.done():
+                yield task.result()
+
+        # remove completed tasks
+        futures[:] = [t for t in futures if not t.done()]
+
+        time.sleep(0.1)
+    executor.shutdown(wait=False)
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 4937946f803..76b69c72d68 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -17,7 +17,7 @@ redis >=5.0.0, < 6.0.0
 walrus==0.9.3
 flask-compress >= 1.0.0
 appdirs >= 1.0.0
-mindsdb-sql ~= 0.16.0
+mindsdb-sql ~= 0.17.0
 pydantic >= 2.7.0
 pydantic-settings >= 2.1.0
 mindsdb-evaluator >= 0.0.7, < 0.1.0
diff --git a/tests/unit/test_executor.py b/tests/unit/test_executor.py
index 0dbd630f3fb..75845d9f12e 100644
--- a/tests/unit/test_executor.py
+++ b/tests/unit/test_executor.py
@@ -307,6 +307,17 @@ def test_ts_predictor(self, mock_handler):
         assert ret_df.shape[0] == 1
         assert ret_df.t.min() == dt.datetime(2020, 1, 2)
 
+        # ----- empty data ------
+
+        ret = self.execute('''
+            select p.* from pg.tasks t
+            join mindsdb.task_model p
+            where t.t > LATEST and t.g = 'wrong'
+        ''')
+
+        ret_df = self.ret_to_df(ret)
+        assert ret_df.shape[0] == 0
+
     @patch('mindsdb.integrations.handlers.postgres_handler.Handler')
     def test_ts_predictor_no_group(self, mock_handler):
         # set integration data
diff --git a/tests/unit/test_project_structure.py b/tests/unit/test_project_structure.py
index 3ffaa2f6ce4..b55bf3d4bc9 100644
--- a/tests/unit/test_project_structure.py
+++ b/tests/unit/test_project_structure.py
@@ -955,6 +955,43 @@ def test_delete_from_table(self):
 
         self.run_sql('delete from tbl1 where a=1', database='dummy_data')
 
+    def test_partition(self):
+        df = pd.DataFrame([
+            {'a': 1, 'b': dt.datetime(2020, 1, 1)},
+            {'a': 2, 'b': dt.datetime(2020, 1, 2)},
+            {'a': 3, 'b': dt.datetime(2020, 1, 3)},
+            {'a': 4, 'b': dt.datetime(2020, 1, 5)},
+            {'a': 5, 'b': dt.datetime(2020, 1, 6)},
+            {'a': 6, 'b': dt.datetime(2020, 1, 7)},
+        ])
+        self.save_file('tasks', df)
+
+        # -- create model --
+        self.run_sql(
+            '''
+                CREATE model mindsdb.task_model
+                from files (select * from tasks)
+                PREDICT a
+                using engine='dummy_ml',
+                join_learn_process=true
+            '''
+        )
+
+        # use model
+        ret = self.run_sql('''
+             SELECT *
+               FROM files.tasks as t
+               JOIN mindsdb.task_model as m
+               using partition_size=2
+        ''')
+
+        # the same rows in output
+        assert len(ret) == 6
+        assert set(df['a']) == set(ret['a'])
+
+        # all predicted
+        assert list(ret.predicted.unique()) == [42]
+
 
 class TestJobs(BaseExecutorDummyML):
 


Code review comments:
diff_hunk: @@ -64,10 +55,98 @@ class MapReduceStepCall(BaseStepCall):
 
     bind = MapReduceStep
 
-    def call(self, step):
+    def call(self, step: MultipleSteps):
         if step.reduce != 'union':
             raise LogicError(f'Unknown MapReduceStep type: {step.reduce}')
 
+        partition = getattr(step, 'partition', None)
+
+        if partition is not None:
+            data = self._reduce_partition(step, partition)
+
+        else:
+            data = self._reduce_vars(step)
+
+        return data
+
+    def _reduce_partition(self, step, partition):
+        if not isinstance(partition, int):
+            raise NotImplementedError('Only integers are supported in partition definition.')
+
+        input_idx = step.values.step_num
+        input_data = self.steps_data[input_idx]
+        input_columns = list(input_data.columns)
+
+        substeps = step.step
+        if not isinstance(substeps, list):
+            substeps = [substeps]
+
+        data = ResultSet()
+
+        df = input_data.get_raw_df()
+
+        # tasks
+        def split_data_f(df):
+            chunk = 0
+            while chunk * partition < len(df):
+                # create results with partition
+                df1 = df[chunk * partition: (chunk + 1) * partition]
path: mindsdb/api/executor/sql_query/steps/map_reduce_step.py
text: @user1:
Just in case, we may wanna control for `partition > 0`, too, because if it's negative this won't behave correctly.

@author:
Added checking value of partition