Skip to main content

linkrs_query/executor/streaming/
slot.rs

1use linkrs_core::DataType;
2use std::collections::HashMap;
3
4pub type SlotId = usize;
5
6/// Optional origin information for a slot (which plan node produced it).
7#[derive(Debug, Clone, PartialEq, Eq, Hash)]
8pub enum SlotOrigin {
9    /// Column from a storage scan
10    StorageColumn(String),
11    /// Computed expression output
12    Computed,
13    /// System-generated column (e.g. _vid, _expand_vertex)
14    System,
15    /// Join output column
16    JoinOutput,
17    /// Aggregation output
18    AggregationOutput,
19    /// Traversal result (vertex/edge/path)
20    TraversalResult,
21}
22
23#[derive(Debug, Clone)]
24pub struct SlotInfo {
25    pub slot_id: SlotId,
26    pub name: String,
27    /// Optional alias (e.g. AS alias_name)
28    pub alias: Option<String>,
29    pub data_type: Option<DataType>,
30    /// Whether the column can contain NULL values
31    pub nullable: bool,
32    /// Where this slot originates from
33    pub origin: Option<SlotOrigin>,
34}
35
36#[derive(Debug, Clone)]
37pub struct SlotLayout {
38    pub slots: Vec<SlotInfo>,
39    pub name_to_slot: HashMap<String, SlotId>,
40}
41
42impl SlotLayout {
43    pub fn new(slots: Vec<SlotInfo>) -> Self {
44        let name_to_slot = slots
45            .iter()
46            .enumerate()
47            .map(|(i, info)| (info.name.clone(), i))
48            .collect();
49        Self {
50            slots,
51            name_to_slot,
52        }
53    }
54
55    pub fn from_names(names: &[String]) -> Self {
56        let slots: Vec<SlotInfo> = names
57            .iter()
58            .enumerate()
59            .map(|(i, name)| SlotInfo {
60                slot_id: i,
61                name: name.clone(),
62                alias: None,
63                data_type: None,
64                nullable: true,
65                origin: None,
66            })
67            .collect();
68        Self::new(slots)
69    }
70
71    pub fn from_names_and_types(names: &[String], types: &[Option<DataType>]) -> Self {
72        let slots: Vec<SlotInfo> = names
73            .iter()
74            .enumerate()
75            .map(|(i, name)| SlotInfo {
76                slot_id: i,
77                name: name.clone(),
78                alias: None,
79                data_type: types.get(i).cloned().unwrap_or(None),
80                nullable: true,
81                origin: None,
82            })
83            .collect();
84        Self::new(slots)
85    }
86
87    /// Create a layout with full slot metadata.
88    pub fn from_slots(slots: Vec<SlotInfo>) -> Self {
89        Self::new(slots)
90    }
91
92    pub fn len(&self) -> usize {
93        self.slots.len()
94    }
95
96    pub fn is_empty(&self) -> bool {
97        self.slots.is_empty()
98    }
99
100    pub fn slot_id(&self, name: &str) -> Option<SlotId> {
101        self.name_to_slot.get(name).copied()
102    }
103
104    pub fn slot_info(&self, id: SlotId) -> Option<&SlotInfo> {
105        self.slots.get(id)
106    }
107
108    /// Return the names of all slots.
109    pub fn names(&self) -> Vec<String> {
110        self.slots.iter().map(|s| s.name.clone()).collect()
111    }
112
113    /// Resolve a slot by name, with alias fallback.
114    pub fn resolve(&self, name_or_alias: &str) -> Option<SlotId> {
115        // First try exact name match
116        if let Some(&id) = self.name_to_slot.get(name_or_alias) {
117            return Some(id);
118        }
119        // Then try alias match
120        self.slots
121            .iter()
122            .find(|s| s.alias.as_deref() == Some(name_or_alias))
123            .map(|s| s.slot_id)
124    }
125}
126
127/// Combine two slot layouts side-by-side (for joins).
128/// Left slots preserve their IDs; right slots get new IDs offset by left count.
129pub fn combine_layouts(left: &SlotLayout, right: &SlotLayout) -> SlotLayout {
130    let mut slots = left.slots.clone();
131    let offset = slots.len();
132    for info in &right.slots {
133        slots.push(SlotInfo {
134            slot_id: info.slot_id + offset,
135            name: info.name.clone(),
136            alias: info.alias.clone(),
137            data_type: info.data_type.clone(),
138            nullable: info.nullable,
139            origin: info.origin.clone(),
140        });
141    }
142    SlotLayout::new(slots)
143}
144
145/// Resolve naming conflicts in a combined layout by appending suffixes.
146pub fn combine_layouts_with_dedup(left: &SlotLayout, right: &SlotLayout) -> SlotLayout {
147    let mut slots = left.slots.clone();
148    let offset = slots.len();
149    let left_names: std::collections::HashSet<&str> =
150        left.slots.iter().map(|s| s.name.as_str()).collect();
151
152    for info in &right.slots {
153        let name = if left_names.contains(info.name.as_str()) {
154            format!("{}_{}", info.name, "right")
155        } else {
156            info.name.clone()
157        };
158        slots.push(SlotInfo {
159            slot_id: info.slot_id + offset,
160            name,
161            alias: info.alias.clone(),
162            data_type: info.data_type.clone(),
163            nullable: info.nullable,
164            origin: info.origin.clone(),
165        });
166    }
167    SlotLayout::new(slots)
168}