Skip to main content

graphrecords_query/operations/string_operations/
pad_start.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::PadStart,
16};
17use graphrecords_core::GraphRecord;
18use std::iter::repeat_n;
19
20pub(super) fn pad(
21    label: &'static str,
22    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    let mut padded = String::new();
37    padded
38        .try_reserve(capacity)
39        .map_err(|error| Failure::new(label, External::new(error)))?;
40    padded.extend(repeat_n(character, padding_length));
41    padded.push_str(&value);
42
43    Ok(padded)
44}
45
46#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
47#[operation(scope = Element)]
48#[explain(label = "PadStart")]
49#[plan(optimizer_hints(empty = if_all))]
50pub struct PadStartOperation<W, C> {
51    #[argument]
52    width: W,
53    #[argument]
54    character: C,
55}
56
57impl<W: Prepare, C: Prepare> Prepare for PadStartOperation<W, C> {
58    type Prepared<'a>
59        = (W::Prepared<'a>, C::Prepared<'a>)
60    where
61        Self: 'a;
62
63    fn prepare<'a>(
64        &'a self,
65        graphrecord: &'a GraphRecord,
66        cache: &'a EvaluationCache<'a>,
67    ) -> QueryResult<Self::Prepared<'a>> {
68        Ok((
69            self.width.prepare(graphrecord, cache)?,
70            self.character.prepare(graphrecord, cache)?,
71        ))
72    }
73}
74
75impl<I, V, W, C> ElementKernel<Indexed<I, V>> for PadStartOperation<W, C>
76where
77    I: IndexDomain,
78    V: StringValue,
79    W: ArgumentSource<Keyed<I>>,
80    W::ValueDomain: IntValue,
81    C: ArgumentSource<Keyed<I>>,
82    C::ValueDomain: StringValue,
83{
84    type Emission = <W::Retention as Retention>::Then<C::Retention>;
85    type OutShape = Indexed<I, V>;
86
87    fn pipeline<'a>(
88        _graphrecord: &'a GraphRecord,
89        prepared: Self::Prepared<'a>,
90    ) -> QueryResult<ElementPipeline<'a, Indexed<I, V>, Self>> {
91        Ok(string_pad_indexed::<_, V, W, C>(prepared, Self::LABEL, pad))
92    }
93
94    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
95        input.with_unknown_distinct()
96    }
97}
98
99impl<V, W, C> ElementKernel<Bare<V>> for PadStartOperation<W, C>
100where
101    V: StringValue + BareValueDomain,
102    W: ArgumentSource<Unaligned>,
103    W::ValueDomain: IntValue,
104    C: ArgumentSource<Unaligned>,
105    C::ValueDomain: StringValue,
106{
107    type Emission = <W::Retention as Retention>::Then<C::Retention>;
108    type OutShape = Bare<V>;
109
110    fn pipeline<'a>(
111        _graphrecord: &'a GraphRecord,
112        prepared: Self::Prepared<'a>,
113    ) -> QueryResult<ElementPipeline<'a, Bare<V>, Self>> {
114        Ok(string_pad_bare::<V, W, C>(prepared, Self::LABEL, pad))
115    }
116
117    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
118        input.with_unknown_distinct()
119    }
120}
121
122impl<O, W, C> PadStart<W, C> for O
123where
124    PadStartOperation<W, C>: Operation,
125    O: Apply<PadStartOperation<W, C>>,
126{
127    type ReturnOperand = O::Output;
128
129    fn pad_start(&self, width: W, character: C) -> Self::ReturnOperand {
130        Self::ReturnOperand::new(OperationContext::new(
131            self.clone(),
132            PadStartOperation { width, character },
133        ))
134    }
135}
136
137operation_manifest! {
138    PadStartOperation<W, C> {
139        method: PadStart<W, C>::pad_start;
140        scope: element;
141
142        kernel {
143            parameters: <I: IndexDomain, V: StringValue>;
144            argument: W: ArgumentSource<Keyed<I>> where W::ValueDomain: IntValue;
145            argument: C: ArgumentSource<Keyed<I>> where C::ValueDomain: StringValue;
146            input: Indexed<I, V>;
147            output: Indexed<I, V>;
148            emission: ArgumentRetention;
149        }
150        kernel {
151            parameters: <V: StringValue + BareValueDomain>;
152            argument: W: ArgumentSource<Unaligned> where W::ValueDomain: IntValue;
153            argument: C: ArgumentSource<Unaligned> where C::ValueDomain: StringValue;
154            input: Bare<V>;
155            output: Bare<V>;
156            emission: ArgumentRetention;
157        }
158    }
159}