/*
   Copyright (c) 2026, MariaDB

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; version 2 of the License.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1335
   USA */

/**
  @file

  Result-row transport for parallel query: the row shape the two ends agree on,
  and the batch implementation of the two ends themselves.

  sql_parallel_transport.h says what a transport is and what the interface
  promises; this is the one that exists today. What a worker computes and what
  the manager does with a row once it has arrived are not here -- they are in
  sql_parallel_execution.cc, which reaches the transport only through
  pwt_row_layout, pwt_row_sink and pwt_row_source.
*/

#include "mariadb.h"
#include "mysqld_error.h"
#include "sql_priv.h"
#include "unireg.h"
#include "sql_class.h"
#include "sql_select.h"
#include "sql_parallel_workers.h"

#ifdef HAVE_PSI_INTERFACE
static PSI_cond_key key_COND_pwt_data_space;
static PSI_cond_info all_pwt_transport_conds[]=
{
  { &key_COND_pwt_data_space, "pwt_batch_source::COND_data_space",  0},
};

void pwt_transport_init_psi_keys(void)
{
  mysql_cond_register("sql", all_pwt_transport_conds,
                      array_elements(all_pwt_transport_conds));
}
#endif /* HAVE_PSI_INTERFACE */


/*****************************************************************************
  pwt_row_layout -- the shape both ends agree on
*****************************************************************************/


/**
  @brief
    result_defn := one clone per shipped column.

  @description
    Clones, so that the query's own items are never bound to a container field,
    and so that the definition can be rebuilt from scratch -- which is what the
    fall-back out of pre-aggregation does, having appended partial columns to a
    definition it then has to take them back off.

  @return true on error (my_error() called).
*/

bool pwt_row_layout::clone_base_defn(THD *thd)
{
  result_defn.empty();
  List_iterator_fast<Item> li(ship_list);
  Item *sel_item;
  while ((sel_item= li++))
  {
    Item *c= sel_item->deep_copy_with_checks(thd);
    if (!c || result_defn.push_back(c, thd->mem_root))
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item));
      return true;
    }
  }
  n_ship_base= result_defn.elements;
  return false;
}


/**
  @description
  Forget that the workers were going to pre-aggregate. Everything downstream
  reads 'grouped' -- the worker's terminal, its setup, and the manager's drain
  -- so clearing it here is the whole of the decision. The allocations stay on
  the statement's mem_root, which outlives all of this.
*/

void pwt_row_layout::forget_aggregates()
{
  /* plan_aggregates stays: the query still groups, the manager just does it. */
  grouped=          false;
  whole_table=      false;
  n_sums=           0;
  mgr_sums=         nullptr;
  partial_items=    nullptr;
  group_defn=       nullptr;
  group_pos=        nullptr;
  n_group=          0;
  group_parts=      0;
  group_length=     0;
  group_null_parts= 0;
}


/**
  @brief
    Which shipped column holds the field 'want'.

  @return  its position in ship_list, or -1 if it is not shipped.
*/

int pwt_row_layout::ship_pos_of(Field *want) const
{
  uint pos= 0;
  Item *it;
  List_iterator_fast<Item> si(const_cast<List<Item>&>(ship_list));
  while ((it= si++))
  {
    if (it->type() == Item::FIELD_ITEM && ((Item_field *) it)->field == want)
      return (int) pos;
    pos++;
  }
  return -1;
}


/**
  @brief
    Remember which shipped column each ORDER element sorts on.

  @description
    Every element names a base-table field -- the gate has already refused an
    order it could not say that of -- and every field the query reads is
    shipped, so each one is found. The positions are what build_sort_order()
    turns into an order over a container.

  @return  true on error (my_error() called).
*/

bool pwt_row_layout::build_sort_positions(THD *thd, ORDER *plan_sort)
{
  n_sort= 0;
  for (ORDER *o= plan_sort; o; o= o->next)
    n_sort++;
  DBUG_ASSERT(n_sort);

  if (!(sort_pos= thd->alloc<uint>(n_sort)))
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) (n_sort * sizeof(uint)));
    return true;
  }

  uint k= 0;
  for (ORDER *o= plan_sort; o; o= o->next, k++)
  {
    Item *real= (*o->item)->real_item();
    DBUG_ASSERT(real->type() == Item::FIELD_ITEM);
    int pos= ship_pos_of(((Item_field *) real)->field);
    if (pos < 0)
    {
      DBUG_ASSERT(0);              // every column the query reads is shipped
      my_error(ER_INTERNAL_ERROR, MYF(0),
               "parallel query: ORDER BY column is not shipped");
      return true;
    }
    sort_pos[k]= (uint) pos;
  }
  plan_order= plan_sort;
  return false;
}


