Skip to main content

graphrecords_query/operations/string_operations/
pad_end.rs

1use super::{padding_character, string_pad_bare, string_pad_indexed};
2use crate::{
3    Bare, BareValueDomain, Explain, External, Failure, IndexDomain, Indexed, Labeled, Operand,
4    Position, QueryResult,
5    capabilities::{IntValue, StringValue},
6    element::Retention,
7    error::string::StringPaddingOverflow,
8    execution::EvaluationCache,
9    operations::{
10        Apply, ArgumentSource, ElementKernel, ElementPipeline, Keyed, Operation, OperationContext,
11        Prepare, Unaligned,
12    },
13    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
14    registry::{describe::ArgumentRetention, operation_manifest},
15    traits::PadEnd,
16};
17use graphrecords_core::GraphRecord;
18use std::iter::repeat_n;
19
20pub(super) fn pad(
21    label: &'static str,
22    mut value: String,
23    width: Position,
24    character: String,
25) -> QueryResult<String> {
26    let character = padding_character(label, character)?;
27    let padding_length = width.saturating_sub(value.chars().count());
28    if padding_length == 0 {
29        return Ok(value);
30    }
31
32    let capacity = padding_length
33        .checked_mul(character.len_utf8())
34        .and_then(|padding_bytes| padding_bytes.checked_add(value.len()))
35        .ok_or_else(|| Failure::new(label, StringPaddingOverflow::new(width)))?;
36    value
37        .try_reserve(capacity - value.len())
38        .map_err(|error| Failure::new(label, External::new(error)))?;
39    value.extend(repeat_n(character, padding_length));
40
41    Ok(value)
42}
43
44#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
45#[operation(scope = Element)]
46#[explain(label = "PadEnd")]
47#[plan(optimizer_hints(empty = if_all))]
48pub struct PadEndOperation<W, C> {
49    #[argument]
50    width: W,
51    #[argument]
52    character: C,
53}
54
55impl<W: Prepare, C: Prepare> Prepare for PadEndOperation<W, C> {
56    type Prepared<'a>
57        = (W::Prepared<'a>, C::Prepared<'a>)
58    where
59        Self: 'a;
60
61    fn prepare<'a>(
62        &'a self,
63        graphrecord: &'a GraphRecord,
64        cache: &'a EvaluationCache<'a>,
65    ) -> QueryResult<Self::Prepared<'a>> {
66        Ok((
67            self.width.prepare(graphrecord, cache)?,
68            self.character.prepare(graphrecord, cache)?,
69        ))
70    }
71}
72
73impl<I, V, W, C> ElementKernel<Indexed<I, V>> for PadEndOperation<W, C>
74where
75    I: IndexDomain,
76    V: StringValue,
77    W: ArgumentSource<Keyed<I>>,
78    W::ValueDomain: IntValue,
79    C: ArgumentSource<Keyed<I>>,
80    C::ValueDomain: StringValue,
81{
82    type Emission = <W::Retention as Retention>::Then<C::Retention>;
83    type OutShape = Indexed<I, V>;
84
85    fn pipeline<'a>(
86        _graphrecord: &'a GraphRecord,
87        prepared: Self::Prepared<'a>,
88    ) -> QueryResult<ElementPipeline<'a, Indexed<I, V>, Self>> {
89        Ok(string_pad_indexed::<_, V, W, C>(prepared, Self::LABEL, pad))
90    }
91
92    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
93        input.with_unknown_distinct()
94    }
95}
96
97impl<V, W, C> ElementKernel<Bare<V>> for PadEndOperation<W, C>
98where
99    V: StringValue + BareValueDomain,
100    W: ArgumentSource<Unaligned>,
101    W::ValueDomain: IntValue,
102    C: ArgumentSource<Unaligned>,
103    C::ValueDomain: StringValue,
104{
105    type Emission = <W::Retention as Retention>::Then<C::Retention>;
106    type OutShape = Bare<V>;
107
108    fn pipeline<'a>(
109        _graphrecord: &'a GraphRecord,
110        prepared: Self::Prepared<'a>,
111    ) -> QueryResult<ElementPipeline<'a, Bare<V>, Self>> {
112        Ok(string_pad_bare::<V, W, C>(prepared, Self::LABEL, pad))
113    }
114
115    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
116        input.with_unknown_distinct()
117    }
118}
119
120impl<O, W, C> PadEnd<W, C> for O
121where
122    PadEndOperation<W, C>: Operation,
123    O: Apply<PadEndOperation<W, C>>,
124{
125    type ReturnOperand = O::Output;
126
127    fn pad_end(&self, width: W, character: C) -> Self::ReturnOperand {
128        Self::ReturnOperand::new(OperationContext::new(
129            self.clone(),
130            PadEndOperation { width, character },
131        ))
132    }
133}
134
135operation_manifest! {
136    PadEndOperation<W, C> {
137        method: PadEnd<W, C>::pad_end;
138        scope: element;
139
140        kernel {
141            parameters: <I: IndexDomain, V: StringValue>;
142            argument: W: ArgumentSource<Keyed<I>> where W::ValueDomain: IntValue;
143            argument: C: ArgumentSource<Keyed<I>> where C::ValueDomain: StringValue;
144            input: Indexed<I, V>;
145            output: Indexed<I, V>;
146            emission: ArgumentRetention;
147        }
148        kernel {
149            parameters: <V: StringValue + BareValueDomain>;
150            argument: W: ArgumentSource<Unaligned> where W::ValueDomain: IntValue;
151            argument: C: ArgumentSource<Unaligned> where C::ValueDomain: StringValue;
152            input: Bare<V>;
153            output: Bare<V>;
154            emission: ArgumentRetention;
155        }
156    }
157}