Skip to main content

rdf_fusion_encoding/string/
encoding.rs

1use crate::encoding::TermEncoding;
2use crate::string::{StringEncodingScalar, StringTermArray};
3use crate::{EncodingName, TermEncoder};
4use datafusion::arrow::array::{Array, ArrayRef, StringBuilder};
5use datafusion::arrow::datatypes::DataType;
6use datafusion::common::{ScalarValue, exec_err};
7use rdf_fusion_common::{DFResult, TermRef, ThinResult};
8use std::sync::{Arc, LazyLock};
9
10/// The instance of the [StringEncoding].
11pub static STRING_ENCODING: LazyLock<StringEncodingRef> =
12    LazyLock::new(|| Arc::new(StringEncoding));
13
14/// A cheaply cloneable reference to a [StringEncoding].
15pub type StringEncodingRef = Arc<StringEncoding>;
16
17#[derive(Debug)]
18pub struct StringEncoding;
19
20impl StringEncoding {
21    /// Returns the type of the [StringEncoding].
22    pub fn data_type() -> DataType {
23        DataType::Utf8
24    }
25
26    /// Encodes the `term` as a [`StringEncodingScalar`].
27    pub fn encode_term(
28        &self,
29        term: ThinResult<TermRef<'_>>,
30    ) -> DFResult<StringEncodingScalar> {
31        let value = match term {
32            Ok(term) => Some(term.to_string()),
33            Err(_) => None,
34        };
35        Ok(StringEncodingScalar::new_unchecked(ScalarValue::Utf8(
36            value,
37        )))
38    }
39}
40
41impl TermEncoding for StringEncoding {
42    type Array = StringTermArray;
43    type Scalar = StringEncodingScalar;
44
45    fn name(&self) -> EncodingName {
46        EncodingName::String
47    }
48
49    fn data_type(&self) -> &DataType {
50        static DATA_TYPE: LazyLock<DataType> = LazyLock::new(StringEncoding::data_type);
51        &DATA_TYPE
52    }
53
54    fn try_new_array(self: &Arc<Self>, array: ArrayRef) -> DFResult<Self::Array> {
55        if array.data_type() != self.data_type() {
56            return exec_err!(
57                "Expected array with StringEncoding (Utf8), got: {}",
58                array.data_type()
59            );
60        }
61        Ok(StringTermArray::new_unchecked(array))
62    }
63
64    fn try_new_scalar(self: &Arc<Self>, scalar: ScalarValue) -> DFResult<Self::Scalar> {
65        if scalar.data_type() != *self.data_type() {
66            return exec_err!(
67                "Expected scalar with StringEncoding (Utf8), got: {}",
68                scalar.data_type()
69            );
70        }
71        Ok(StringEncodingScalar::new_unchecked(scalar))
72    }
73}
74
75impl TermEncoder<StringEncoding> for StringEncoding {
76    type Term<'data> = TermRef<'data>;
77
78    fn encode_terms<'data>(
79        &self,
80        terms: impl IntoIterator<Item = ThinResult<Self::Term<'data>>>,
81    ) -> DFResult<StringTermArray> {
82        let mut builder = StringBuilder::new();
83        for term in terms {
84            match term {
85                Ok(term) => builder.append_value(term.to_string()),
86                Err(_) => builder.append_null(),
87            }
88        }
89        Ok(StringTermArray::new_unchecked(Arc::new(builder.finish())))
90    }
91
92    fn encode_term(
93        &self,
94        term: ThinResult<Self::Term<'_>>,
95    ) -> DFResult<StringEncodingScalar> {
96        self.encode_term(term)
97    }
98}
99
100/// A list of arrays with the [StringEncoding].
101#[derive(Debug, Clone)]
102pub struct StringArgs {
103    arrays: Vec<StringTermArray>,
104}
105
106impl StringArgs {
107    /// Creates a new [StringArgs] from a list of arrays.
108    pub fn new_unchecked(arrays: Vec<StringTermArray>) -> Self {
109        Self { arrays }
110    }
111
112    /// Returns the array at `quad_table`.
113    pub fn get(&self, index: usize) -> &StringTermArray {
114        &self.arrays[index]
115    }
116}