ORDER *pwt_row_layout::build_sort_order(THD *thd, TABLE *container)
{
  DBUG_ASSERT(plan_sorts && plan_order && n_sort);
  ORDER *so= thd->alloc<ORDER>(n_sort);
  if (!so)
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) (n_sort * sizeof(ORDER)));
    return nullptr;
  }
  uint k= 0;
  for (ORDER *o= plan_order; o; o= o->next, k++)
  {
    /*
      A field of the container, not of the base table the plan named: this
      order sorts the rows the workers produced, which is where they are.
      Same column, because sort_pos[k] is the position the plan's field was
      shipped in and the container was built from the same definition.
    */
    Item *f= new (thd->mem_root) Item_field(thd,
                                            container->field[sort_pos[k]]);
    if (!f)
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item_field));
      return nullptr;
    }
    bzero((char *) &so[k], sizeof(ORDER));
    so[k].item_ptr=  f;
    so[k].item=      &so[k].item_ptr;
    so[k].direction= o->direction;
    so[k].next=      k + 1 < n_sort ? &so[k + 1] : nullptr;
  }
  return so;
}


/**
  @brief
    Work out what travels, define the columns, and build the manager's
    receiving container.

  @description
    What travels is every base-table column the query reads, in table order.
    See the class comment for why that rather than the projected select list.

    The definition is a list of clones, so the query's own items are never
    bound to a container field, and the manager and every worker build the
    identical layout from it.

  @return  true on error (my_error() has been called).
*/

bool pwt_row_layout::build(THD *thd, JOIN *join_arg, TABLE **tables,
                           uint n_tables, pwt_preagg_kind preagg,
                           ORDER *plan_group, ORDER *plan_sort)
{
  DBUG_ENTER("pwt_row_layout::build");
  join= join_arg;
  plan_aggregates= grouped= (preagg != PWT_PREAGG_NONE);
  whole_table= (preagg == PWT_PREAGG_WHOLE);
  DBUG_ASSERT((preagg == PWT_PREAGG_GROUPED) == (plan_group != nullptr));
  plan_sorts= (plan_sort != nullptr);

  for (uint t= 0; t < n_tables; t++)
  {
    TABLE *tbl= tables[t];
    for (Field **f= tbl->field; *f; f++)
    {
      if (!bitmap_is_set(tbl->read_set, (*f)->field_index))
        continue;
      Item *itf= new (thd->mem_root) Item_field(thd, *f);
      if (!itf || ship_list.push_back(itf, thd->mem_root))
      {
        my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item_field));
        DBUG_RETURN(true);
      }
    }
  }

  /*
    A query reading no column of any table still needs a row shape, because the
    transport measures its rows in record images. "SELECT COUNT(*)" is that
    query: it reads no column, and the count is the number of images that
    arrive.
  */
  if (ship_list.is_empty())
  {
    Item *one= new (thd->mem_root) Item_int(thd, (longlong) 1, 1);
    if (!one || ship_list.push_back(one, thd->mem_root))
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item_int));
      DBUG_RETURN(true);
    }
  }

  if (!(copy_back= new (thd->mem_root) Copy_field[ship_list.elements]))
  {
    my_error(ER_OUTOFMEMORY, MYF(0),
             (int) (ship_list.elements * sizeof(Copy_field)));
    DBUG_RETURN(true);
  }

  if (clone_base_defn(thd))
    DBUG_RETURN(true);

  if (plan_sorts && build_sort_positions(thd, plan_sort))
    DBUG_RETURN(true);

  /*
    Define grouping information in our per worker aggregate container.
    Create and instantiate our temporary tables
  */

  if (grouped && build_aggregates(thd, plan_group))
    DBUG_RETURN(true);

  /*
    Twice at most: once wanting to pre-aggregate, and once more having found
    that this plan cannot. See the foot of the loop.
  */
  for (;;)
  {
    ORDER *group= nullptr;
    if (grouped && !(group= clone_group_defn(thd)))
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) (n_group * sizeof(ORDER)));
      DBUG_RETURN(true);
    }
    if (make_container(thd, &recv, group))
    {
      my_error(ER_INTERNAL_ERROR, MYF(0),
               "parallel query: failed to build the result row container");
      DBUG_RETURN(true);
    }

    /*
      A group key create_tmp_table() could not build as an index is not one a
      worker can look a group up in: it falls back to a unique constraint, which
      is a hash of the row in a hidden column, enforced on write, with nothing
      to read a group back out of. accumulate_group() needs the index.

      Whether it managed one can only be asked here, of the finished table.
      can_run_query_in_workers() asks what it can of the SQL -- the key parts
      are plain columns, none wider than a varchar, the whole no wider than a
      blob -- but whether that becomes an index depends on the engine
      create_tmp_table() picked and that engine's own key limits, and it picks
      the engine from the row it is building.

      So give up pre-aggregating and build the layout the query would have had
      without it. The query still runs in the workers; they ship their rows and
      the manager groups them, which is what every query did before there were
      partials to ship.
    */
    bool no_group_index= !(recv.table->group && recv.table->s->keys == 1 &&
                           !recv.table->s->have_unique_constraint());
    /*
      Nothing this tree can express reaches the real thing: the gate already
      refuses the two shapes that make create_tmp_table() fall back for a heap
      container, and its engine-limit test is only in the on-disk branch. So
      the fall-back is reached on demand instead, or it would never be run.
    */
    DBUG_EXECUTE_IF("pwt_no_group_index", no_group_index= true;);

    if (!grouped || !no_group_index)
      break;

    DBUG_PRINT("info", ("no group index on the result container, "
                        "shipping rows instead"));
    free_container(thd, &recv);
    forget_aggregates();
    if (clone_base_defn(thd))
      DBUG_RETURN(true);
  }
  reclength= recv.table->s->reclength;

  /*
    Pair each container column with the base-table field it was projected from,
    so that receiving a row is a copy per column back into the manager's own
    records. Position i of ship_list is column i of the container, which is how
    make_container() built it. The one item that is not an Item_field is the
    filler shipped for a query that reads no column, and it has nowhere to go
    back to.
  */
  {
    List_iterator_fast<Item> si(ship_list);
    Item *it;
    n_copy_back= 0;
    for (uint i= 0; (it= si++); i++)
      if (it->type() == Item::FIELD_ITEM)
        copy_back[n_copy_back++].set(((Item_field*) it)->field,
                                     recv.table->field[i], false);
  }

  /* add our Item_field partial_items, for MIN/MAX's direct_add(). */
  for (uint i= 0; i < n_sums; i++)
    if (!(partial_items[i]= new (thd->mem_root)
                          Item_field(thd, recv.table->field[n_ship_base + i])))
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item_field));
      DBUG_RETURN(true);
    }
  DBUG_RETURN(false);
}


