polars-stream 0.55.2

use std::collections::BTreeMap;
use std::ops::{Range, RangeBounds};

use polars_core::frame::DataFrame;
use polars_core::prelude::*;
use polars_core::schema::SchemaRef;
use polars_core::series::Series;
use polars_utils::range::check_range;

use crate::pipe::PortReceiver;

#[derive(Clone, Debug)]
pub(super) struct DataFrameSearchBuffer {
    schema: SchemaRef,
    dfs_at_offsets: BTreeMap<usize, DataFrame>,
    total_rows: usize,
    skip_rows: usize,
    frozen: bool,
}

impl DataFrameSearchBuffer {
    pub(super) fn empty_with_schema(schema: SchemaRef) -> Self {
        DataFrameSearchBuffer {
            schema,
            dfs_at_offsets: BTreeMap::new(),
            total_rows: 0,
            skip_rows: 0,
            frozen: false,
        }
    }

    pub(super) fn height(&self) -> usize {
        self.total_rows
    }

    /// Get the `row_index`th value from the `column`.
    ///
    /// SAFETY: Caller must ensure that `row_index` is within bounds.
    pub(super) unsafe fn get_unchecked(&self, column: &str, row_index: usize) -> AnyValue<'_> {
        unsafe { self.get_bypass_validity(column, row_index, false) }
    }

    /// Get the `row_index`th value from the `column` potentially bypassing its
    /// validity bitmap.
    ///
    /// SAFETY: Caller must ensure that `row_index` is within bounds.
    pub(super) unsafe fn get_bypass_validity(
        &self,
        column: &str,
        row_index: usize,
        bypass_validity: bool,
    ) -> AnyValue<'_> {
        debug_assert!(row_index < self.total_rows);
        let first_offset = match self.dfs_at_offsets.first_key_value() {
            Some((offset, _)) => *offset,
            None => 0,
        };
        let buf_index = self.skip_rows + first_offset + row_index;
        let (df_offset, df) = self.dfs_at_offsets.range(..=buf_index).next_back().unwrap();
        let series_index = buf_index - df_offset;
        let series = df.column(column).unwrap().as_materialized_series();
        unsafe { series_get_bypass_validity(series, series_index, bypass_validity) }
    }

    pub(super) fn push_df(&mut self, df: DataFrame) {
        assert!(!self.frozen);
        let added_rows = df.height();
        let offset = match self.dfs_at_offsets.last_key_value() {
            Some((last_key, last_df)) => last_key + last_df.height(),
            None => 0,
        };
        self.dfs_at_offsets.insert(offset, df);
        self.total_rows += added_rows;
    }

    pub(super) fn split_at(&mut self, mut at: usize) -> Self {
        at = at.clamp(0, self.total_rows);
        let mut top = self.clone();
        top.total_rows = at;
        top.frozen = true;
        self.skip_rows += at;
        self.total_rows -= at;
        self.gc();
        top
    }

    pub(super) fn slice(mut self, offset: usize, len: usize) -> Self {
        self.skip_rows += offset;
        self.total_rows -= offset;
        self.total_rows = usize::min(self.total_rows, len);
        self.frozen = true;
        self.gc();
        self
    }

    pub(super) fn into_df(self) -> DataFrame {
        let mut acc = DataFrame::empty_with_schema(&self.schema);
        for df in self.dfs_at_offsets.into_values() {
            acc.vstack_mut_owned(df).unwrap();
        }
        acc.slice(self.skip_rows as i64, self.total_rows)
    }

    fn gc(&mut self) {
        while let Some((_, df)) = self.dfs_at_offsets.first_key_value()
            && self.skip_rows > df.height()
        {
            let (_, gc_df) = self.dfs_at_offsets.pop_first().unwrap();
            self.skip_rows -= gc_df.height();
        }
    }

    pub(super) fn is_empty(&self) -> bool {
        self.total_rows == 0
    }

    /// Find the index of the first item in the buffer that satisfies `predicate`,
    /// assuming it is first always false and then always true.
    pub(super) fn binary_search<P, R>(&self, predicate: P, key_col_name: &str, range: R) -> usize
    where
        P: Fn(&AnyValue<'_>) -> bool,
        R: RangeBounds<usize>,
    {
        self.binary_search_binary_offset_bypass_validity(predicate, key_col_name, range, false)
    }

    /// Find the index of the first item in the buffer that satisfies `predicate`,
    /// assuming it is first always false and then always true.
    pub(super) fn binary_search_binary_offset_bypass_validity<P, R>(
        &self,
        predicate: P,
        key_col_name: &str,
        range: R,
        binary_offset_bypass_validity: bool,
    ) -> usize
    where
        P: Fn(&AnyValue<'_>) -> bool,
        R: RangeBounds<usize>,
    {
        let Range {
            start: mut lower,
            end: mut upper,
        } = check_range(range, ..self.height());
        while lower < upper {
            let mid = (lower + upper) / 2;
            let mid_val = unsafe {
                self.get_bypass_validity(key_col_name, mid, binary_offset_bypass_validity)
            };
            if predicate(&mid_val) {
                upper = mid;
            } else {
                lower = mid + 1;
            }
        }
        lower
    }

    pub(super) async fn stop_and_buffer_from_pipe(&mut self, port: Option<&mut PortReceiver>) {
        stop_and_buffer_pipe_contents(port, &mut |df| self.push_df(df)).await
    }

    pub(super) fn select<I, S>(&self, columns: I) -> Self
    where
        I: IntoIterator<Item = S> + Clone,
        S: AsRef<str>,
    {
        let select_map = |df: &DataFrame| df.select(columns.clone()).expect("projection failed");
        let dfs_at_offsets = self
            .dfs_at_offsets
            .iter()
            .map(|(offset, df)| (*offset, select_map(df)))
            .collect();
        DataFrameSearchBuffer {
            schema: self
                .schema
                .try_project(columns)
                .expect("projection failed")
                .into(),
            dfs_at_offsets,
            total_rows: self.total_rows,
            skip_rows: self.skip_rows,
            frozen: self.frozen,
        }
    }
}

/// Tell the sender to this port to stop, and buffer everything that is still in the pipe.
pub(super) async fn stop_and_buffer_pipe_contents<F>(
    port: Option<&mut PortReceiver>,
    buffer_morsel: &mut F,
) where
    F: FnMut(DataFrame),
{
    let Some(port) = port else {
        return;
    };

    while let Ok(morsel) = port.recv().await {
        morsel.source_token().stop();
        let (sf, _, _, _) = morsel.into_inner();
        buffer_morsel(sf.into_df().await);
    }
}

/// Get value from series bypassing the validity bitmap.
///
/// SAFETY: Caller must ensure that `index` is within bounds of `s`.
unsafe fn series_get_bypass_validity<'a>(
    s: &'a Series,
    index: usize,
    binary_offset_bypass_validity: bool,
) -> AnyValue<'a> {
    debug_assert!(index < s.len());
    if binary_offset_bypass_validity {
        let arr = s.binary_offset().unwrap();
        unsafe { arr.get_any_value_bypass_validity(index) }
    } else {
        unsafe { s.get_unchecked(index) }
    }
}