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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
21pub enum QuadStorageEncoding {
22 PlainTerm,
27 ObjectId(ObjectIdEncodingRef),
29 String,
34}
35
36impl QuadStorageEncoding {
37 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
49pub enum QuadStorageEncodingName {
50 PlainTerm,
52 ObjectId,
54 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 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 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 pub fn object_id_encoding(&self) -> Option<&ObjectIdEncodingRef> {
136 match &self {
137 QuadStorageEncoding::ObjectId(encoding) => Some(encoding),
138 _ => None,
139 }
140 }
141
142 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 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
187fn 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}