/**
  @brief
    Define the partial columns and the group accumulation key

  @return true on error (my_error() called).
*/

bool pwt_row_layout::build_aggregates(THD *thd, ORDER *plan_group)
{
  Item *one= nullptr;
  n_sums= 0;
  for (Item_sum **s= join->sum_funcs; *s; s++)
    n_sums++;
  if (!(mgr_sums= thd->alloc<Item_sum*>(n_sums)) ||
      !(partial_items= thd->alloc<Item*>(n_sums)))
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) (n_sums * sizeof(void*)));
    return true;
  }

  uint i= 0;
  for (Item_sum **s= join->sum_funcs; *s; s++, i++)
  {
    mgr_sums[i]= *s;
    Item *c= (*s)->deep_copy_with_checks(thd);
    if (!c || result_defn.push_back(c, thd->mem_root))
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item));
      return true;
    }
  }

  /*
    The key. Each entry has to name a result_defn item while create_tmp_table()
    runs, because that is how it finds which column of the table it is building
    the key part covers; group_pos remembers which column that was, so a copy
    of the key can later be re-pointed at a container's own fields.
  */
  n_group= 0;
  for (ORDER *g= plan_group; g; g= g->next)
    n_group++;
  /*
    Aggregating over the whole table is one group, and the column it is keyed
    on is a constant this layout adds: nothing in the query distinguishes one
    row from another, so anything constant does. The definition carries it like
    any other column, after the partials, and from here on this shape is the
    grouped one with a single group in it -- which is what lets the worker
    accumulate through end_update() into a keyed container, and keeps the
    container keyed, which is the form Item_sum_avg::create_tmp_field() packs
    its count into.
  */
  if (whole_table)
  {
    DBUG_ASSERT(!n_group);
    one= new (thd->mem_root) Item_int(thd, (longlong) 1, 1);
    if (!one || result_defn.push_back(one, thd->mem_root))
    {
      my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(Item_int));
      return true;
    }
    n_group= 1;
  }
  if (!(group_defn= thd->alloc<ORDER>(n_group)) ||
      !(group_pos= thd->alloc<uint>(n_group)))
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) (n_group * sizeof(ORDER)));
    return true;
  }
  if (whole_table)
  {
    bzero((char *) &group_defn[0], sizeof(ORDER));
    group_defn[0].item_ptr=  one;
    group_defn[0].item=      &group_defn[0].item_ptr;
    group_defn[0].direction= ORDER::ORDER_ASC;
    group_defn[0].next=      nullptr;
    group_pos[0]=            result_defn.elements - 1;
  }

  uint k= 0;
  for (ORDER *g= plan_group; g; g= g->next, k++)
  {
    Field *want= ((Item_field *) (*g->item)->real_item())->field;

    uint pos= 0;
    Item *it;
    List_iterator_fast<Item> si(ship_list);
    while ((it= si++))
    {
      if (it->type() == Item::FIELD_ITEM && ((Item_field *) it)->field == want)
        break;
      pos++;
    }
    if (!it)
    {
      DBUG_ASSERT(0);            // every column the query reads is shipped
      return true;
    }

    List_iterator_fast<Item> di(result_defn);
    Item *defn_item= nullptr;
    for (uint j= 0; j <= pos; j++)
      defn_item= di++;

    bzero((char *) &group_defn[k], sizeof(ORDER));
    group_defn[k].item_ptr=  defn_item;
    group_defn[k].item=      &group_defn[k].item_ptr;
    group_defn[k].direction= ORDER::ORDER_ASC;
    group_defn[k].next=      k + 1 < n_group ? &group_defn[k + 1] : nullptr;
    group_pos[k]=            pos;
  }

  /*
    Measure the key once, here, and hand the sizes to every container.
    create_tmp_table() would otherwise re-derive them per container from items
    that have since been re-pointed at that container's own fields.
  */
  TMP_TABLE_PARAM measure;
  measure.init();
  calc_group_buffer(&measure, group_defn);
  group_parts=      measure.group_parts;
  group_length=     measure.group_length;
  group_null_parts= measure.group_null_parts;
  return false;
}


