exon_fcs/
config.rs

1// Copyright 2023 WHERE TRUE Technologies.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use arrow::datatypes::{Field, Schema, SchemaRef};
18use exon_common::DEFAULT_BATCH_SIZE;
19use object_store::ObjectStore;
20
21/// Configuration for a FCS datasource.
22pub struct FCSConfig {
23    /// The number of records to read at a time.
24    pub batch_size: usize,
25    /// The object store to use.
26    pub object_store: Arc<dyn ObjectStore>,
27    /// The file schema to use.
28    pub file_schema: Arc<arrow::datatypes::Schema>,
29    /// Any projections to apply to the resulting batches.
30    pub projection: Option<Vec<usize>>,
31}
32
33impl FCSConfig {
34    /// Create a new FCS configuration.
35    pub fn new(object_store: Arc<dyn ObjectStore>, file_schema: SchemaRef) -> Self {
36        Self {
37            batch_size: DEFAULT_BATCH_SIZE,
38            object_store,
39            file_schema,
40            projection: None,
41        }
42    }
43
44    /// Set the batch size.
45    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
46        self.batch_size = batch_size;
47        self
48    }
49
50    /// Set the projection.
51    pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
52        self.projection = Some(projection);
53        self
54    }
55}
56
57#[derive(Debug, Default)]
58pub struct FCSSchemaBuilder {
59    file_fields: Vec<Field>,
60    partition_fields: Vec<Field>,
61}
62
63impl FCSSchemaBuilder {
64    pub fn new() -> Self {
65        Self {
66            file_fields: vec![],
67            partition_fields: vec![],
68        }
69    }
70
71    pub fn add_file_fields(&mut self, fields: Vec<Field>) {
72        self.file_fields.extend(fields)
73    }
74
75    /// Add fields to the schema builder.
76    pub fn add_partition_fields(&mut self, fields: Vec<Field>) {
77        self.partition_fields.extend(fields)
78    }
79
80    /// Add build the schema and projection.
81    pub fn build(self) -> (Schema, Vec<usize>) {
82        let mut fields = self.file_fields.clone();
83        fields.extend(self.partition_fields);
84
85        let schema = Schema::new(fields);
86
87        let projection = (0..self.file_fields.len()).collect();
88
89        (schema, projection)
90    }
91}
exon_fcs/config.rs

exon_fcs/
config.rs
