zeppelin-embed 0.6.0

In-process vector, lexical, and hybrid search for macOS and Apple silicon
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

//! Conservative physical choices inside one already validated pattern region.

use super::super::plan::{MAX_PLAN_DEPTH, Operator, OperatorKind, PlanNodeId};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum JoinStrategy {
    Nested,
    Hash,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum BuildSide {
    Left,
    Right,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) struct JoinPlan {
    pub(super) strategy: JoinStrategy,
    pub(super) build: BuildSide,
}

pub(super) fn join_plan(
    shared_slots: usize,
    operators: &[Operator<'_>],
    left: PlanNodeId,
    right: PlanNodeId,
) -> JoinPlan {
    #[cfg(test)]
    if let Some(choice) = FORCED_JOIN.with(std::cell::Cell::get) {
        return JoinPlan {
            strategy: choice,
            build: BuildSide::Right,
        };
    }
    let strategy = if shared_slots == 0 {
        JoinStrategy::Nested
    } else {
        JoinStrategy::Hash
    };
    let build = match (
        safe_source_rank(operators, left, 0),
        safe_source_rank(operators, right, 0),
    ) {
        (Some(left), Some(right)) if left < right => BuildSide::Left,
        _ => BuildSide::Right,
    };
    JoinPlan { strategy, build }
}

fn safe_source_rank(operators: &[Operator<'_>], node: PlanNodeId, depth: usize) -> Option<u8> {
    if depth >= MAX_PLAN_DEPTH {
        return None;
    }
    let operator = operators.get(node.0 as usize)?;
    match operator.kind {
        OperatorKind::LookupNode { .. }
        | OperatorKind::LookupRelationship { .. }
        | OperatorKind::LookupKey { .. } => {
            safe_input_chain(operators, operator, depth).then_some(0)
        }
        OperatorKind::ScanNodes { label: Some(_), .. } => {
            safe_input_chain(operators, operator, depth).then_some(1)
        }
        OperatorKind::ScanNodes { label: None, .. } => {
            safe_input_chain(operators, operator, depth).then_some(2)
        }
        OperatorKind::Collect | OperatorKind::Expand { .. } => operator
            .inputs
            .first()
            .and_then(|input| safe_source_rank(operators, *input, depth + 1)),
        OperatorKind::BoundedExpand {
            edge_predicate: None,
            completed_edge_predicate: None,
            ..
        } => operator
            .inputs
            .first()
            .and_then(|input| safe_source_rank(operators, *input, depth + 1)),
        _ => None,
    }
}

fn safe_input_chain(operators: &[Operator<'_>], operator: &Operator<'_>, depth: usize) -> bool {
    let Some(input) = operator.inputs.first().copied() else {
        return false;
    };
    if depth + 1 >= MAX_PLAN_DEPTH {
        return false;
    }
    let Some(input) = operators.get(input.0 as usize) else {
        return false;
    };
    match input.kind {
        OperatorKind::Unit => input.inputs.is_empty(),
        OperatorKind::LookupNode { .. }
        | OperatorKind::LookupRelationship { .. }
        | OperatorKind::LookupKey { .. }
        | OperatorKind::ScanNodes { .. } => safe_input_chain(operators, input, depth + 1),
        OperatorKind::Collect | OperatorKind::Expand { .. } => {
            safe_input_chain(operators, input, depth + 1)
        }
        OperatorKind::BoundedExpand {
            edge_predicate: None,
            completed_edge_predicate: None,
            ..
        } => safe_input_chain(operators, input, depth + 1),
        _ => false,
    }
}

#[cfg(test)]
std::thread_local! {
    static FORCED_JOIN: std::cell::Cell<Option<JoinStrategy>> = const { std::cell::Cell::new(None) };
}

#[cfg(test)]
pub(super) fn force_join_strategy(choice: Option<JoinStrategy>) {
    FORCED_JOIN.with(|forced| forced.set(choice));
}