ORDER *pwt_row_layout::clone_group_defn(THD *thd)
{
  DBUG_ASSERT(group_defn && n_group);
  ORDER *copy= thd->alloc<ORDER>(n_group);
  if (!copy)
    return nullptr;
  for (uint k= 0; k < n_group; k++)
  {
    copy[k]= group_defn[k];
    copy[k].item= &copy[k].item_ptr;             // at this copy's own storage
    copy[k].next= k + 1 < n_group ? &copy[k + 1] : nullptr;
  }
  return copy;
}


/*
  @brief
    Build one container of this layout: the record image both ends copy.

  @description
    A tmp table used for its record buffer and its fields. Whether rows are
    ever written into it through the storage engine is the transport's
    business: the batch transport writes none, and ships images of record[0].

  @return  true on error, false on success (*out set).
*/

bool pwt_row_layout::make_container(THD *thd, pwt_row_container *out,
                                    ORDER *group)
{
  return make_container_from(thd, result_defn, out, group);
}


bool pwt_row_layout::make_container_from(THD *thd, List<Item> &defn,
                                         pwt_row_container *out,
                                         ORDER *group)
{
  /*
    out needs its own TMP_TABLE_PARAM because
    1) create_tmp_table() overwrites param->func_count with the number
       of items it actually has to copy, so a second table built from the same
       param allocates fewer fields than the layout needs.
    2) create_tmp_table() allocates param->start_recinfo out of the table's
       own mem_root, so the column descriptions belong to one container and
       die with it; sharing a param would leave every container but the last
       describing a table that has been freed.
  */
  TMP_TABLE_PARAM *param= new (thd->mem_root) TMP_TABLE_PARAM;
  if (!param)
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(TMP_TABLE_PARAM));
    return true;
  }
  param->init();
  count_field_types(join->select_lex, param, defn, false);
  param->skip_create_table= true;
  /*
    This container is written by a worker and read by the manager, so it is
    instantiated cross-thread below. A rebuild on disk has to be opened the
    same way, and for a pre-aggregating worker that rebuild happens inside
    end_update(), which reads the decision from here.
  */
  param->cross_thread= true;
  /* The key sizes measured once in build_aggregates(), not re-derived here. */
  param->group_parts=      group_parts;
  param->group_length=     group_length;
  param->group_null_parts= group_null_parts;

  /*
    TMP_TABLE_ALL_COLUMNS makes create_tmp_table() give every item of the list a
    field, including the constant ones. By default it skips constants, which is
    right for a query that materialises its result and can evaluate them once
    outside the table, but wrong here: the layout has to mirror the definition
    one for one, because the worker projects item i into field i and ships the
    record image, and the manager copies field i back to the column it came
    from. Without this a list holding a constant built a table with fewer fields
    than the projection walks over.
  */
  const ulonglong opts= join->select_options | TMP_TABLE_ALL_COLUMNS;
  TABLE *t= create_tmp_table(thd, param, defn,
                             group, false, grouped,
                             opts, HA_POS_ERROR,
                             &empty_clex_str, true, false);
  if (!t)
    return true;
  if (instantiate_tmp_table(t, param->keyinfo, param->start_recinfo,
                            &param->recinfo, opts, true /*cross_thread*/))
  {
    free_tmp_table(thd, t);
    return true;
  }
  /*
    The whole transport is positional, so a layout that does not match the
    definition item for item would have the two ends reading different columns,
    or walking past the end of the field array. Refuse instead, whatever the
    reason turns out to be.
  */
  DBUG_ASSERT(t->s->fields == defn.elements);
  if (t->s->fields != defn.elements)
  {
    free_tmp_table(thd, t);
    return true;
  }
  out->table= t;
  out->param= param;
  return false;
}


