1use 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#[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#[derive(Debug, Clone)]
46pub struct CompiledQuery {
47 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 pub notes: Vec<String>,
56}
57
58impl CompiledQuery {
59 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
114pub 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
211fn 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 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
281pub 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 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}