Polars Publishes a Join's Build Side as a Runtime Filter

//! Let a hash join tell the parquet scan under its probe side which keys the build
//! side holds.
//!
//! A join whose one side is known to be much smaller than a scan under its other
//! side gets that side forced as build side; one whose side is only estimated to be
//! much smaller gets it preferred. Every probe key that
//! reads a parquet scan column unchanged gets a dynamic predicate on that scan. The
//! join publishes the range of its build keys once the build is done. The scan
//! skips the row groups outside the range and drops the rows outside it. When the
//! build keys are estimated to be a small share of the scan's distinct keys the
//! join also publishes a bloom filter over them, which the scan probes per row
//! instead.

pub(super) fn attach_join_runtime_filters(
    root: Node,
    ir_arena: &mut Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
) {
    if !polars_config::config().join_runtime_filters() {
        return;
    }
    // Inputs before their join, so a join lower in a probe chain is forced before
    // an outer one tries to carry a predicate through it.
    let mut joins = Vec::new();
    let mut has_pruning_scan = false;
    let mut stack = vec![root];
    while let Some(node) = stack.pop() {
        let ir = ir_arena.get(node);
        match ir {
            IR::Join { .. } => joins.push(node),
            IR::Scan { scan_type, .. } => has_pruning_scan |= skips_batches(scan_type),
            _ => {},
        }
        ir.copy_inputs(&mut stack);
    }
    if !has_pruning_scan {
        return;
    }

fn skips_batches(scan_type: &FileScanIR) -> bool {
    scan_type
        .flags()
        .contains(ScanFlags::SKIPS_BATCHES_BY_STATISTICS)
}

/// Whether the join may publish a filter or be crossed by one: an inner, semi or
/// anti equi join the streaming engine can run as a hash join that blocks its
/// probe side until the build is done. Sorted inputs may still make it a merge
/// join, which drops the filter. A key that may evaluate differently each time
/// gives no filter, as the join evaluates it again when it builds.
fn is_eligible_join(options: &JoinOptionsIR, expr_arena: &Arena<AExpr>) -> bool {
    let args = &options.args;
    let JoinTypeOptionsIR::Equi { on, .. } = &options.options else {
        return false;
    };
    !on.is_empty()
        && on.iter().all(|(left, right)| {
            !is_inherently_nondeterministic(left.node(), expr_arena)
                && !is_inherently_nondeterministic(right.node(), expr_arena)
        })
        && (args.how == JoinType::Inner || args.how.is_semi_anti())
        && args.maintain_order == MaintainOrderJoin::None
        && !args.nulls_equal
        && args.slice.is_none()
        && !args.validation.needs_checks()
}

//! A semi join publishes from either side and an anti join from its left side
//! only, as the rows of its right side that match no left key change nothing. A
//! semi join may prefer its right side on an estimate; a left side is only built
//! when forced, as building it keeps its rows. It is not forced when the right keys
//! are known to be much fewer than its estimated rows, as a right build keeps only
//! those.

/// The ways a filtered side may be built from: forced when it has a bound, and
/// preferred on its estimate.
fn build_candidates(left: bool, stats: &NodeStats) -> Vec<BuildCandidate> {
    if stats.filtered >= stats.unfiltered {
        return Vec::new();
    }
    let forced = stats.max_rows().map(|rows| BuildCandidate {
        left,
        rows,
        forced: true,
    });
    let preferred = BuildCandidate {
        left,
        rows: stats.filtered,
        forced: false,
    };
    forced.into_iter().chain([preferred]).collect()
}

/// A side has to be this many times smaller to be picked as the build side.
pub(super) const LOPSIDED_FACTOR: f64 = 8.0;

        (!filters.is_empty() && *probe_rows >= LOPSIDED_FACTOR * side.rows).then_some(side)

fn trace_probe_keys(
    probe_input: Node,
    probe_keys: Vec<Option<PlSmallStr>>,
    ir_arena: &Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
    scratch: &mut UnitVec<Node>,
) -> Vec<TracedKey> {
    let probe_schema = ir_arena.get(probe_input).schema(ir_arena).into_owned();
    let mut traced = Vec::new();
    for (key_idx, name) in probe_keys.into_iter().enumerate() {
        let Some(name) = name else { continue };
        if !probe_schema.get(&name).is_some_and(supports_runtime_range) {
            continue;
        }
        let column = expr_arena.add(AExpr::Column(name));
        let (dyn_node, pred) = new_dynamic_pred(column, expr_arena, None);
        let predicate = ExprIR::from_node(dyn_node, expr_arena);
        let origins = scan_origins(probe_input, predicate, ir_arena, expr_arena, scratch);

                IR::SimpleProjection { .. }
                | IR::Filter { .. }
                | IR::Select { .. }
                | IR::HStack { .. } => {
                    match push_past(
                        node,
                        &mut predicate,
                        ir_arena,
                        expr_arena,
                        scratch,
                        maintain_errors,
                    ) {
                        Ok(Some(input)) => node = input,
                        _ => continue 'paths,
                    }
                },
                // A slice over the union would take other rows once some are pruned.
                IR::Union { inputs, options } if options.slice.is_none() => {
                    // Each input rewrites its own expression nodes; the filter they
                    // refer to stays the one the join publishes.
                    for &input in inputs.iter().rev() {
                        let copy = deep_clone_ae(predicate.node(), expr_arena);
                        paths.push((input, ExprIR::from_node(copy, expr_arena)));
                    }
                    continue 'paths;
                },

                IR::Filter {
                    predicate: filter, ..
                } if has_window(std::slice::from_ref(filter), expr_arena) => continue 'paths,

fn attach_to_scan(
    scan: Node,
    predicate: ExprIR,
    ir_arena: &mut Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
) {
    let IR::Scan {
        predicate: existing,
        predicate_file_skip_applied,
        ..
    } = ir_arena.get_mut(scan)
    else {
        unreachable!()
    };
    // Files are skipped by the scan predicate after these filters are attached.
    debug_assert!(predicate_file_skip_applied.is_none());
    *existing = Some(match existing.take() {
        None => predicate,
        Some(existing) => {
            let node = expr_arena.add(AExpr::BinaryExpr {
                left: existing.node(),
                op: Operator::And,
                right: predicate.node(),
            });
            ExprIR::from_node(node, expr_arena)
        },
    });
}

/// A predicate over `node` whose value a producer sets at run time, evaluated
/// per row wherever it lands. A reader reads in `read_order` when it can.
pub fn new_dynamic_pred(
    node: Node,
    arena: &mut Arena<AExpr>,
    read_order: Option<ReadOrder>,
) -> (Node, DynamicPred) {
    let pred = DynamicPred::new(read_order);
    let function = IRFunctionExpr::DynamicPred {
        pred: pred.downgrade(),
    };
    let options = function.function_options();
    let aexpr = AExpr::Function {
        input: vec![ExprIR::from_node(node, arena)],
        function,
        options,
    };

    (arena.add(aexpr), pred)
}

    pub fn set(&self, pred: Arc<dyn PredicateExpr>) {
        {
            let mut guard = self.inner.pred.write().unwrap();
            *guard = Some(pred);
        }
        self.inner.is_set.store(true, Ordering::Release);
    }

    pub fn is_set(&self) -> bool {
        self.inner.is_set.load(Ordering::Acquire)
    }

/// The build-side key at `key_idx` of the join's `on`, published through
/// `pred` for the probe side: its range, and a bloom filter over its values
/// when `bloom_keys` gives the number of distinct keys to size it for.
#[derive(Clone, Debug, PartialEq, Hash)]
#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
pub struct RuntimeFilter {
    pub key_idx: usize,
    pub pred: DynamicPred,
    pub bloom_keys: Option<usize>,
    /// The probe side's distinct keys, when the plan can estimate them.
    pub probe_distinct: Option<usize>,
}

/// Largest ratio of distinct build keys to distinct probe keys for which a
/// bloom filter is published.
pub const MAX_BUILD_PROBE_DISTINCT_RATIO: f64 = 0.3;

pub(super) struct KeyFilterBuilder {
    range: KeyRange,
    bloom: Option<BloomBuilder>,
}

impl KeyRange {
    /// Widen the range to cover the non-null values of `column`.
    fn extend(&mut self, column: &Column) -> PolarsResult<()> {
        let min = column.min_reduce()?;
        let max = column.max_reduce()?;
        if !min.is_null() && !max.is_null() {
            self.merge(Self {
                bounds: Some((min, max)),
            });
        }
        Ok(())
    }

impl PredicateExpr for KeyRange {
    /// Whether each value lies in the range. Nulls never do.
    fn evaluate(&self, columns: &[Column]) -> PolarsResult<Option<Column>> {
        let Some(bounds) = &self.bounds else {
            return Ok(None);
        };
        let column = columns[0].as_materialized_series();
        let Some((lo, hi)) = bounds_for(bounds, column.dtype()) else {
            return Ok(None);
        };
        let mask = (column.gt_eq(&lo)? & column.lt_eq(&hi)?).fill_null_with_values(false)?;
        Ok(Some(mask.into_column()))
    }

    fn evaluate_stats(
        &self,
        min: &Column,
        max: &Column,
        _null_count: &Column,
    ) -> PolarsResult<Option<Column>> {
        let Some(bounds) = &self.bounds else {
            return Ok(Some(Column::new_scalar(
                min.name().clone(),
                Scalar::from(true),
                min.len(),
            )));
        };
        let min = min.as_materialized_series();
        let max = max.as_materialized_series();
        let Some((lo, hi)) = bounds_for(bounds, min.dtype()) else {
            return Ok(None);
        };
        // A batch is skipped when it lies entirely below or above the range. An
        // unknown statistic is null and settles nothing.
        let skip = max.lt(&lo)? | min.gt(&hi)?;
        Ok(Some(skip.fill_null_with_values(false)?.into_column()))
    }

    fn try_transition_to_build(
        &mut self,
        recv: &[PortState],
        params: &mut EquiJoinParams,
        state: &StreamingExecutionState,
        spill_ctx: &MostRecentSpillContext,
    ) -> PolarsResult<Option<BuildState>> {
        if let Some(left) = self.only_side {
            let idx = if left { 0 } else { 1 };
            let len = self.len(left);
            if len >= params.sample_limit {
                if config::verbose() {
                    eprintln!("preferred build side reached the sample limit, sampling both sides");
                }
            } else if recv[idx] == PortState::Done {
                if config::verbose() {
                    eprintln!("preferred build side done with {len} rows, publishing its ranges");
                }
                self.publish_runtime_filters(left, params, state)?;

    /// Hand every build key to the runtime filters. Filters set from a sample
    /// keep their value, whichever side is built; filters of a side that was
    /// not built get a predicate that skips nothing.
    fn publish_runtime_filters(&mut self, params: &mut EquiJoinParams) {
        if !params.publishes_runtime_filters() {
            params.runtime_filters.publish_nothing();
            return;
        }
        let locals = self
            .local_builders
            .iter_mut()
            .map(|l| std::mem::take(&mut l.key_filters));
        params.runtime_filters.publish_merged(locals);
    }

        let locals = RAYON.install(|| {
            morsels
                .par_chunks(chunk_len)
                .map(|chunk| {
                    let mut builders = self.new_builders();
                    for morsel in chunk {
                        let df = morsel.df_blocking();
                        let keys = ASYNC.block_on(select_key_columns(&df, key_selectors, state))?;
                        self.extend(&keys, &mut builders)?;
                    }
                    PolarsResult::Ok(builders)
                })
                .collect::<PolarsResult<Vec<_>>>()
        })?;
        self.publish_merged(locals);

/// Bits per key a bloom filter is sized for.
const BLOOM_BITS_PER_KEY: usize = 8;
/// A bloom filter with fewer bits per distinct key than this is not published.
const BLOOM_MIN_BITS_PER_KEY: usize = 4;
/// Smallest bloom filter that is built.
const BLOOM_MIN_BYTES: usize = 64 << 10;
/// Largest bloom filter that is published.
const BLOOM_MAX_BYTES: usize = 32 << 20;

        let distinct = sketch.estimate();
        let weak = spec
            .probe_distinct
            .is_some_and(|probe| distinct as f64 > probe as f64 * MAX_BUILD_PROBE_DISTINCT_RATIO);
        let bloom = if weak {
            None
        } else if planned {

    fn evaluate(&self, columns: &[Column]) -> PolarsResult<Option<Column>> {
        let Some(KeyBloom { spec, bloom }) = &self.bloom else {
            return self.range.evaluate(columns);
        };
        let column = &columns[0];
        if column.dtype() != &spec.dtype {
            return Ok(None);
        }
        let mut mask = BitmapBuilder::with_capacity(column.len());
        spec.hash_keys(column).for_each_hash(|_, hash| {
            mask.push(hash.is_some_and(|hash| bloom.contains(hash)));
        });
        Ok(Some(
            BooleanChunked::from_bitmap(column.name().clone(), mask.freeze()).into_column(),
        ))
    }

        F::DynamicPred { pred, .. } => {
            map_as_slice!(misc::dynamic_pred, &pred)
        },

                IRFunctionExpr::DynamicPred { pred, .. } => {
                    let target = resolve_stat_target(input[0].node(), arena, schema)?;
                    let dtype = target_leaf_dtype(&target, schema)?;
                    if !supports_runtime_range(dtype) {
                        return None;
                    }
                    let function = IRFunctionExpr::DynamicSkipBatch { pred: pred.clone() };
                    let options = function.function_options();
                    let input = [
                        target.min(arena),
                        target.max(arena),
                        target.null_count(arena),
                    ]
                    .into_iter()
                    .map(|b| ExprIR::from_node(b.node(), arena))
                    .collect();

/// Whether a runtime key range may skip batches of a column with this dtype.
///
/// The range is compared against the batch's min/max with polars' ordering, so the
/// type must be one whose parquet statistics are ordered the same way. Floats are out
/// because the statistics leave out NaN; 128-bit integers because their statistics
/// are byte-wise.
pub(crate) fn supports_runtime_range(dtype: &DataType) -> bool {
    use DataType as D;
    match dtype {
        D::Boolean
        | D::UInt8
        | D::UInt16
        | D::UInt32
        | D::UInt64
        | D::Int8
        | D::Int16
        | D::Int32
        | D::Int64
        | D::String
        | D::Binary
        | D::Date
        | D::Datetime(..)
        | D::Duration(_)
        | D::Time => true,
        #[cfg(feature = "dtype-decimal")]
        D::Decimal(..) => true,
        _ => false,
    }
}

pub enum RuntimeRange {
    /// Not published yet. Every batch is kept; a later file may see a range.
    Pending,
    /// Never published. Every batch is kept.
    Disabled,
    /// No value can match. Every batch is skipped.
    Empty,
    /// Only values in `lo..=hi` can match, a missing bound does not limit that
    /// side. Nulls match when `nulls_match`.
    Range {
        lo: Option<Scalar>,
        hi: Option<Scalar>,
        nulls_match: bool,
    },
}