void pwt_row_layout::free_container(THD *thd, pwt_row_container *c)
{
  if (c->table)
  {
    free_tmp_table(thd, c->table);
    c->table= nullptr;
  }
  c->param= nullptr;    // points into c->table
}


void pwt_row_layout::cleanup(THD *thd)
{
  free_container(thd, &recv);
}


/**
  @brief
    Hand each of the query's aggregates the partial now in recv.

  @description
    One direct_add() per aggregate, choosing the overload by what the aggregate
    is. A NULL partial comes from a worker that saw no non-NULL value in this
    group, and has to be skipped rather than read as a zero or taken as the
    extreme -- which is what each overload does with it, given that it is told.
    The two SUM overloads are told differently: the real one by a flag, the
    decimal one by a null pointer, where a pointer to decimal zero would mean
    the value 0.

    The value stays pending on the aggregate until something consumes it, which
    here is the plan's own terminal: the update_field() or reset_field() it
    calls once it knows which group the row belongs to takes the direct value
    in place of the row.
*/

void pwt_row_layout::direct_add_partials()
{
  for (uint i= 0; i < n_sums; i++)
  {
    Item_sum *a= mgr_sums[i];
    Field *f= recv.table->field[n_ship_base + i];

    switch (a->sum_func()) {
    case Item_sum::COUNT_FUNC:
      ((Item_sum_count *) a)->direct_add(f->is_null() ? 0 : f->val_int());
      break;
    case Item_sum::SUM_FUNC:
      if (a->result_type() == DECIMAL_RESULT)
      {
        my_decimal buf;
        ((Item_sum_sum *) a)->direct_add(f->is_null() ? (my_decimal *) NULL
                                                      : f->val_decimal(&buf));
      }
      else
        ((Item_sum_sum *) a)->direct_add(f->is_null() ? 0.0 : f->val_real(),
                                         f->is_null());
      break;
    case Item_sum::MIN_FUNC:
    case Item_sum::MAX_FUNC:
      ((Item_sum_min_max *) a)->direct_add(partial_items[i]);
      break;
    case Item_sum::AVG_FUNC:
    {
      /*
        An average does not merge, so what a worker computed for a group is
        both halves of one: the sum of its rows and how many there were.
        Item_sum_avg::create_tmp_field() packs exactly that into one field --
        the value followed by an 8-byte count -- and update_field() has been
        maintaining it there all along, so the partial column already holds
        both and this only has to take them apart again the same way.
      */
      Item_sum_avg *avg= (Item_sum_avg *) a;
      const uchar *res= f->ptr;
      if (a->result_type() == DECIMAL_RESULT)
      {
        my_decimal partial;
        binary2my_decimal(E_DEC_FATAL_ERROR, res, &partial,
                          avg->f_precision, avg->f_scale);
        const longlong cnt= sint8korr(res + avg->dec_bin_size);
        /* No rows behind it: a null sum, so that it adds nothing. */
        avg->direct_add(cnt ? &partial : (my_decimal *) NULL,
                        (ulonglong) cnt);
      }
      else
      {
        double val;
        float8get(val, res);
        const longlong cnt= sint8korr(res + sizeof(double));
        avg->direct_add(val, cnt == 0, (ulonglong) cnt);
      }
      break;
    }
    default:
      DBUG_ASSERT(0);                     // the gate accepts no others
      break;
    }
  }
}


