Skip to main content

rdf_fusion_encoding/
quad_storage_encoding.rs

1use crate::TermEncoding;
2use crate::encoding::EncodingScalar;
3use crate::object_id::{ObjectIdEncoding, ObjectIdEncodingRef};
4use crate::plain_term::{PLAIN_TERM_ENCODING, PlainTermEncoding};
5use crate::string::STRING_ENCODING;
6use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
7use datafusion::common::{DFSchema, DFSchemaRef, ScalarValue};
8use rdf_fusion_common::quads::{COL_GRAPH, COL_OBJECT, COL_PREDICATE, COL_SUBJECT};
9use rdf_fusion_common::{DFResult, TermRef};
10use std::collections::HashMap;
11use std::fmt::{Display, Formatter};
12use std::str::FromStr;
13use std::sync::{Arc, LazyLock};
14use thiserror::Error;
15
16/// Defines which encoding is used for retrieving quads from the storage.
17///
18/// Defining this is necessary such that the query planner knows what type should be assigned to the
19/// schema of quad pattern logical nodes.
20#[derive(Debug, Clone, PartialEq, Eq, Hash)]
21pub enum QuadStorageEncoding {
22    /// Uses the plain term encoding.
23    ///
24    /// Currently, the plain term encoding is not parameterizable. Therefore, this variant has no
25    /// further information.
26    PlainTerm,
27    /// Uses the provided object id encoding.
28    ObjectId(ObjectIdEncodingRef),
29    /// Uses the string encoding.
30    ///
31    /// Currently, the string encoding is not parameterizable. Therefore, this variant has no
32    /// further information.
33    String,
34}
35
36impl QuadStorageEncoding {
37    /// Returns the name of the encoding.
38    pub fn name(&self) -> QuadStorageEncodingName {
39        match self {
40            QuadStorageEncoding::PlainTerm => QuadStorageEncodingName::PlainTerm,
41            QuadStorageEncoding::ObjectId(_) => QuadStorageEncodingName::ObjectId,
42            QuadStorageEncoding::String => QuadStorageEncodingName::String,
43        }
44    }
45}
46
47/// A version of [`QuadStorageEncoding`] that only reflects the name of the encoding.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
49pub enum QuadStorageEncodingName {
50    /// See [`QuadStorageEncoding::PlainTerm`]
51    PlainTerm,
52    /// See [`QuadStorageEncoding::ObjectId`]
53    ObjectId,
54    /// See [`QuadStorageEncoding::String`]
55    String,
56}
57
58impl Display for QuadStorageEncodingName {
59    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
60        match self {
61            QuadStorageEncodingName::PlainTerm => f.write_str("PlainTerm"),
62            QuadStorageEncodingName::ObjectId => f.write_str("ObjectId"),
63            QuadStorageEncodingName::String => f.write_str("String"),
64        }
65    }
66}
67
68#[derive(Debug, Error)]
69#[error(
70    "Invalid quad storage encoding: {0}. Supported encodings: plain-term, string, object-id."
71)]
72pub struct QuadStorageEncodingNameParserError(String);
73
74impl FromStr for QuadStorageEncodingName {
75    type Err = QuadStorageEncodingNameParserError;
76
77    fn from_str(s: &str) -> Result<Self, Self::Err> {
78        match s.to_lowercase().as_str() {
79            "plain-term" | "plainterm" => Ok(QuadStorageEncodingName::PlainTerm),
80            "object-id" | "objectid" => Ok(QuadStorageEncodingName::ObjectId),
81            "string" => Ok(QuadStorageEncodingName::String),
82            _ => Err(QuadStorageEncodingNameParserError(s.to_string())),
83        }
84    }
85}
86
87static PLAIN_TERM_QUAD_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
88    SchemaRef::new(Schema::new(vec![
89        Field::new(COL_GRAPH, PlainTermEncoding::data_type(), true),
90        Field::new(COL_SUBJECT, PlainTermEncoding::data_type(), false),
91        Field::new(COL_PREDICATE, PlainTermEncoding::data_type(), false),
92        Field::new(COL_OBJECT, PlainTermEncoding::data_type(), false),
93    ]))
94});
95
96static PLAIN_TERM_QUAD_DFSCHEMA: LazyLock<DFSchemaRef> = LazyLock::new(|| {
97    DFSchemaRef::new(DFSchema::try_from(PLAIN_TERM_QUAD_SCHEMA.clone()).unwrap())
98});
99
100static STRING_QUAD_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
101    SchemaRef::new(Schema::new(vec![
102        Field::new(COL_GRAPH, STRING_ENCODING.data_type().clone(), true),
103        Field::new(COL_SUBJECT, STRING_ENCODING.data_type().clone(), false),
104        Field::new(COL_PREDICATE, STRING_ENCODING.data_type().clone(), false),
105        Field::new(COL_OBJECT, STRING_ENCODING.data_type().clone(), false),
106    ]))
107});
108
109static STRING_QUAD_DFSCHEMA: LazyLock<DFSchemaRef> = LazyLock::new(|| {
110    DFSchemaRef::new(DFSchema::try_from(STRING_QUAD_SCHEMA.clone()).unwrap())
111});
112
113impl QuadStorageEncoding {
114    /// Returns the data type of a single term column, given the current encoding.
115    pub fn term_type(&self) -> &DataType {
116        match self {
117            QuadStorageEncoding::PlainTerm => PLAIN_TERM_ENCODING.data_type(),
118            QuadStorageEncoding::ObjectId(enc) => enc.data_type(),
119            QuadStorageEncoding::String => STRING_ENCODING.data_type(),
120        }
121    }
122
123    /// Returns the schema of an entire quad, given the current encoding.
124    pub fn quad_schema(&self) -> DFSchemaRef {
125        match self {
126            QuadStorageEncoding::PlainTerm => PLAIN_TERM_QUAD_DFSCHEMA.clone(),
127            QuadStorageEncoding::ObjectId(encoding) => object_id_quad_schema(encoding),
128            QuadStorageEncoding::String => STRING_QUAD_DFSCHEMA.clone(),
129        }
130    }
131
132    /// Returns an optional reference to the contained [ObjectIdEncoding].
133    ///
134    /// Returns [None] otherwise.
135    pub fn object_id_encoding(&self) -> Option<&ObjectIdEncodingRef> {
136        match &self {
137            QuadStorageEncoding::ObjectId(encoding) => Some(encoding),
138            _ => None,
139        }
140    }
141
142    /// Encodes the given term into a [ScalarValue] that can be used for filtering.
143    ///
144    /// If `Ok(None)` is returned, the mapping for the object id does not exist.
145    pub async fn try_encode_term(
146        &self,
147        term: TermRef<'_>,
148    ) -> DFResult<Option<ScalarValue>> {
149        match self {
150            QuadStorageEncoding::PlainTerm => Ok(Some(
151                PLAIN_TERM_ENCODING
152                    .encode_term(Ok(term))?
153                    .into_scalar_value(),
154            )),
155            QuadStorageEncoding::ObjectId(enc) => {
156                let pt_scalar = PLAIN_TERM_ENCODING.encode_term(Ok(term))?;
157                let mapping = enc.mapping();
158                let val = mapping.try_get_object_id(&pt_scalar).await?;
159                Ok(val)
160            }
161            QuadStorageEncoding::String => Ok(Some(
162                STRING_ENCODING.encode_term(Ok(term))?.into_scalar_value(),
163            )),
164        }
165    }
166
167    /// Creates a null scalar for the given storage encoding.
168    pub fn create_null_scalar(&self) -> DFResult<ScalarValue> {
169        ScalarValue::try_new_null(self.term_type())
170    }
171}
172
173impl Display for QuadStorageEncoding {
174    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
175        match self {
176            QuadStorageEncoding::PlainTerm => write!(f, "PlainTerm"),
177            QuadStorageEncoding::ObjectId(encoding) => {
178                write!(f, "ObjectId({})", encoding.object_id_data_type())
179            }
180            QuadStorageEncoding::String => {
181                write!(f, "String")
182            }
183        }
184    }
185}
186
187/// Computes the quad schema based on the given [ObjectIdEncoding].
188fn object_id_quad_schema(encoding: &ObjectIdEncoding) -> DFSchemaRef {
189    let data_type = encoding.data_type();
190    Arc::new(
191        DFSchema::from_unqualified_fields(
192            Fields::from(vec![
193                Field::new(COL_GRAPH, data_type.clone(), true),
194                Field::new(COL_SUBJECT, data_type.clone(), false),
195                Field::new(COL_PREDICATE, data_type.clone(), false),
196                Field::new(COL_OBJECT, data_type.clone(), false),
197            ]),
198            HashMap::new(),
199        )
200        .expect("Fields are fixed"),
201    )
202}
203
204#[cfg(test)]
205mod test {
206    use crate::QuadStorageEncodingName;
207
208    #[test]
209    fn test_quad_storage_name_parsing_round_trip() {
210        test_parsing_round_trip(QuadStorageEncodingName::PlainTerm);
211        test_parsing_round_trip(QuadStorageEncodingName::ObjectId);
212
213        fn test_parsing_round_trip(name: QuadStorageEncodingName) {
214            let serialized = name.to_string();
215            assert_eq!(serialized.parse::<QuadStorageEncodingName>().unwrap(), name);
216        }
217    }
218}