Skip to main content

oxilite_core/
query.rs

1//! Query compilation entry point, execution job and result decoding.
2//!
3// @lat: [[architecture#SPARQL to SQL compiler]]
4
5use crate::compiler::expr::format_number;
6use crate::compiler::{Block, Col, Compiler, QueryOptions, VAL_FIELDS};
7use crate::encoding::{blank_node_id, tag_of, Tag, DEFAULT_GRAPH_ID};
8use crate::error::{Error, Result};
9use crate::job::{Job, Step};
10use crate::resolve::{join, TermResolver};
11use crate::sql::{Capabilities, Request, Response, SqlValue, Statement};
12use crate::stats::Stats;
13use oxrdf::{BlankNode, Literal, NamedNode, NamedOrBlankNode, Term, Triple, Variable};
14use spargebra::algebra::GraphPattern;
15use spargebra::term::{NamedNodePattern, TermPattern, TriplePattern};
16use spargebra::Query;
17use std::collections::{HashMap, HashSet};
18
19/// Decoded query results.
20#[derive(Debug, Clone, PartialEq)]
21pub enum QueryOutput {
22    Solutions {
23        variables: Vec<Variable>,
24        rows: Vec<Vec<Option<Term>>>,
25    },
26    Boolean(bool),
27    Graph(Vec<Triple>),
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31enum Out {
32    Id,
33    Val,
34}
35
36#[derive(Debug, Clone)]
37enum Form {
38    Select,
39    Ask,
40    Construct(Vec<TriplePattern>),
41    Describe,
42}
43
44/// A query compiled to SQL.
45#[derive(Debug, Clone)]
46pub struct CompiledQuery {
47    /// The main SQL statement.
48    pub sql: String,
49    form: Form,
50    variables: Vec<Variable>,
51    layout: Vec<Out>,
52    constants: HashMap<i64, Term>,
53    union_default_graph: bool,
54    /// Planner decisions and warnings.
55    pub notes: Vec<String>,
56}
57
58impl CompiledQuery {
59    /// Human-readable description for `explain()`: the SQL, join orders with estimated
60    /// cardinalities, and warnings.
61    pub fn explain(&self) -> String {
62        let mut out = format!(
63            "-- oxilite: fully compiled to SQL ({} output variable(s))\n",
64            self.variables.len()
65        );
66        for n in &self.notes {
67            out.push_str("-- ");
68            out.push_str(n);
69            out.push('\n');
70        }
71        out.push_str(&self.sql);
72        out
73    }
74}
75
76fn projected(p: &GraphPattern) -> Option<&[Variable]> {
77    match p {
78        GraphPattern::Project { variables, .. } => Some(variables),
79        GraphPattern::Slice { inner, .. }
80        | GraphPattern::Distinct { inner }
81        | GraphPattern::Reduced { inner }
82        | GraphPattern::OrderBy { inner, .. } => projected(inner),
83        _ => None,
84    }
85}
86
87fn template_vars(template: &[TriplePattern], out: &mut Vec<Variable>) {
88    fn tp(t: &TermPattern, out: &mut Vec<Variable>) {
89        match t {
90            TermPattern::Variable(v) if !out.contains(v) => out.push(v.clone()),
91            TermPattern::Triple(t) => {
92                tp(&t.subject, out);
93                if let NamedNodePattern::Variable(v) = &t.predicate {
94                    if !out.contains(v) {
95                        out.push(v.clone());
96                    }
97                }
98                tp(&t.object, out);
99            }
100            _ => {}
101        }
102    }
103    for t in template {
104        tp(&t.subject, out);
105        if let NamedNodePattern::Variable(v) = &t.predicate {
106            if !out.contains(v) {
107                out.push(v.clone());
108            }
109        }
110        tp(&t.object, out);
111    }
112}
113
114/// Compiles a SPARQL query to SQL. Returns [`Error::Unsupported`] for anything that cannot
115/// be expressed in SQL on this backend.
116pub fn compile_query(
117    query: &Query,
118    stats: &Stats,
119    caps: &Capabilities,
120    options: &QueryOptions,
121) -> Result<CompiledQuery> {
122    let (dataset, pattern, base, form) = match query {
123        Query::Select {
124            dataset,
125            pattern,
126            base_iri,
127        } => (dataset, pattern, base_iri, Form::Select),
128        Query::Ask {
129            dataset,
130            pattern,
131            base_iri,
132        } => (dataset, pattern, base_iri, Form::Ask),
133        Query::Construct {
134            template,
135            dataset,
136            pattern,
137            base_iri,
138        } => (
139            dataset,
140            pattern,
141            base_iri,
142            Form::Construct(template.clone()),
143        ),
144        Query::Describe {
145            dataset,
146            pattern,
147            base_iri,
148        } => (dataset, pattern, base_iri, Form::Describe),
149    };
150    let mut c = Compiler::new(
151        stats,
152        caps,
153        options,
154        dataset.as_ref(),
155        base.as_ref().map(|b| b.as_str().to_string()),
156    );
157    let block = c.pattern(pattern)?;
158    let variables: Vec<Variable> = match &form {
159        Form::Construct(t) => {
160            let mut v = Vec::new();
161            template_vars(t, &mut v);
162            v
163        }
164        _ => match projected(pattern) {
165            Some(v) => v.to_vec(),
166            None => block
167                .cols
168                .keys()
169                .map(|i| c.var_names[*i].clone())
170                .filter(|v| !v.as_str().starts_with('\u{1}'))
171                .collect(),
172        },
173    };
174    let idxs: Vec<usize> = variables.iter().map(|v| c.var(v)).collect();
175    let layout = idxs
176        .iter()
177        .map(|i| match block.cols.get(i).map(|b| &b.col) {
178            Some(Col::Val(_)) => Out::Val,
179            _ => Out::Id,
180        })
181        .collect();
182    let sql = match form {
183        Form::Ask => {
184            let inner = block.to_select(Some(&[]), false);
185            format!("SELECT EXISTS ({inner})")
186        }
187        _ => final_select(&block, &idxs, caps.int64_as_text),
188    };
189    if sql.len() > caps.max_sql_len {
190        return Err(Error::unsupported(format!(
191            "generated SQL is {} bytes, over the backend limit of {}",
192            sql.len(),
193            caps.max_sql_len
194        )));
195    }
196    Ok(CompiledQuery {
197        sql,
198        form,
199        variables,
200        layout,
201        constants: c.constants,
202        union_default_graph: options.union_default_graph,
203        notes: c.notes,
204    })
205}
206
207fn final_select(block: &Block, idxs: &[usize], text_ids: bool) -> String {
208    block.to_select(Some(idxs), text_ids)
209}
210
211/// Decodes a computed value from its value columns.
212fn decode_val(cells: &[SqlValue]) -> Option<Term> {
213    let get = |name: &str| {
214        let i = VAL_FIELDS.iter().position(|f| *f == name).expect("field");
215        &cells[i]
216    };
217    let kind = get("k").as_i64()?;
218    let lex = get("l").clone().into_string();
219    let dt = get("d").clone().into_string();
220    let lang = get("g").clone().into_string();
221    let num = get("n").as_f64();
222    Some(match Tag::from_u8(kind as u8)? {
223        Tag::Iri => NamedNode::new(lex?).ok()?.into(),
224        Tag::BlankNode => BlankNode::new_unchecked(lex?).into(),
225        Tag::String => Literal::new_simple_literal(lex?).into(),
226        Tag::LangString => Literal::new_language_tagged_literal(lex?, lang?)
227            .ok()?
228            .into(),
229        Tag::DirLangString => {
230            let lang = lang?;
231            let (tag, dir) = lang.split_once("--")?;
232            let dir = if dir == "rtl" {
233                oxrdf::BaseDirection::Rtl
234            } else {
235                oxrdf::BaseDirection::Ltr
236            };
237            Literal::new_directional_language_tagged_literal(lex?, tag, dir)
238                .ok()?
239                .into()
240        }
241        Tag::Integer => match lex {
242            Some(l) => Literal::new_typed_literal(l, oxrdf::vocab::xsd::INTEGER).into(),
243            None => Literal::from(num? as i64).into(),
244        },
245        Tag::Boolean => Literal::from(get("b").as_i64()? != 0).into(),
246        Tag::Typed => {
247            let dt = dt?;
248            if let Some((sum, count)) = get("x")
249                .as_str()
250                .and_then(|a| a.strip_prefix("avg:"))
251                .and_then(|a| a.split_once('/'))
252            {
253                // Exact integer AVG: decimal division like Oxigraph.
254                let sum = oxsdatatypes::Decimal::from(sum.parse::<i64>().ok()?);
255                let count = oxsdatatypes::Decimal::from(count.parse::<i64>().ok()?);
256                return Some(Literal::from(sum.checked_div(count)?).into());
257            }
258            match lex {
259                Some(l) => Literal::new_typed_literal(l, NamedNode::new_unchecked(dt)).into(),
260                None => format_number(num?, &dt).into(),
261            }
262        }
263        Tag::Triple | Tag::Default => return None,
264    })
265}
266
267enum Cell {
268    Id(i64),
269    Term(Term),
270    Unbound,
271}
272
273enum State {
274    Start,
275    Main,
276    Resolving,
277    Describe,
278    DescribeResolving,
279}
280
281/// Runs a compiled query.
282pub struct QueryJob {
283    compiled: CompiledQuery,
284    caps: Capabilities,
285    state: State,
286    resolver: TermResolver,
287    cells: Vec<Vec<Cell>>,
288    describe_done: HashSet<i64>,
289    describe_quads: Vec<[i64; 3]>,
290}
291
292impl QueryJob {
293    pub fn new(compiled: CompiledQuery, caps: Capabilities) -> Self {
294        let resolver = TermResolver::with_constants(compiled.constants.clone());
295        Self {
296            compiled,
297            caps,
298            state: State::Start,
299            resolver,
300            cells: Vec::new(),
301            describe_done: HashSet::new(),
302            describe_quads: Vec::new(),
303        }
304    }
305
306    pub fn compiled(&self) -> &CompiledQuery {
307        &self.compiled
308    }
309
310    fn absorb_main(&mut self, response: Response) -> Result<()> {
311        let rs = response
312            .into_iter()
313            .next()
314            .ok_or_else(|| Error::backend("empty response"))?;
315        for row in rs.rows {
316            let mut out = Vec::with_capacity(self.compiled.layout.len());
317            let mut i = 0;
318            for l in &self.compiled.layout {
319                match l {
320                    Out::Id => {
321                        let cell = match row.get(i).and_then(SqlValue::as_i64) {
322                            Some(id) => {
323                                self.resolver.want(id);
324                                Cell::Id(id)
325                            }
326                            None => Cell::Unbound,
327                        };
328                        out.push(cell);
329                        i += 1;
330                    }
331                    Out::Val => {
332                        let cells = row
333                            .get(i..i + VAL_FIELDS.len())
334                            .ok_or_else(|| Error::corrupted("short row"))?;
335                        out.push(decode_val(cells).map_or(Cell::Unbound, Cell::Term));
336                        i += VAL_FIELDS.len();
337                    }
338                }
339            }
340            self.cells.push(out);
341        }
342        Ok(())
343    }
344
345    fn solutions(&self) -> Result<Vec<Vec<Option<Term>>>> {
346        self.cells
347            .iter()
348            .map(|row| {
349                row.iter()
350                    .map(|c| match c {
351                        Cell::Id(id) => self.resolver.get(*id).map(Some),
352                        Cell::Term(t) => Ok(Some(t.clone())),
353                        Cell::Unbound => Ok(None),
354                    })
355                    .collect()
356            })
357            .collect()
358    }
359
360    fn describe_request(&mut self, ids: Vec<i64>) -> Option<Request> {
361        let ids: Vec<i64> = ids
362            .into_iter()
363            .filter(|id| {
364                matches!(tag_of(*id), Some(Tag::Iri | Tag::BlankNode))
365                    && self.describe_done.insert(*id)
366            })
367            .collect();
368        if ids.is_empty() {
369            return None;
370        }
371        let c = |x: &str| {
372            if self.caps.int64_as_text {
373                format!("CAST({x} AS TEXT)")
374            } else {
375                x.into()
376            }
377        };
378        let graph = if self.compiled.union_default_graph {
379            String::new()
380        } else {
381            format!(" AND g = {DEFAULT_GRAPH_ID}")
382        };
383        let stmts = ids
384            .chunks(400)
385            .map(|chunk| {
386                Statement::new(format!(
387                    "SELECT DISTINCT {}, {}, {} FROM quads WHERE s IN ({}){graph}",
388                    c("s"),
389                    c("p"),
390                    c("o"),
391                    join(chunk)
392                ))
393            })
394            .collect();
395        Some(Request::read(stmts))
396    }
397
398    fn finish(&mut self) -> Result<QueryOutput> {
399        let rows = self.solutions()?;
400        Ok(match &self.compiled.form {
401            Form::Select => QueryOutput::Solutions {
402                variables: self.compiled.variables.clone(),
403                rows,
404            },
405            Form::Ask => unreachable!("handled on the main response"),
406            Form::Construct(template) => {
407                let mut seen = HashSet::new();
408                let mut out = Vec::new();
409                for row in rows {
410                    let mut bnodes = HashMap::new();
411                    for t in template {
412                        if let Some(triple) =
413                            instantiate(t, &self.compiled.variables, &row, &mut bnodes)
414                        {
415                            if seen.insert(triple.clone()) {
416                                out.push(triple);
417                            }
418                        }
419                    }
420                }
421                QueryOutput::Graph(out)
422            }
423            Form::Describe => {
424                let mut out = Vec::new();
425                let mut seen = HashSet::new();
426                for [s, p, o] in &self.describe_quads {
427                    let (Ok(s), Ok(p), Ok(o)) = (
428                        self.resolver.get(*s),
429                        self.resolver.get(*p),
430                        self.resolver.get(*o),
431                    ) else {
432                        continue;
433                    };
434                    let (Ok(s), Term::NamedNode(p)) = (crate::encoding::to_subject(s), p) else {
435                        continue;
436                    };
437                    let t = Triple::new(s, p, o);
438                    if seen.insert(t.clone()) {
439                        out.push(t);
440                    }
441                }
442                QueryOutput::Graph(out)
443            }
444        })
445    }
446
447    fn after_resolution(&mut self) -> Result<Step<QueryOutput>> {
448        if let Some(r) = self.resolver.request(&self.caps) {
449            return Ok(Step::Execute(r));
450        }
451        if let (Form::Describe, State::Resolving) = (&self.compiled.form, &self.state) {
452            let ids: Vec<i64> = self
453                .cells
454                .iter()
455                .flatten()
456                .filter_map(|c| match c {
457                    Cell::Id(id) => Some(*id),
458                    Cell::Term(t) => match t {
459                        Term::NamedNode(n) => Some(crate::encoding::named_node_id(n.as_str())),
460                        Term::BlankNode(b) => Some(blank_node_id(b.as_str())),
461                        _ => None,
462                    },
463                    Cell::Unbound => None,
464                })
465                .collect();
466            self.state = State::Describe;
467            if let Some(r) = self.describe_request(ids) {
468                return Ok(Step::Execute(r));
469            }
470        }
471        Ok(Step::Done(self.finish()?))
472    }
473}
474
475fn instantiate(
476    t: &TriplePattern,
477    vars: &[Variable],
478    row: &[Option<Term>],
479    bnodes: &mut HashMap<String, BlankNode>,
480) -> Option<Triple> {
481    fn term(
482        t: &TermPattern,
483        vars: &[Variable],
484        row: &[Option<Term>],
485        bnodes: &mut HashMap<String, BlankNode>,
486    ) -> Option<Term> {
487        Some(match t {
488            TermPattern::NamedNode(n) => n.clone().into(),
489            TermPattern::Literal(l) => l.clone().into(),
490            TermPattern::BlankNode(b) => bnodes
491                .entry(b.as_str().to_string())
492                .or_default()
493                .clone()
494                .into(),
495            TermPattern::Variable(v) => row[vars.iter().position(|x| x == v)?].clone()?,
496            TermPattern::Triple(tp) => instantiate(tp, vars, row, bnodes)?.into(),
497        })
498    }
499    let s = match term(&t.subject, vars, row, bnodes)? {
500        Term::NamedNode(n) => NamedOrBlankNode::from(n),
501        Term::BlankNode(b) => NamedOrBlankNode::from(b),
502        _ => return None,
503    };
504    let p = match &t.predicate {
505        NamedNodePattern::NamedNode(n) => n.clone(),
506        NamedNodePattern::Variable(v) => match row[vars.iter().position(|x| x == v)?].clone()? {
507            Term::NamedNode(n) => n,
508            _ => return None,
509        },
510    };
511    let o = term(&t.object, vars, row, bnodes)?;
512    Some(Triple::new(s, p, o))
513}
514
515impl Job for QueryJob {
516    type Output = QueryOutput;
517
518    fn step(&mut self, response: Option<Response>) -> Result<Step<QueryOutput>> {
519        match (&self.state, response) {
520            (State::Start, _) => {
521                self.state = State::Main;
522                Ok(Step::Execute(Request::read(vec![Statement::new(
523                    self.compiled.sql.clone(),
524                )])))
525            }
526            (State::Main, Some(r)) => {
527                if let Form::Ask = self.compiled.form {
528                    let b = r
529                        .first()
530                        .and_then(|rs| rs.rows.first())
531                        .and_then(|row| row.first())
532                        .and_then(SqlValue::as_i64)
533                        .unwrap_or(0);
534                    return Ok(Step::Done(QueryOutput::Boolean(b != 0)));
535                }
536                self.absorb_main(r)?;
537                self.state = State::Resolving;
538                self.after_resolution()
539            }
540            (State::Resolving, Some(r)) => {
541                self.resolver.absorb(r)?;
542                self.after_resolution()
543            }
544            (State::Describe, Some(r)) => {
545                let mut next = Vec::new();
546                for rs in r {
547                    for row in rs.rows {
548                        let ids: Vec<i64> = row.iter().filter_map(SqlValue::as_i64).collect();
549                        if let [s, p, o] = ids[..] {
550                            self.resolver.want(s);
551                            self.resolver.want(p);
552                            self.resolver.want(o);
553                            self.describe_quads.push([s, p, o]);
554                            if tag_of(o) == Some(Tag::BlankNode) {
555                                next.push(o);
556                            }
557                        }
558                    }
559                }
560                // Concise bounded description: follow blank-node objects.
561                if let Some(r) = self.describe_request(next) {
562                    return Ok(Step::Execute(r));
563                }
564                self.state = State::DescribeResolving;
565                match self.resolver.request(&self.caps) {
566                    Some(r) => Ok(Step::Execute(r)),
567                    None => Ok(Step::Done(self.finish()?)),
568                }
569            }
570            (State::DescribeResolving, Some(r)) => {
571                self.resolver.absorb(r)?;
572                match self.resolver.request(&self.caps) {
573                    Some(r) => Ok(Step::Execute(r)),
574                    None => Ok(Step::Done(self.finish()?)),
575                }
576            }
577            (_, None) => Err(Error::Other("query job resumed without a response".into())),
578        }
579    }
580}