/**
  @brief
    Put the row now in recv_record() back where the query expects to read it.

    One copy per shipped column, container field to the base-table field it was
    projected from. After this the manager's records hold what a serial scan
    would have left there.
*/

void pwt_row_layout::copy_back_row()
{
  for (uint i= 0; i < n_copy_back; i++)
    (*copy_back[i].do_copy)(&copy_back[i]);
}


/**
  @brief
    Make the manager's own records fit to receive rows.

  @description
    The manager's tables were opened but never read, so clear the flags a
    reader would have left. Copy_field captures &table->null_row, and a stale
    null_row would make every copied field read as NULL.

    Writing into these records is not something a SELECT's write_set allows,
    and Field::store() asserts on that, so mark the fields writable for the
    duration. A record is being filled here in place of the reader that would
    normally have filled it, which is what the helper is for.

  @return  true on error.
*/

bool pwt_row_layout::begin_receive(THD *thd, TABLE **tables, uint n_tables)
{
  if (!(saved_write_set= (MY_BITMAP**)
                         thd->alloc(n_tables * sizeof(MY_BITMAP*))))
    return true;
  for (uint t= 0; t < n_tables; t++)
  {
    tables[t]->status= 0;
    tables[t]->null_row= false;
    saved_write_set[t]= dbug_tmp_use_all_columns(tables[t],
                                                 &tables[t]->write_set);
  }
  return false;
}


void pwt_row_layout::end_receive(TABLE **tables, uint n_tables)
{
  if (!saved_write_set)
    return;
  for (uint t= 0; t < n_tables; t++)
    dbug_tmp_restore_column_map(&tables[t]->write_set, saved_write_set[t]);
  saved_write_set= nullptr;
}


/*
  @description
  Runs in the worker's thread, before its first row. Take the container: the
  engine accounts a write to TABLE::in_use's status counters, and a projection
  into a field asks in_use for the session's time zone and sql_mode. 
*/

bool pwt_tmp_table_sink::begin()
{
  container->table->in_use= current_thd;
  container->table->file->rebind_to_thread();
  return false;
}


/*
  @brief  Has the consumer asked us to stop?

  @description
    Read under the lock rather than peeked at, and only once every so many rows
    -- the answer only ever changes once, and reading it costs the same mutex
    the batch transport takes to hand a buffer over. Stopping is not required
    for correctness here: the manager cannot ask before it has drained
    everything it wants, so this only saves a worker from finishing work
    nobody will read.
*/

bool pwt_tmp_table_sink::stop_requested()
{
  mysql_mutex_lock(&manager->LOCK_data);
  bool stop= manager->workers_must_stop;
  mysql_mutex_unlock(&manager->LOCK_data);
  return stop;
}


/*
  @brief  Keep one finished row.

  @description
    The row is already in the container's record buffer -- the projection put
    it there -- so storing it is one engine write and no copy.
*/

int pwt_tmp_table_sink::emit_row(const uchar *rec)
{
  TABLE *table= container->table;
  DBUG_ASSERT(rec == table->record[0]);
  (void) rec;

  int err= table->file->ha_write_tmp_row(table->record[0]);
  if (unlikely(err))
  {
    if (err == HA_ERR_RECORD_FILE_FULL)
    {
      /*
        The heap container is full: rebuild it on disk and carry on there. The
        column descriptions it is driven from are this container's own, which
        is what pairing each container with its own TMP_TABLE_PARAM was for --
        a shared one would describe some other container, in a mem_root that
        may already be gone.

        The row that did not fit is written by the conversion itself, from
        record[0], so there is nothing to re-emit here.

        Run on this worker's thread and with this worker's THD, which is what
        makes the thread questions live: the conversion re-opens the table, and
        Aria binds an open handle to the opening thread's my_thread_var, freed
        when this worker's THD is destroyed. The manager reads this container
        afterwards. Being worked through.
      */
      THD *worker_thd= current_thd;
      const int64 before= worker_thd->status_var.local_memory_used;

      if (create_internal_tmp_table_from_heap(worker_thd, table,
                                              container->param->start_recinfo,
                                              &container->param->recinfo,
                                              err, 0, NULL,
                                              /*cross_thread=*/ true))
        return PWT_EMIT_ERROR;              // already reported

      /*
        What the rebuild allocated into the container is thread-specific memory
        charged to this thread, and the manager is the thread that will free it.
        Move the charge with the thing it accounts for: take it off our books,
        so ~THD finds them square, and hand it to the manager in cleanup(),
        which runs with every worker joined and before free_containers().
        Adding to the manager's counter from here would race with the manager,
        which is draining at the same time.
      */
      spilled_memory+= worker_thd->status_var.local_memory_used - before;
      worker_thd->status_var.local_memory_used= before;
    }
    else
    {
      table->file->print_error(err, MYF(0));
      return PWT_EMIT_ERROR;
    }
  }

  if (++since_check == PWT_ROW_GANULARITY)
  {
    since_check= 0;
    if (stop_requested())
      return PWT_EMIT_STOP;
  }
  return PWT_EMIT_OK;
}


