Polars Hash-Shuffles Rows Between Streaming Stages

    /// Converts a hash to a partition. It is guaranteed that the output is
    /// in the range [0, n_partitions), and that independent HashPartitioners
    /// that we initialized with the same num_partitions and seed return the same
    /// partition.
    #[inline(always)]
    pub fn hash_to_partition(&self, hash: u64) -> usize {
        // Assuming r is a 64-bit random number, we note that
        // r / 2^64 is almost a uniform random number in [0, 1), and thus
        // floor(r * n_partitions / 2^64) is almost a uniform random integer in
        // [0, n_partitions). Despite being written with u128 multiplication this
        // compiles to a single mul / mulhi instruction on x86-x64/aarch64.
        let shuffled = hash.wrapping_mul(self.seed);
        ((shuffled as u128 * self.num_partitions as u128) >> 64) as usize
    }

    #[inline]
    pub fn new(num_partitions: usize, mut seed: u64) -> Self {
        assert!(num_partitions > 0);
        // Make sure seeds bits are properly randomized, and is odd.
        const ARBITRARY1: u64 = 0x85921e81c41226a0;
        const ARBITRARY2: u64 = 0x3bc1d0faba166294;
        const ARBITRARY3: u64 = 0xfbde893e21a73756;
        seed = folded_multiply(seed ^ ARBITRARY1, ARBITRARY2);
        seed = folded_multiply(seed, ARBITRARY3);
        seed |= 1;
        Self {
            seed,
            num_partitions,
        }
    }

        let partitioner = HashPartitioner::new(state.num_pipelines, 0);

