harmos_storage/query/transform/
project.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Storage Arrow Query Contract
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16pub use harmos_domain::model::storage::query::transform::project::{Project, Projection};
17
18/* === Imports === */
19
20use crate::storage::query::transform::FieldPathExecution;
21// Standard Library Imports
22use std::sync::Arc;
23// Third Party Imports
24use arrow_array::RecordBatch;
25use arrow_schema::Schema;
26// Local Imports
27use super::{FieldPath, invalid};
28use crate::Result;
29
30/* === Constants === */
31
32/* === Macros === */
33
34/* === Type Aliases === */
35
36/* === Structs === */
37
38/* === Enums === */
39
40/* === Traits === */
41
42pub trait ProjectExecution {
43    fn output_schema(&self, input: &Schema) -> Result<Schema>;
44
45    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
46}
47
48/* === Implementations === */
49
50impl ProjectExecution for Project {
51    fn output_schema(&self, input: &Schema) -> Result<Schema> {
52        if self.fields.is_empty() || self.fields.len() > 4096 {
53            return Err(invalid("projection requires bounded nonempty fields"));
54        }
55        let mut fields = Vec::with_capacity(self.fields.len());
56        for projected in &self.fields {
57            if projected.name.is_empty() || self.fields.iter().filter(|field| field.name == projected.name).count() != 1 {
58                return Err(invalid("projection names must be nonempty and unique"));
59            }
60            let field = projected.field.field(input)?.with_name(&projected.name);
61            fields.push(field);
62        }
63        Ok(Schema::new_with_metadata(fields, input.metadata().clone()))
64    }
65
66    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
67        let schema = Arc::new(self.output_schema(input.schema().as_ref())?);
68        let mut columns = Vec::with_capacity(self.fields.len());
69        for projected in &self.fields {
70            check()?;
71            // A nested projection must propagate parent Struct validity.
72            let mut array = Arc::clone(projected.field.array(input)?);
73            for depth in 1..projected.field.fields.len() {
74                let parent = FieldPath { fields: projected.field.fields[..depth].to_vec() };
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