/**
  @brief
    For the tmp table transport, this only signals the manager that we are done
*/

bool pwt_tmp_table_sink::flush()
{
  /*
    Past this the manager may scan the container, on its own thread, whenever
    it gets to it. Anything this producer had open on the table has to be
    closed by now: handler::ha_rnd_init() asserts that it is, and a worker that
    aggregated into this container left its index open until something closed
    it. Asserted here rather than left to the assert in the manager, which
    fires only when the manager happens to win the race.
  */
  DBUG_ASSERT(!container || !container->table ||
              container->table->file->inited == handler::NONE);
  mysql_mutex_lock(&manager->LOCK_data);
  done= true;
  mysql_cond_signal(&manager->COND_data_avail);
  mysql_mutex_unlock(&manager->LOCK_data);
  return false;
}


/**
  @description
  Give the container back to the manager before it frees it. Called on the
  manager's thread with every worker joined, so a container this worker took
  and never returned -- it failed, or was killed -- would otherwise be freed
  naming a THD that no longer exists.
*/

void pwt_tmp_table_sink::cleanup()
{
  if (container)
  {
    if (container->table)
    {
      container->table->in_use= manager->thd;
      container->table->file->rebind_to_thread();
    }
    container= nullptr;
  }
  /*
    The memory the worker allocated into the container and the manager is about
    to free. Safe to touch the manager's counter here: every worker has been
    joined, so this runs on the manager's own thread.
  */
  if (spilled_memory && manager->thd)
  {
    manager->thd->status_var.local_memory_used+= spilled_memory;
    spilled_memory= 0;
  }
}


/**
  @description
    pwt_tmp_table_source -- the consuming end of the temporary-table transport
    allocate our per worker sink pointers
*/

bool pwt_tmp_table_source::init(THD *thd, pwt_manager *mgr, uint n_workers,
                                uint reclength_arg)
{
  manager=   mgr;
  n_sinks=   n_workers;
  reclength= reclength_arg;
  if (!(sinks= thd->alloc<pwt_tmp_table_sink*>(n_workers)))
    return true;
  for (uint i= 0; i < n_workers; i++)
    sinks[i]= nullptr;
  return false;
}


/**
  @brief
    allocate our per worker sinks
*/

pwt_row_sink *pwt_tmp_table_source::make_sink(THD *thd, uint worker_nr,
                                              pwt_row_container *container)
{
  DBUG_ASSERT(worker_nr < n_sinks);
#ifndef DBUG_OFF
  /*
    No two containers may share their column descriptions. Rebuilding a full
    container on disk is driven from them, and they are allocated out of the
    table's own mem_root, so one param between containers describes only the
    last one built and points into memory freed with it. Every container is the
    same shape, so nothing about the values would give that away -- this is the
    check that fails on the code this replaced.
  */
  for (uint j= 0; j < n_sinks; j++)
    DBUG_ASSERT(!sinks[j] || sinks[j]->container->param != container->param);
#endif
  pwt_tmp_table_sink *s= new (thd->mem_root) pwt_tmp_table_sink;
  if (!s || s->init(manager, this, container))
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(pwt_tmp_table_sink));
    return nullptr;
  }
  sinks[worker_nr]= s;
  return s;
}


