uqa-execution 0.4.6

Volcano physical operators with row-batch pipelines
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
121
122
123
124
125
126
127

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Exact cache changes reserve complete key and identity nodes before publishing a staged row.

use super::{
    keys::{document_key, ExactKey, FieldSet},
    resource_error, CommandStoredDocument, DocId, SQLError, StorageReadControl,
};
use uqa_core::memory::{BudgetedMap, PreparedMapEntry};

pub(super) struct CommandExactIndex {
    doc_ids_by_key: BudgetedMap<ExactKey, KeyRows>,
}

struct KeyRows {
    ids: BudgetedMap<DocId, ()>,
}

pub(super) struct PreparedChange {
    previous: Option<ExactKey>,
    replacement: Option<Replacement>,
}

enum Replacement {
    Existing {
        key: ExactKey,
        entry: PreparedMapEntry<DocId, ()>,
    },
    New(PreparedMapEntry<ExactKey, KeyRows>),
}

impl CommandExactIndex {
    pub(super) fn build(
        documents: &BudgetedMap<DocId, Option<CommandStoredDocument>>,
        fields: &FieldSet,
        control: &StorageReadControl,
    ) -> Result<Self, SQLError> {
        let mut index = Self {
            doc_ids_by_key: BudgetedMap::new(control.memory()),
        };
        for (id, document) in documents {
            control.check().map_err(resource_error)?;
            let Some(document) = document else { continue };
            let change = index.prepare(*id, None, Some(document), fields, control)?;
            index.apply(*id, change);
        }
        Ok(index)
    }

    pub(super) fn candidates(&self, key: &[u8]) -> Option<&BudgetedMap<DocId, ()>> {
        self.doc_ids_by_key.get(key).map(|rows| &rows.ids)
    }

    pub(super) fn prepare(
        &self,
        id: DocId,
        previous: Option<&CommandStoredDocument>,
        replacement: Option<&CommandStoredDocument>,
        fields: &FieldSet,
        control: &StorageReadControl,
    ) -> Result<Option<PreparedChange>, SQLError> {
        let previous = previous
            .map(|document| document_key(document.fields.as_ref(), fields, control))
            .transpose()?;
        let replacement = replacement
            .map(|document| document_key(document.fields.as_ref(), fields, control))
            .transpose()?;
        if previous == replacement {
            return Ok(None);
        }
        let replacement = replacement
            .map(|key| {
                if let Some(rows) = self.doc_ids_by_key.get(key.bytes()) {
                    let entry = rows.ids.prepare_entry(id, ()).map_err(resource_error)?;
                    Ok(Replacement::Existing { key, entry })
                } else {
                    let mut rows = KeyRows {
                        ids: BudgetedMap::new(control.memory()),
                    };
                    rows.ids.insert(id, ()).map_err(resource_error)?;
                    self.doc_ids_by_key
                        .prepare_entry(key, rows)
                        .map(Replacement::New)
                        .map_err(resource_error)
                }
            })
            .transpose()?;
        Ok(Some(PreparedChange {
            previous,
            replacement,
        }))
    }

    pub(super) fn apply(&mut self, id: DocId, change: Option<PreparedChange>) {
        let Some(change) = change else { return };
        if let Some(key) = change.previous {
            let rows = self
                .doc_ids_by_key
                .get_mut(key.bytes())
                .expect("cached previous key");
            assert!(
                rows.ids.remove(&id).is_some(),
                "cached previous document identity"
            );
            if rows.ids.is_empty() {
                self.doc_ids_by_key.remove(key.bytes());
            }
        }
        match change.replacement {
            Some(Replacement::Existing { key, entry }) => {
                self.doc_ids_by_key
                    .get_mut(key.bytes())
                    .expect("prepared existing key")
                    .ids
                    .insert_prepared(entry);
            }
            Some(Replacement::New(entry)) => {
                self.doc_ids_by_key.insert_prepared(entry);
            }
            None => {}
        }
    }
}