rushdb-query 3.0.8

//! Record properties read by point reads instead of rebuilding the record
//!: the ordering properties of an ordered page, the properties a
//! residual predicate names, and the columns a SQL label scan reads.

use crate::aggregate_rows::Cell;
use crate::planner::{family_id, pids_of};
use crate::{access, normalize, reconstruct, Expr, FieldPred, Json, QueryError, QuerySnapshot};
use rushdb_lmpg::reader::TryRead;
use serde_json::Map;

/// Record properties read with one point read per typed definition instead
/// of rebuilding the record (READ-1: sorting 667K matches, or checking an
/// unindexed `age < 25` on them, rebuilt every one). Covers scalar record
/// properties only; [`Self::resolve`] answers `None` for a system field or a
/// property stored as an array anywhere, and the caller rebuilds the record.
///
/// A record holds one value per property: a write deletes the value of
/// every typed definition of the property before it stores the new one
/// (also when one batch writes the property twice). So the first definition
/// holding a value is the only one, and definitions are probed in the order
/// most records hold them: by descending owner support (under the read's
/// labels). A CSV column inferred as both Int64 and Float64 is read with one
/// point read per row for most rows, not two.
pub struct PointFields(Vec<(String, Vec<u64>)>);

impl PointFields {
    /// Point reads for `names`, each row keyed by the name as given; `None`
    /// when one of them needs the whole record. Definitions are probed by
    /// their tenant-wide support.
    ///
    /// # Errors
    /// Returns storage read errors from the property catalog.
    pub fn resolve<'n>(
        snap: &QuerySnapshot,
        tenant: rushdb_types::ids::TenantId,
        names: impl IntoIterator<Item = &'n str>,
    ) -> Result<Option<Self>, QueryError> {
        Self::resolve_in(snap, tenant, &[], names)
    }

    /// [`Self::resolve`] for records of `labels` (every label when empty):
    /// definitions are probed by their support under those labels.
    ///
    /// # Errors
    /// Returns storage read errors from the property catalog.
    pub fn resolve_in<'n>(
        snap: &QuerySnapshot,
        tenant: rushdb_types::ids::TenantId,
        labels: &[String],
        names: impl IntoIterator<Item = &'n str>,
    ) -> Result<Option<Self>, QueryError> {
        let mut fields = Vec::new();
        for name in names {
            let Some(pids) = scalar_pids(snap, tenant, labels, name)? else {
                return Ok(None);
            };
            fields.push((name.to_owned(), pids));
        }
        Ok(Some(Self(fields)))
    }

    /// The properties `predicate` reads, or `None` when it needs the whole
    /// record.
    pub(super) fn for_predicate(
        snap: &QuerySnapshot,
        tenant: rushdb_types::ids::TenantId,
        labels: &[String],
        predicate: &Expr,
    ) -> Result<Option<Self>, QueryError> {
        let mut names = std::collections::BTreeSet::new();
        if !predicate_fields(predicate, &mut names) {
            return Ok(None);
        }
        Self::resolve_in(snap, tenant, labels, names.iter().map(String::as_str))
    }

    /// The fields of record `ord`: each one's value, as the rebuilt record
    /// would carry it.
    ///
    /// # Errors
    /// Returns storage read errors and corrupt-occurrence errors.
    pub fn row(
        &self,
        snap: &QuerySnapshot,
        tenant: rushdb_types::ids::TenantId,
        ord: u64,
    ) -> Result<Map<String, Json>, QueryError> {
        let mut row = Map::new();
        for (field, pids) in &self.0 {
            if let Some(occ) = held(snap, tenant, ord, pids)? {
                row.insert(
                    field.clone(),
                    reconstruct::decode_value(occ.value, &occ.payload)?,
                );
            }
        }
        Ok(row)
    }

    /// The fields of record `ord` as typed cells, one per field in the order
    /// the fields were resolved (`None`: the record does not hold it).
    pub(crate) fn cells(
        &self,
        snap: &QuerySnapshot,
        tenant: rushdb_types::ids::TenantId,
        ord: u64,
        out: &mut Vec<Option<Cell>>,
    ) -> Result<(), QueryError> {
        out.clear();
        for (_, pids) in &self.0 {
            out.push(match held(snap, tenant, ord, pids)? {
                Some(occ) => Some(Cell::decode(occ.value, &occ.payload)?),
                None => None,
            });
        }
        Ok(())
    }
}

