harmos_storage/query/transform/
project.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14pub use harmos_domain::model::storage::query::transform::project::{Project, Projection};
15
16/* === Imports === */
17
18use crate::storage::query::transform::FieldPathExecution;
19// Standard Library Imports
20use std::sync::Arc;
21// Third Party Imports
22use arrow_array::RecordBatch;
23use arrow_schema::Schema;
24// Local Imports
25use super::{FieldPath, invalid};
26use crate::Result;
27
28/* === Constants === */
29
30/* === Macros === */
31
32/* === Type Aliases === */
33
34/* === Structs === */
35
36/* === Enums === */
37
38/* === Traits === */
39
40pub trait ProjectExecution {
41    fn output_schema(&self, input: &Schema) -> Result<Schema>;
42
43    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
44}
45
46/* === Implementations === */
47
48impl ProjectExecution for Project {
49    fn output_schema(&self, input: &Schema) -> Result<Schema> {
50        if self.fields.is_empty() || self.fields.len() > 4096 {
51            return Err(invalid("projection requires bounded nonempty fields"));
52        }
53        let mut fields = Vec::with_capacity(self.fields.len());
54        for projected in &self.fields {
55            if projected.name.is_empty() || self.fields.iter().filter(|field| field.name == projected.name).count() != 1 {
56                return Err(invalid("projection names must be nonempty and unique"));
57            }
58            let field = projected.field.field(input)?.with_name(&projected.name);
59            fields.push(field);
60        }
61        Ok(Schema::new_with_metadata(fields, input.metadata().clone()))
62    }
63
64    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
65        let schema = Arc::new(self.output_schema(input.schema().as_ref())?);
66        let mut columns = Vec::with_capacity(self.fields.len());
67        for projected in &self.fields {
68            check()?;
69            // A nested projection must propagate parent Struct validity.
70            let mut array = Arc::clone(projected.field.array(input)?);
71            for depth in 1..projected.field.fields.len() {
72                let parent = FieldPath {
73                    fields: projected.field.fields[..depth].to_vec(),
74                };
75                let parent = parent.array(input)?;
76                if let Some(nulls) = parent.nulls() {
77                    let combined = arrow_buffer::NullBuffer::union(array.nulls(), Some(nulls));
78                    let data = array
79                        .to_data()
80                        .into_builder()
81                        .nulls(combined)
82                        .build()
83                        .map_err(|_| invalid("projected parent validity is invalid"))?;
84                    array = arrow_array::make_array(data);
85                }
86            }
87            columns.push(array);
88        }
89        RecordBatch::try_new(schema, columns).map_err(|_| invalid("projected batch is invalid"))
90    }
91}
92
93/* === Public Functions === */
94
95/* === Private Functions === */
96
97/* === Tests === */
harmos_storage/query/transform/project.rs

harmos_storage/query/transform/
project.rs
