search_sorted in an order of magnitude slower when single element chunk vstacked to the original dataframe

[MichalLebeda]

Checks

• I have checked that this issue has not already been reported.
• I have confirmed this bug exists on the latest version of Polars.

Reproducible example

use std::time::Instant;

use polars::df;
use polars::frame::DataFrame;
use polars::prelude::{col, IntoLazy, SearchSortedSide};

const LEN: i64 = 10_000_000;
const STEP: usize = 10_000;

fn main() {
    // Single chunk
    let data_vec: Vec<i64> = (0..LEN).collect();
    let df = df!("values" => data_vec).unwrap();
    benchmark(&df);

    // Two chunks
    let data_vec: Vec<i64> = (0..LEN).collect();
    let mut df = df!("values" => data_vec).unwrap();
    df.vstack_mut(&df!("values" => [LEN]).unwrap()).unwrap();
    benchmark(&df);
}

fn benchmark(df: &DataFrame) {
    let now = Instant::now();
    let mut n_iters = 0;
    for i in (0..LEN).step_by(STEP) {
        df.clone().lazy().select(vec![col(&"values").search_sorted(i, SearchSortedSide::Left)]).collect().unwrap();
        n_iters += 1;
    }
    println!("df_shape: ({},{}), chunks: {}, n_iters: {:?}, took: {:?}, iter_avg: {:?}", df.width(), df.height(), df.n_chunks(), n_iters, now.elapsed(), now.elapsed()/n_iters);
}

Log output

df_shape: (1,10000000), chunks: 1, n_iters: 1000, took: 168.590893ms, iter_avg: 168.59µs
df_shape: (1,10000001), chunks: 2, n_iters: 1000, took: 8.171000876s, iter_avg: 8.171ms

Issue description

Vstackiing a single element chunk to the dataframe makes search_sorted ~50x slower

Expected behavior

very little performance difference

Installed versions

polars = {version = "0.40.0", features = ["lazy","search_sorted"]}

[ritchie46]

Because we currently need to rechunk the search array. This is much more expensive than the binary search.

[MichalLebeda]

@ritchie46 Thank you for explanation and @orlp for the quick fix!