/**
  @brief
    Claim the next complete result set, waiting for one if need be.

  @description
    Claimed rather than merely found: the sink is marked taken under the same
    lock that made it visible, so it cannot be scanned twice.

    This is the manager's only wait, so it is also where the team's own state
    is noticed -- a worker killed, a worker failed, and "nothing left to read
    and nobody still running" as end of data. Waiting on active_workers rather
    than on every sink being done is what covers a worker that exits without
    flushing at all.

  @return  0 = *out claimed,  -1 = end of data,  1 = error (reported).
*/

int pwt_tmp_table_source::claim_next_result(pwt_tmp_table_sink **out)
{
  THD *thd= manager->thd;
  PSI_stage_info old_stage;
  struct timespec wait;
  wait.tv_nsec= 0;

  mysql_mutex_lock(&manager->LOCK_data);
  for (;;)
  {
    for (uint i= 0; i < n_sinks; i++)
    {
      if (sinks[i] && sinks[i]->done && !sinks[i]->taken)
      {
        sinks[i]->taken= true;
        mysql_mutex_unlock(&manager->LOCK_data);
        *out= sinks[i];
        return 0;
      }
    }

    int res;
    if ((res= manager->locked__process_manager_wakeup()) ||
        (res=(thd->killed != NOT_KILLED)))
    {
      mysql_mutex_unlock(&manager->LOCK_data);
      return res;
    }
    /*
      wait for a result set, a finishing worker, or a 1s tick to re-check
      killed. ENTER_COND/EXIT_COND publish the "Reading data from parallel
      workers" stage and register the cond so a KILL of the manager wakes it.
    */
    wait.tv_sec= time(0) + 1;
    thd->ENTER_COND(&manager->COND_data_avail, &manager->LOCK_data,
                    &stage_reading_data_from_parallel_worker, &old_stage);
    mysql_cond_timedwait(&manager->COND_data_avail, &manager->LOCK_data, &wait);
    thd->EXIT_COND(&old_stage);                     // unlocks LOCK_data
    mysql_mutex_lock(&manager->LOCK_data);          // re-lock for the next pass
  }
}


/**
  @brief
    Copy the next result row's record image into dst.

  @description
    Scans one finished worker's container to its end, then claims the next.
    Rows therefore arrive in whole-worker runs; the interface promises no
    order and the manager's drain loop assumes none.

  @return
    0 = row copied into dst,  -1 = end of data,  1 = error.
*/

int pwt_tmp_table_source::next_row(uchar *dst)
{
  DBUG_ENTER("pwt_tmp_table_source::next_row");

  for (;;)
  {
    if (cur)
    {
      TABLE *t= cur->container->table;
      int err= t->file->ha_rnd_next(t->record[0]);
      if (likely(!err))
      {
        memcpy(dst, t->record[0], reclength);
        DBUG_RETURN(0);
      }
      if (err != HA_ERR_END_OF_FILE)
      {
        t->file->print_error(err, MYF(0));
        DBUG_RETURN(1);
      }
      release_position();                    // this result set is exhausted
    }

    pwt_tmp_table_sink *next;
    int rc= claim_next_result(&next);
    if (rc)
      DBUG_RETURN(rc);                       // -1 end of data, 1 error

    /*
      Take the container back from the worker that filled it before reading:
      the engine accounts what we read to TABLE::in_use. Safe without further
      locking -- claim_next_result() saw 'done', which that worker published
      under LOCK_data as the last thing it did to this container.
    */
    TABLE *t= next->container->table;
    t->in_use= manager->thd;
    t->file->rebind_to_thread();
    if (int err= t->file->ha_rnd_init(true))
    {
      t->file->print_error(err, MYF(0));
      DBUG_RETURN(1);
    }
    cur= next;
    scan_open= true;
  }
}


/**
  @brief
  End the open scan. The containers are freed just after the workers are
  reaped, and a scan left open would be one the manager still holds on a table
  about to go.
*/

void pwt_tmp_table_source::release_position()
{
  if (cur && scan_open)
    cur->container->table->file->ha_rnd_end();
  cur= nullptr;
  scan_open= false;
}


/*
  @brief
  Create our transport.  Only temporary table transport supported at this time
  pwt_batch_source left commented out for now.
*/

pwt_row_source *pwt_create_transport(THD *thd, pwt_manager *mgr,
                                     uint n_workers, uint reclength)
{
  pwt_tmp_table_source *src= new (thd->mem_root) pwt_tmp_table_source;
  if (!src || src->init(thd, mgr, n_workers, reclength))
  {
    my_error(ER_OUTOFMEMORY, MYF(0), (int) sizeof(pwt_tmp_table_source));
    return nullptr;
  }
  return src;
}