/// The occurrence record `ord` holds under one of `pids` (probed in order).
fn held(
    snap: &QuerySnapshot,
    tenant: rushdb_types::ids::TenantId,
    ord: u64,
    pids: &[u64],
) -> Result<Option<access::Occurrence>, QueryError> {
    for &pid in pids {
        if let Some(occ) = access::occurrence(
            snap,
            tenant,
            rushdb_types::ids::OwnerKind::Record,
            ord,
            pid,
            0,
        )? {
            return Ok(Some(occ));
        }
    }
    Ok(None)
}

/// Collect the property names `expr` reads into `out`; `false` when it needs
/// the whole record (`$id`). A traversal reads none of the record's
/// properties: it starts from the record's ordinal.
pub(super) fn predicate_fields(expr: &Expr, out: &mut std::collections::BTreeSet<String>) -> bool {
    match expr {
        Expr::And(items) | Expr::Or(items) | Expr::Nor(items) => {
            items.iter().all(|item| predicate_fields(item, out))
        }
        Expr::Xor(a, b) => predicate_fields(a, out) && predicate_fields(b, out),
        Expr::Not(inner) => predicate_fields(inner, out),
        Expr::Related(_) => true,
        Expr::Field(FieldPred::Id(_)) => false,
        Expr::Field(
            FieldPred::Eq(name, _)
            | FieldPred::Ne(name, _)
            | FieldPred::Gt(name, _)
            | FieldPred::Gte(name, _)
            | FieldPred::Lt(name, _)
            | FieldPred::Lte(name, _)
            | FieldPred::Contains(name, _)
            | FieldPred::StartsWith(name, _)
            | FieldPred::EndsWith(name, _)
            | FieldPred::Exists(name, _)
            | FieldPred::IsNull(name, _)
            | FieldPred::In(name, _)
            | FieldPred::Nin(name, _)
            | FieldPred::Type(name, _)
            | FieldPred::DatetimeComponent { field: name, .. },
        ) => {
            out.insert(name.clone());
            true
        }
    }
}

/// Typed definitions holding record property `name` as a scalar, in probe
/// order: most owners (under `labels`, every label when empty) first, ties
/// by id. `None` when the name is a system field or any label stores the
/// property as an array: those need the whole record.
fn scalar_pids<R: TryRead>(
    snap: &R,
    tenant: rushdb_types::ids::TenantId,
    labels: &[String],
    name: &str,
) -> Result<Option<Vec<u64>>, QueryError> {
    let norm = normalize(name);
    if norm.is_empty() || rushdb_types::value::is_reserved_name(&norm) {
        return Ok(None);
    }
    let Some(fid) = family_id(snap, tenant, &norm)? else {
        return Ok(Some(Vec::new()));
    };
    let kind = rushdb_types::ids::OwnerKind::Record;
    let mut by_support = Vec::new();
    for pid in pids_of(snap, tenant, fid)? {
        let (arrays, _) = rushdb_lmpg::derived::array_markers(snap, tenant, kind, pid)?;
        if arrays {
            return Ok(None);
        }
        let support = if labels.is_empty() {
            rushdb_lmpg::derived::prop_support_total(snap, tenant, kind, pid)?
        } else {
            let mut support = 0u64;
            let distinct: std::collections::BTreeSet<&String> = labels.iter().collect();
            for label in distinct {
                support = support.saturating_add(rushdb_lmpg::derived::prop_support_label(
                    snap, tenant, kind, pid, label,
                )?);
            }
            support
        };
        by_support.push((std::cmp::Reverse(support), pid));
    }
    by_support.sort_unstable();
    Ok(Some(by_support.into_iter().map(|(_, pid)| pid).collect()))
}

/// The properties a collect with a `fields` projection reads of each child:
/// the projected sources, the ordering fields and the filter's fields, by
/// either spelling. `None` (whole child) without `fields`, or when the
/// filter needs the whole record (`$id`, a traversal).
pub(super) fn collect_selection(
    spec: &crate::CollectSpec,
) -> Option<std::collections::HashSet<String>> {
    if spec.fields.is_empty() {
        return None;
    }
    let mut names = std::collections::BTreeSet::new();
    if let Some(filter) = &spec.filter {
        if !predicate_fields(filter, &mut names) {
            return None;
        }
    }
    for (_, source) in &spec.fields {
        names.insert(source.strip_prefix("$record.").unwrap_or(source).to_owned());
    }
    for (field, _) in &spec.order_by {
        names.insert(field.clone());
    }
    Some(
        names
            .into_iter()
            .flat_map(|name| [normalize(&name), name])
            .collect(),
    )
}