/// Represents a DataFrame plus a hash per row, intended for keys in grouping
/// or joining. The hashes may or may not actually be physically pre-computed,
/// this depends per type.
#[derive(Clone, Debug)]
pub enum HashKeys {
    RowEncoded(RowEncodedKeys),
    KeyRows(KeyRowKeys),
    Binview(BinviewKeys),
    Single(SingleKeys),
}

    /// After this call partitions will be extended with the partition for each
    /// hash. Nulls are assigned IdxSize::MAX or a specific partition depending
    /// on whether partition_nulls is true.
    pub fn gen_partitions(
        &self,
        partitioner: &HashPartitioner,
        partitions: &mut Vec<IdxSize>,
        partition_nulls: bool,
    ) {
        unsafe {
            let null_p = if partition_nulls | self.null_is_valid() {
                partitioner.null_partition() as IdxSize
            } else {
                IdxSize::MAX
            };
            partitions.reserve(self.len());
            self.for_each_hash(|_idx, opt_h| {
                partitions.push_unchecked(
                    opt_h
                        .map(|h| partitioner.hash_to_partition(h) as IdxSize)
                        .unwrap_or(null_p),
                );
            });
        }
    }

    fn gen_idxs_per_partition_impl<const BUILD_SKETCHES: bool, const BUILD_F2: bool>(
        &self,
        partitioner: &HashPartitioner,
        partition_idxs: &mut [Vec<IdxSize>],
        sketches: &mut [CardinalitySketch],
        mut f2_sketch: Option<&mut F2Sketch>,
        partition_nulls: bool,
    ) {

/// Groups the rows of `keys` (only those set in `mask`, if given) by partition: one pass for the
/// partitions and counts, a prefix sum, one pass placing the indices. Rows without a hash go to
/// partition 0. `None` with one partition and no mask: all rows, nothing computed.
pub fn split(
    keys: &HashKeys,
    mask: Option<&Bitmap>,
    partitioner: &HashPartitioner,
) -> Option<SplitRows> {
    let num_partitions = partitioner.num_partitions();
    if let Some(mask) = mask {
        assert_eq!(mask.len(), keys.len());
    }

    if num_partitions == 1 {
        let mask = mask?;
        let idxs: Vec<IdxSize> = mask.true_idx_iter().map(|i| i as IdxSize).collect();
        let offsets = vec![0, idxs.len() as IdxSize];
        return Some(SplitRows { idxs, offsets });
    }

    let mut partitions = Vec::new();
    keys.gen_partitions(partitioner, &mut partitions, true);
    // SAFETY: there is one partition per row, and the mask has one bit per row.
    let in_mask = |i: usize| mask.is_none_or(|m| unsafe { m.get_bit_unchecked(i) });

    let mut offsets: Vec<IdxSize> = vec![0; num_partitions + 1];
    for (i, p) in partitions.iter().enumerate() {
        if in_mask(i) {
            offsets[*p as usize + 1] += 1;
        }
    }
    let mut total = 0;
    for offset in &mut offsets {
        total += *offset;
        *offset = total;
    }

    let mut next = offsets[..num_partitions].to_vec();
    let mut idxs: Vec<IdxSize> = vec![0; total as usize];
    for (i, p) in partitions.iter().enumerate() {
        if in_mask(i) {
            let pos = &mut next[*p as usize];
            idxs[*pos as usize] = i as IdxSize;
            *pos += 1;
        }
    }

    Some(SplitRows { idxs, offsets })
}

/// Row indices grouped by partition: one array, increasing within each partition (first/last
/// rely on it), and `num_partitions + 1` offsets into it.
pub struct SplitRows {
    idxs: Vec<IdxSize>,
    offsets: Vec<IdxSize>,
}

impl SplitRows {
    /// The rows of partition `p`, in increasing order.
    pub fn partition(&self, p: usize) -> &[IdxSize] {
        &self.idxs[self.offsets[p] as usize..self.offsets[p + 1] as usize]
    }
}

/// Rows with their hashed keys, and an optional mask of the rows that belong to it.
pub struct FrameWithKeys {
    pub frame: DataFrame,
    pub keys: HashKeys,
    pub mask: Option<Bitmap>,
}

/// A `FrameWithKeys` that can spill, like `SpillFrame`: one `SpillToken`, its size known while
/// spilled.
pub struct SpillFrameWithKeys {
    token: SpillToken<FrameWithKeys>,
    num_rows: usize,
    estimated_size: usize,
}

    /// Appends the rows `rows` of `frame` and `keys`. Panics unless the columns of `frame` are
    /// single chunks.
    ///
    /// # Safety
    /// The rows must be in-bounds of `frame` and `keys`, which have the same number of rows.
    pub unsafe fn gather_extend(&mut self, frame: &DataFrame, keys: &HashKeys, rows: &[IdxSize]) {
        unsafe {
            self.frame.gather_extend(frame, rows, ShareStrategy::Never);
            self.keys.gather_extend(keys, rows);
        }
    }

    /// Stores the cold rows `cold_idxs` of a slice `df` of a morsel of `input_idx`, and adds
    /// their hashes to the sketches. A whole morsel that may be stored as it is and has mostly
    /// cold rows becomes a frame with a mask of its cold rows; otherwise the cold rows are
    /// gathered into the builder of the input, which becomes a frame once it holds a morsel's
    /// worth of rows.
    async fn store_cold_rows(
        &mut self,
        input_idx: usize,
        df: DataFrame,
        keys: HashKeys,
        cold_idxs: &[IdxSize],
        may_store_whole: bool,
        spill_ctx: &MostRecentSpillContext,
    ) {

    /// Builds the results once all input is received. With at most one partition per pipeline
    /// all partitions are built here; otherwise the cold rows are stored per partition, and the
    /// partitions are built while the results are sent.
    fn start_flush(
        &self,
        sink: GroupBySink,
        state: &StreamingExecutionState,
    ) -> PolarsResult<Flush> {
        let (stats, frames, pre_aggs) = sink.finish_input(&self.params, &self.spill_ctx, state)?;
        let num_partitions = self.tuning.num_partitions(&stats, self.num_pipelines);
        let partitioner = HashPartitioner::new(num_partitions, 0);
        let pre_aggs = SplitPreAggs::new(pre_aggs, &partitioner, state)?;
        let num_frames = frames.len();
        let batches = make_batches(frames, self.tuning.batch_bytes());
        let reserve_groups = (stats.estimated_groups * 5 / 4 / num_partitions as u64) as usize;
        let build_at_once = num_partitions <= self.num_pipelines;

    /// The number of partitions: enough for every pipeline to have work, and enough that the
    /// results of any `num_pipelines` partitions fit in memory. Above `num_pipelines` it is a
    /// multiple of `num_pipelines`.
    pub fn num_partitions(&self, stats: &InputStats, num_pipelines: usize) -> usize {
        let work = stats.work();
        if work == 0 {
            return 1;
        }
        if let Some(n) = self.num_partitions {
            return n;
        }

        let t = num_pipelines.max(1) as u64;
        let n_rows = (work / self.min_rows_per_partition).clamp(1, t);
        let result_bytes = (2 * t)
            .saturating_mul(stats.estimated_groups)
            .saturating_mul(stats.group_bytes());
        let n_memory = result_bytes.div_ceil((budget() / 2).max(1));
        if n_memory <= t {
            return n_rows.max(n_memory) as usize;
        }

        let max_partitions = MAX_PARTITIONS.min(self.batch_bytes() as u64 / MIN_STORED_FRAME_BYTES);
        let n_cap = (max_partitions / t * t).max(2 * t);
        let n = n_memory.div_ceil(t).saturating_mul(t).min(n_cap);
        usize::try_from(n).unwrap()
    }

        let (outputs, stored, pre_aggs) = if build_at_once {
            let outputs =
                self.build_partitions(pre_aggs, batches, &partitioner, reserve_groups, state)?;
            actual_groups.add(outputs.iter().map(|out| out.height() as i64).sum());
            (outputs, Vec::new(), SplitPreAggs::default())
        } else {
            let stored = self.store_partitions(batches, &partitioner, state)?;
            (Vec::new(), stored, pre_aggs)
        };

    /// Gathers this partition's rows of every frame of `batch` by index into one frame per input:
    /// reserved exactly, registered for spilling without waiting. Inputs without rows here are
    /// skipped. Each input's columns are those of its first frame in the batch. The frames must
    /// be split.
    pub fn write_batch(&mut self, batch: &[LoadedFrame]) {
        let p = self.index;
        for input in 0..self.frames_per_input.len() {
            let rows_per_frame = || {
                batch
                    .iter()
                    .filter(move |f| f.input == input)
                    .map(move |f| (f, f.partition_rows(p).expect("frames must be split")))
            };
            let num_rows: usize = rows_per_frame().map(|(_, rows)| rows.len()).sum();
            if num_rows == 0 {
                continue;
            }

            let (first, _) = rows_per_frame().next().unwrap();
            let schema = first.frame.frame.schema().clone();
            let mut builder = FrameWithKeysBuilder::new(schema, &first.frame.keys);
            builder.reserve(num_rows);
            for (f, rows) in rows_per_frame() {
                if !rows.is_empty() {
                    unsafe { builder.gather_extend(&f.frame.frame, &f.frame.keys, rows) };
                }
            }

            let frame = SpillFrameWithKeys::new_unregistered(builder.freeze_reset());
            self.spill_ctx.register_no_spill_check(&frame);
            self.estimated_size += frame.estimated_size() as u64;
            self.frames_per_input[input].push_back(frame);
        }
    }

/// Orders frames (those still in memory first, then by input) and cuts them into batches of
/// about `batch_bytes`, at least one frame each.
pub fn make_batches(
    frames: Vec<(usize, SpillFrameWithKeys)>,
    batch_bytes: usize,
) -> Vec<Vec<(usize, SpillFrameWithKeys)>> {

/// Calls `build` on every partition once per batch, with the frames of that batch loaded. While
/// low-priority tasks build the partitions, high-priority tasks load, prepare and split the next
/// batch (each frame once) and free the previous one.
pub fn build_from_batches<S: Send>(

                TaskState::Building { table, stored } => {
                    if let Some((input, frame)) = stored.pop_frame() {
                        let frame = frame.load().await;
                        let frame = self.params.prepare_frame(input, frame, state).await?;
                        table.add_rows(input, &frame, None)?;
                    } else {
                        let frame = table.finish()?;
                        self.actual_groups.add(frame.height() as i64);
                        *task = if frame.height() > 0 {
                            TaskState::Sending { frame, offset: 0 }
                        } else {
                            TaskState::Idle
                        };
                    }
                    memory_manager().spill().await;
                },

    async fn partition_and_sink(
        mut recv: PortReceiver,
        local: &mut LocalBuilder,
        partitioner: HashPartitioner,
        params: &EquiJoinParams,
        state: &StreamingExecutionState,
        spill_ctx: &MostRecentSpillContext,
    ) -> PolarsResult<()> {

struct LocalBuilder {
    // The complete list of morsels and their computed hashes seen by this builder.
    morsels: Vec<(MorselSeq, SpillFrame, HashKeys)>,

    // A cardinality sketch per partition for the keys seen by this builder.
    sketch_per_p: Vec<CardinalitySketch>,

    // The key of each runtime filter seen by this builder.
    key_filters: Vec<KeyFilterBuilder>,

    // morsel_idxs_values_per_p[p][start..stop] contains the offsets into morsels[i]
    // for partition p, where start, stop are:
    // let start = morsel_idxs_offsets[i * num_partitions + p];
    // let stop = morsel_idxs_offsets[(i + 1) * num_partitions + p];
    morsel_idxs_values_per_p: Vec<Vec<IdxSize>>,
    morsel_idxs_offsets_per_p: Vec<usize>,
}

                            let p_morsel_idxs = &l.morsel_idxs_values_per_p[p]
                                [p_morsel_idxs_start..p_morsel_idxs_stop];
                            p_table.insert_keys_subset(keys, p_morsel_idxs, track_unmatchable);
                            p_payload.gather_extend(&payload, p_morsel_idxs, ShareStrategy::Never);

    async fn partition_and_probe(
        mut recv: PortReceiver,
        mut send: PortSender,
        partitions: &[ProbeTable],
        unordered_morsel_seq: &AtomicU64,
        partitioner: HashPartitioner,
        params: &EquiJoinParams,
        state: &StreamingExecutionState,
    ) -> PolarsResult<MorselSeq> {

                    for p in partition_idxs.iter_mut() {
                        p.clear();
                    }
                    hash_keys.gen_idxs_per_partition(
                        &partitioner,
                        &mut partition_idxs,
                        &mut [],
                        None,
                        emit_unmatched,
                    );

                if params.preserve_order_probe {
                    // To preserve the order we can't do bulk probes per partition and must follow
                    // the order of the probe morsel. We can still group probes that are
                    // consecutively on the same partition.
                    probe_partitions.clear();
                    hash_keys.gen_partitions(&partitioner, &mut probe_partitions, emit_unmatched);