diff --git a/Cargo.lock b/Cargo.lock index 792a1365..a9d5a59e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -537,6 +537,7 @@ version = "1.2.1" dependencies = [ "anyhow", "arrow", + "futures", "hifitime", "indexmap 2.10.0", "nd-arrow-array", @@ -545,6 +546,7 @@ dependencies = [ "netcdf-sys", "num-traits", "regex", + "serde", "tempfile", "thiserror 2.0.15", "tracing", diff --git a/Cargo.toml b/Cargo.toml index d2ca803b..c4a92503 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ utoipa-axum = "0.2.0" utoipa-scalar = { version = "0.3.0", features = ["axum"] } utoipa-swagger-ui = { version = "9.0.0", features = ["axum"] } -serde = { version = "=1.0.200", features = ["rc"] } +serde = { version = "=1.0.200", features = ["rc", "derive"] } serde_json = "=1.0.120" anyhow = "1.0.95" thiserror = "2.0.12" @@ -27,13 +27,13 @@ tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } glob = "0.3.2" tempfile = "3.15.0" typetag = "0.2.19" -indexmap = "2.7.1" +indexmap = { version = "2.7.1", features = ["serde"]} chrono = { version = "0.4.41", features = ["serde"] } datafusion = "49.0.0" object_store = { version = "0.12.3", features = ["aws"] } -arrow = { version = "^55.2.0" } +arrow = { version = "^55.2.0", features = ["prettyprint"]} arrow-schema = { version = "^55.2.0", features = ["serde"] } parquet = { version = "^55.2.0", features = ["async"] } geoarrow = { version = "=0.4.0" } diff --git a/beacon-api/Cargo.toml b/beacon-api/Cargo.toml index cd1280b3..b44abebe 100644 --- a/beacon-api/Cargo.toml +++ b/beacon-api/Cargo.toml @@ -3,10 +3,12 @@ name = "beacon-api" version = "1.2.0" edition = "2021" +[target.'cfg(not(windows))'.dependencies] +tikv-jemallocator = "0.6.0" + [dependencies] axum = { version = "0.8.1", features = ["tracing"] } tower-http = { version = "0.6.1", features = ["trace", "cors"] } -tikv-jemallocator = "0.6.0" base64 = "0.22.1" futures = { workspace = true } diff --git a/beacon-arrow-netcdf/Cargo.toml b/beacon-arrow-netcdf/Cargo.toml index 6f2c9fc3..889e78b3 100644 --- a/beacon-arrow-netcdf/Cargo.toml +++ b/beacon-arrow-netcdf/Cargo.toml @@ -10,10 +10,12 @@ nd-arrow-array = { git = "https://github.com/maris-development/nd-arrow-array.gi arrow = { workspace=true } anyhow = { workspace=true } tempfile = { workspace=true} +futures = { workspace=true } indexmap = { workspace=true } thiserror = { workspace=true } ndarray = "0.16.1" hifitime = "4.0.2" regex = "1.11.1" num-traits = "0.2.19" -tracing = { workspace = true } \ No newline at end of file +tracing = { workspace = true } +serde = { workspace = true } \ No newline at end of file diff --git a/beacon-arrow-netcdf/src/chunked_stream.rs b/beacon-arrow-netcdf/src/chunked_stream.rs new file mode 100644 index 00000000..a38abf66 --- /dev/null +++ b/beacon-arrow-netcdf/src/chunked_stream.rs @@ -0,0 +1,527 @@ +use std::{collections::HashMap, pin::Pin, sync::Arc}; + +use arrow::datatypes::SchemaRef; +use indexmap::IndexMap; +use nd_arrow_array::{batch::NdRecordBatch, NdArrowArray}; +use netcdf::Variable; + +use crate::{ + error::ArrowNetCDFError, + reader::{global_attribute, read_variable, variable_attribute}, + NcResult, +}; + +fn is_string_dimension(dimension_name: &str) -> bool { + let dimension_name = dimension_name.to_lowercase(); + dimension_name.starts_with("string") + || dimension_name.starts_with("strlen") + || dimension_name.starts_with("strnlen") +} + +pub struct Stream { + chunk_sizes: IndexMap, + chunk_step_state: IndexMap, + dimension_lengths: IndexMap, + file: Arc, + projected_schema: SchemaRef, + is_done: bool, +} + +#[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)] +pub enum Chunking { + Auto { + target_chunk_size: usize, + }, + ChunkSizes(IndexMap), + #[default] + None, +} + +impl Stream { + pub fn new( + file: Arc, + chunking: Option, + projected_schema: SchemaRef, + ) -> NcResult { + let chunking = chunking.unwrap_or_default(); + + // Get all the dimensions used by the variables in the projected schema + let variables = projected_schema + .fields() + .iter() + .filter_map(|f| { + if f.name().contains('.') { + None + } else { + file.variable(f.name()) + } + }) + .collect::>(); + + // For each variable, get its dimensions and lengths. Create a unique list of dimensions. + let mut dimensions: Vec<(String, usize)> = Vec::new(); + let mut dim_set: HashMap = HashMap::new(); + + for variable in &variables { + let var_dims = variable.dimensions(); + for (i, dim) in var_dims.iter().enumerate() { + let dim_name = dim.name(); + let dim_len = dim.len(); + if i == var_dims.len() - 1 && is_string_dimension(&dim_name) { + // Skip string length dimensions at the end + continue; + } + if dim_set.contains_key(&dim_name) { + continue; + } + dimensions.push((dim_name.clone(), dim_len)); + dim_set.insert(dim_name, dim_len); + } + } + + // Validate that each variable shares the same dimensions or consists of only 1 dimension in the list + for variable in &variables { + let var_dims = variable.dimensions(); + if var_dims.len() > 1 { + if var_dims.len() == dimensions.len() { + for dim in var_dims { + let dim_name = dim.name(); + if !dim_set.contains_key(&dim_name) { + return Err(ArrowNetCDFError::Stream(format!( + "Variable {} has dimension {} which is not in the dimension list.", + variable.name(), + dim_name + ))); + } + } + } else { + return Err(ArrowNetCDFError::Stream(format!( + "Variable {} has dimensions {:?} which do not match the overall dimension list {:?}.", + variable.name(), + var_dims.iter().map(|d| d.name()).collect::>(), + dimensions.iter().map(|(n, _)| n).collect::>(), + ))); + } + } else if var_dims.len() == 1 { + let dim_name = var_dims[0].name(); + if !dim_set.contains_key(&dim_name) { + return Err(ArrowNetCDFError::Stream(format!( + "Variable {} has dimension {} which is not in the dimension list.", + variable.name(), + dim_name + ))); + } + } // Scalars are always valid + } + + let chunk_sizes = match chunking { + Chunking::Auto { target_chunk_size } => { + Self::balanced_chunk_sizes(&dimensions, target_chunk_size) + } + Chunking::ChunkSizes(map) => { + // Validate that all dimensions in hash_map are in dimensions + for dim_name in map.keys() { + if !dim_set.contains_key(dim_name) { + return Err(ArrowNetCDFError::Stream(format!( + "Chunk size specified for dimension {} which is not in the dimension list.", + dim_name + ))); + } + } + map + } + Chunking::None => { + // By default, set chunk size to full dimension length + dimensions + .iter() + .map(|(n, l)| (n.clone(), *l)) + .collect::>() + } + }; + + Ok(Self { + chunk_step_state: chunk_sizes.keys().map(|n| (n.clone(), 0)).collect(), + chunk_sizes: IndexMap::from_iter(chunk_sizes), + dimension_lengths: IndexMap::from_iter(dimensions), + file: file.clone(), + projected_schema, + is_done: false, + }) + } + + fn balanced_chunk_sizes( + dimensions: &[(String, usize)], + target_chunk_size: usize, + ) -> IndexMap { + let total_volume: usize = dimensions.iter().map(|(_, len)| *len).product(); + if total_volume == 0 { + return dimensions.iter().map(|(d, _)| (d.clone(), 0)).collect(); + } + + // Only count dimensions > 1 for balancing + let n_active = dimensions.iter().filter(|(_, l)| *l > 1).count().max(1); + + // Scale factor for balancing + let scale = (target_chunk_size as f64 / total_volume as f64).powf(1.0 / n_active as f64); + + let mut chunk_sizes = IndexMap::new(); + for (dim_name, len) in dimensions { + let chunk_size = if *len <= 1 { + 1 + } else { + let est = (*len as f64 * scale).ceil() as usize; + est.clamp(1, *len) + }; + chunk_sizes.insert(dim_name.clone(), chunk_size); + } + + chunk_sizes + } + + fn is_done(&self) -> bool { + self.is_done + } + + fn advance_chunk_state(&mut self) { + // Advance like a multi-dimensional counter + for (dim, step) in self.chunk_step_state.iter_mut() { + if let Some(d) = self.file.dimension(dim) { + let size = self.chunk_sizes[dim]; + *step += size; + if *step < d.len() { + return; // still within bounds + } else { + *step = 0; // reset and carry to next dimension + } + } + } + self.is_done = true; // All dimensions have been processed + } + + fn generate_hyper_slab(&self, dimensions: &[String]) -> Vec { + let mut hyper_slab = Vec::new(); + for dim in dimensions { + let chunk_count = self.chunk_sizes.get(dim).cloned().unwrap_or_else(|| { + self.dimension_lengths.get(dim).cloned().unwrap() // Default to full length + }); + let step = self.chunk_step_state.get(dim).cloned().unwrap_or(0); + let min_chunk_count = self.dimension_lengths[dim] + .saturating_sub(step) + .min(chunk_count); + + hyper_slab.push(DimensionHyperSlab { + start: step, + count: min_chunk_count, + }); + } + hyper_slab + } + + fn read_variable_scalar(variable: &Variable) -> NcResult { + let values = read_variable(variable, None)?; + let array = values.into_nd_arrow_array().unwrap(); + Ok(array) + } + + fn read_variable_attribute( + variable: &Variable, + attribute_name: &str, + ) -> NcResult { + let variable_attribute = variable_attribute(variable, attribute_name)?; + if let Some(attr) = variable_attribute { + Ok(attr.into_nd_arrow_array().unwrap()) + } else { + Err(ArrowNetCDFError::Stream(format!( + "Attribute {} not found for variable {}", + attribute_name, + variable.name() + ))) + } + } + + fn read_global_attribute(file: &netcdf::File, attribute_name: &str) -> NcResult { + let global_attribute = global_attribute(file, attribute_name)?; + if let Some(attr) = global_attribute { + Ok(attr.into_nd_arrow_array().unwrap()) + } else { + Err(ArrowNetCDFError::Stream(format!( + "Global attribute {} not found", + attribute_name + ))) + } + } + + fn read_variable_hyper_slab( + variable: &Variable, + hyper_slab: &[DimensionHyperSlab], + ) -> NcResult { + let mut start: Vec = Vec::new(); + let mut count: Vec = Vec::new(); + + for dim_slab in hyper_slab { + start.push(dim_slab.start); + count.push(dim_slab.count); + } + + let values = read_variable(variable, Some((start, count)))?; + let array = values.into_nd_arrow_array().unwrap(); + + Ok(array) + } +} + +#[derive(Debug, Clone, Copy)] +struct DimensionHyperSlab { + start: usize, + count: usize, +} + +impl Iterator for Stream { + type Item = NcResult; + + fn next(&mut self) -> Option { + if self.is_done() { + return None; + } + + let mut arrays = Vec::new(); + let mut fields = Vec::new(); + + for field in self.projected_schema.fields() { + fields.push(field.clone()); + let array_result = if let Some((var_name, attr_name)) = field.name().split_once('.') { + if var_name.is_empty() { + // Global attribute + Self::read_global_attribute(&self.file, attr_name) + } else { + // Variable attribute + match self.file.variable(var_name) { + Some(variable) => Self::read_variable_attribute(&variable, attr_name), + None => { + return Some(Err(ArrowNetCDFError::InvalidFieldName( + field.name().to_string(), + ))) + } + } + } + } else { + match self.file.variable(field.name()) { + Some(variable) => { + let var_dims: Vec<_> = + variable.dimensions().iter().map(|d| d.name()).collect(); + if var_dims.is_empty() { + Self::read_variable_scalar(&variable) + } else { + let hyper_slab = self.generate_hyper_slab(&var_dims); + println!( + "Reading variable {} with hyper slab: {:?}", + variable.name(), + hyper_slab + ); + Self::read_variable_hyper_slab(&variable, &hyper_slab) + } + } + None => { + return Some(Err(ArrowNetCDFError::InvalidFieldName( + field.name().to_string(), + ))) + } + } + }; + + match array_result { + Ok(array) => arrays.push(array), + Err(e) => return Some(Err(e)), + } + } + + self.advance_chunk_state(); + + let maybe_nd_batch = nd_arrow_array::batch::NdRecordBatch::new( + fields.into_iter().map(|f| f.as_ref().clone()).collect(), + arrays, + ) + .map_err(|e| ArrowNetCDFError::Stream(format!("Failed to create NdRecordBatch: {}", e))); + + Some(maybe_nd_batch) + } +} + +impl futures::Stream for Stream { + type Item = NcResult; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = Pin::get_mut(self); + std::task::Poll::Ready(this.next()) + } +} + +#[cfg(test)] +mod tests { + use std::vec; + + use crate::reader::NetCDFArrowReader; + + use super::*; + + #[test] + fn test_auto_chunk_with_aligned_dimensions() { + let reader = NetCDFArrowReader::new_with_aligned_dimensions( + "test_files/gridded-example.nc", + vec!["time".to_string(), "lat".to_string(), "lon".to_string()], + ) + .unwrap(); + + let stream = reader.read_as_stream::>(None, None); + assert!(stream.is_ok()); + } + + #[test] + fn test_stream_read_auto_chunk() { + let reader = NetCDFArrowReader::new("test_files/gridded-example.nc").unwrap(); + + let temp_idx = reader.schema().index_of("analysed_sst").unwrap(); + let lat = reader.schema().index_of("lat").unwrap(); + let lon = reader.schema().index_of("lon").unwrap(); + + let mut stream = reader + .read_as_stream::>( + Some(vec![temp_idx, lat, lon]), + Some(Chunking::Auto { + target_chunk_size: 128000, + }), + ) + .unwrap(); + + let mut total_rows = 0; + while let Some(batch_result) = stream.next() { + match batch_result { + Ok(batch) => { + let batch = batch.to_arrow_record_batch().unwrap(); + println!("Batch: {:?}", batch); + total_rows += batch.num_rows(); + } + Err(e) => { + eprintln!("Error reading batch: {}", e); + } + } + } + println!("Total rows read: {}", total_rows); + } + + #[test] + fn test_stream_read_no_chunking() { + let reader = NetCDFArrowReader::new("test_files/gridded-example.nc").unwrap(); + + let temp_idx = reader.schema().index_of("analysed_sst").unwrap(); + let lat = reader.schema().index_of("lat").unwrap(); + let lon = reader.schema().index_of("lon").unwrap(); + + let mut stream = reader + .read_as_stream::>(Some(vec![temp_idx, lat, lon]), Some(Chunking::None)) + .unwrap(); + + let mut total_rows = 0; + while let Some(batch_result) = stream.next() { + match batch_result { + Ok(batch) => { + let batch = batch.to_arrow_record_batch().unwrap(); + println!("Batch: {:?}", batch); + total_rows += batch.num_rows(); + } + Err(e) => { + eprintln!("Error reading batch: {}", e); + } + } + } + println!("Total rows read: {}", total_rows); + } + + #[test] + fn test_stream_read_custom_chunking() { + let reader = NetCDFArrowReader::new("test_files/gridded-example.nc").unwrap(); + + let temp = reader.schema().index_of("analysed_sst").unwrap(); + let temp_units = reader.schema().index_of("analysed_sst.units").unwrap(); + let lat = reader.schema().index_of("lat").unwrap(); + let lon = reader.schema().index_of("lon").unwrap(); + let processing_level = reader.schema().index_of(".processing_level").unwrap(); + + let mut stream = reader + .read_as_stream::>( + Some(vec![temp, temp_units, lat, lon, processing_level]), + Some(Chunking::ChunkSizes( + [("lat".to_string(), 500), ("lon".to_string(), 500)] + .into_iter() + .collect(), + )), + ) + .unwrap(); + + let mut total_rows = 0; + while let Some(batch_result) = stream.next() { + match batch_result { + Ok(batch) => { + let batch = batch.to_arrow_record_batch().unwrap(); + println!("Batch: {:?}", batch); + total_rows += batch.num_rows(); + } + Err(e) => { + eprintln!("Error reading batch: {}", e); + } + } + } + println!("Total rows read: {}", total_rows); + } + + #[test] + fn test_flat_read() { + let reader = NetCDFArrowReader::new("test_files/gridded-example.nc").unwrap(); + + let temp = reader.schema().index_of("analysed_sst").unwrap(); + let temp_units = reader.schema().index_of("analysed_sst.units").unwrap(); + let lat = reader.schema().index_of("lat").unwrap(); + let lon = reader.schema().index_of("lon").unwrap(); + let processing_level = reader.schema().index_of(".processing_level").unwrap(); + + let mut batch = reader + .read_as_batch::>(Some(vec![temp, temp_units, lat, lon, processing_level])) + .unwrap(); + + println!("Total rows read: {}", batch.num_rows()); + } + + #[test] + fn test_chunked_column_read() { + let reader = NetCDFArrowReader::new_with_aligned_dimensions( + "test_files/gridded-example.nc", + vec!["time".to_string(), "lat".to_string(), "lon".to_string()], + ) + .unwrap(); + + let chunking = Chunking::ChunkSizes( + [("lat".to_string(), 500), ("lon".to_string(), 500)] + .into_iter() + .collect(), + ); + let mut temp_chunked_column = reader + .read_column_as_stream("analysed_sst", Some(chunking)) + .unwrap(); + + while let Some(batch_result) = temp_chunked_column.next() { + match batch_result { + Ok(batch) => { + let batch = batch.to_arrow_record_batch().unwrap(); + // println!("Batch: {:?}", batch); + println!("Batch rows: {}", batch.num_rows()); + } + Err(e) => { + eprintln!("Error reading batch: {}", e); + } + } + } + } +} diff --git a/beacon-arrow-netcdf/src/error.rs b/beacon-arrow-netcdf/src/error.rs index af7def17..032c462a 100644 --- a/beacon-arrow-netcdf/src/error.rs +++ b/beacon-arrow-netcdf/src/error.rs @@ -50,4 +50,8 @@ pub enum ArrowNetCDFError { IpcBufferCloseError(ArrowError), #[error("Ipc Buffer Failed to Open for reading: {0}")] IpcBufferOpenError(ArrowError), + #[error("Stream Error: {0}")] + Stream(String), + #[error("Reader Error: {0}")] + Reader(String), } diff --git a/beacon-arrow-netcdf/src/lib.rs b/beacon-arrow-netcdf/src/lib.rs index edf4a8fe..40cf81c1 100644 --- a/beacon-arrow-netcdf/src/lib.rs +++ b/beacon-arrow-netcdf/src/lib.rs @@ -3,6 +3,7 @@ use std::ffi::CString; use netcdf::{types::NcVariableType, NcTypeDescriptor}; pub mod cf_time; +pub mod chunked_stream; pub mod encoders; pub mod error; pub mod nc_array; diff --git a/beacon-arrow-netcdf/src/nc_array.rs b/beacon-arrow-netcdf/src/nc_array.rs index 06ac0e46..ac663c31 100644 --- a/beacon-arrow-netcdf/src/nc_array.rs +++ b/beacon-arrow-netcdf/src/nc_array.rs @@ -22,6 +22,12 @@ pub struct Dimension { pub size: usize, } +impl Dimension { + pub fn new(name: String, size: usize) -> Self { + Self { name, size } + } +} + pub struct NetCDFNdArray { pub dims: Vec, pub array: NetCDFNdArrayInner, diff --git a/beacon-arrow-netcdf/src/reader.rs b/beacon-arrow-netcdf/src/reader.rs index 510d58d7..1146cee0 100644 --- a/beacon-arrow-netcdf/src/reader.rs +++ b/beacon-arrow-netcdf/src/reader.rs @@ -1,6 +1,6 @@ use std::{path::Path, sync::Arc}; -use arrow::array::RecordBatch; +use arrow::{array::RecordBatch, datatypes::Schema}; use nd_arrow_array::NdArrowArray; use ndarray::{ArrayBase, ArrayD}; use netcdf::{ @@ -10,6 +10,7 @@ use netcdf::{ use crate::{ cf_time::{decode_cf_time_variable, is_cf_time_variable}, + chunked_stream::{Chunking, Stream}, error::ArrowNetCDFError, nc_array::{Dimension, NetCDFNdArray, NetCDFNdArrayBase, NetCDFNdArrayInner}, NcChar, NcResult, @@ -17,14 +18,82 @@ use crate::{ pub struct NetCDFArrowReader { file_schema: arrow::datatypes::SchemaRef, - file: netcdf::File, + file: Arc, } impl NetCDFArrowReader { pub fn new>(path: P) -> NcResult { let file = netcdf::open(path)?; let file_schema = Arc::new(arrow_schema(&file)?); - Ok(Self { file_schema, file }) + Ok(Self { + file_schema, + file: Arc::new(file), + }) + } + + pub fn new_with_aligned_dimensions>( + path: P, + dimensions: Vec, + ) -> NcResult { + let file = netcdf::open(path)?; + let file_schema = arrow_schema(&file)?; + + // Align dimensions + file.dimensions().try_for_each(|dim| { + if dimensions.contains(&dim.name().to_string()) { + Ok(()) + } else { + Err(ArrowNetCDFError::Reader(format!( + "Dimension '{}' not found in NetCDF file.", + dim.name() + ))) + } + })?; + + // Check all the variables, and keep only the variables which have one of the specified dimensions or all of them + // Scalars (0D variables) are always kept + let mut removable_variables = vec![]; + for variable in file.variables() { + let variable_dimensions = variable.dimensions(); + + if variable_dimensions.is_empty() { + continue; // Scalar variable, keep it + } + if variable.dimensions().len() == 1 { + let dimension_name = variable_dimensions[0].name(); + if !dimensions.contains(&dimension_name) { + removable_variables.push(variable.name().to_string()); + } + } else if !variable + .dimensions() + .iter() + .all(|d| dimensions.contains(&d.name().to_string())) + { + removable_variables.push(variable.name().to_string()); + } + } + + // Create a new schema excluding the removable variables + let fields: Vec = file_schema + .fields() + .iter() + .filter(|field| { + let field_name = field.name(); + // Remove the field if in the removable variables list + !removable_variables.contains(field_name) + // Also remove variable attributes if the parent variable is removed + || (field_name.contains('.') && { + let parts: Vec<&str> = field_name.split('.').collect(); + parts.len() == 2 && !removable_variables.contains(&parts[0].to_string()) + }) + }) + .map(|f| f.as_ref().clone()) + .collect::>(); + + Ok(Self { + file_schema: Arc::new(Schema::new(fields)), + file: Arc::new(file), + }) } pub fn dimensions(&self) -> Vec { @@ -51,69 +120,44 @@ impl NetCDFArrowReader { } else { self.file_schema.clone() }; + let mut stream = Stream::new(self.file.clone(), None, projected_schema)?; + let batch = stream.next().transpose().unwrap().ok_or_else(|| { + ArrowNetCDFError::Stream("No data available in the NetCDF file.".to_string()) + })?; + batch.to_arrow_record_batch().map_err(|e| { + ArrowNetCDFError::Stream(format!("Failed to flatten to RecordBatch: {}", e)) + }) + } - let mut columns = indexmap::IndexMap::new(); - for field in projected_schema.fields() { - let name = field.name(); - if name.contains('.') { - let parts = name.split('.').collect::>(); - if parts.len() != 2 { - return Err(ArrowNetCDFError::InvalidFieldName(name.to_string())); - } - if parts[0].is_empty() { - //Global attribute - let attr_name = parts[1]; - let attr_value = global_attribute(&self.file, attr_name)? - .expect("Attribute not found but was in schema."); - columns.insert( - field.clone(), - attr_value - .into_nd_arrow_array() - .map_err(ArrowNetCDFError::NdArrowError)?, - ); - } else { - //Variable attribute - let variable = self - .file - .variable(parts[0]) - .expect("Variable not found but was in schema."); - columns.insert( - field.clone(), - variable_attribute(&variable, parts[1])? - .expect("Attribute not found but was in schema.") - .into_nd_arrow_array() - .map_err(ArrowNetCDFError::NdArrowError)?, - ); - } - } else { - let variable = self - .file - .variable(name) - .expect("Variable not found but was in schema."); - let array = read_variable(&variable) - .map_err(|e| ArrowNetCDFError::VariableReadError(Box::new(e)))?; - columns.insert( - field.clone(), - array - .into_nd_arrow_array() - .map_err(ArrowNetCDFError::NdArrowError)?, - ); - } - } - - let mut fields = vec![]; - let mut arrays = vec![]; - for (field, array) in columns { - fields.push(field.as_ref().clone()); - arrays.push(array); - } - - let nd_batch = nd_arrow_array::batch::NdRecordBatch::new(fields, arrays).unwrap(); - let record_batch = nd_batch - .to_arrow_record_batch() - .map_err(ArrowNetCDFError::NdArrowError)?; + pub fn read_as_stream>( + &self, + projection: Option

, + chunking: Option, + ) -> NcResult { + let projected_schema = if let Some(projection) = projection { + Arc::new( + self.file_schema + .project(projection.as_ref()) + .map_err(ArrowNetCDFError::ArrowSchemaProjectionError)?, + ) + } else { + self.file_schema.clone() + }; + Stream::new(self.file.clone(), chunking, projected_schema) + } - Ok(record_batch) + pub fn read_column_as_stream>( + &self, + column_name: P, + chunking: Option, + ) -> NcResult { + let column_name = column_name.as_ref(); + let column_index = self + .file_schema + .index_of(column_name) + .map_err(|_| ArrowNetCDFError::InvalidFieldName(column_name.to_string()))?; + + self.read_as_stream(Some(&[column_index]), chunking) } pub fn read_column(&self, column_name: &str) -> NcResult { @@ -150,7 +194,7 @@ impl NetCDFArrowReader { .file .variable(column_name) .expect("Variable not found but was in schema."); - let array = read_variable(&variable) + let array = read_variable(&variable, None) .map_err(|e| ArrowNetCDFError::VariableReadError(Box::new(e)))?; Ok(array.into_nd_arrow_array().unwrap()) } @@ -158,19 +202,22 @@ impl NetCDFArrowReader { } macro_rules! create_netcdf_ndarray { - ($var:ident, $t:ty, $inner_variant:ident, $as_ref:ident) => {{ - let array = $var.get::<$t, _>(Extents::All)?; - let dims = $var - .dimensions() - .iter() - .map(|d| Dimension { - name: d.name().to_string(), - size: d.len(), - }) - .collect::>(); + ($var:ident, $t:ty, $inner_variant:ident, $as_ref:ident, $dims:expr, $extents:expr) => {{ + let array = $var.get::<$t, _>($extents)?; + let dims = if let Some(dims) = $dims { + dims + } else { + $var.dimensions() + .iter() + .map(|d| Dimension { + name: d.name().to_string(), + size: d.len(), + }) + .collect::>() + }; let mut fill_value = $var.fill_value::<$t>()?; - if fill_value.is_none() { + { fill_value = read_fill_value_attribute(&$var)?.and_then(|fv| fv.$as_ref()) } @@ -187,7 +234,10 @@ macro_rules! create_netcdf_ndarray { }}; } -pub fn read_variable(variable: &Variable) -> NcResult { +pub fn read_variable( + variable: &Variable, + mut extents_start_count: Option<(Vec, Vec)>, +) -> NcResult { if is_cf_time_variable(variable) { if let Some(array) = decode_cf_time_variable(variable) .map_err(|e| ArrowNetCDFError::TimeVariableReadError(Box::new(e)))? @@ -196,41 +246,56 @@ pub fn read_variable(variable: &Variable) -> NcResult { } } + let mut extents = extents_start_count + .clone() + .map(|e| e.try_into()) + .unwrap_or(Ok(Extents::All)) + .unwrap(); + let mut dims = vec![]; + for (i, dim) in variable.dimensions().iter().enumerate() { + match extents_start_count.as_ref() { + Some((_, count)) => { + dims.push(Dimension::new(dim.name(), count[i])); + } + None => { + dims.push(Dimension::new(dim.name(), dim.len())); + } + } + } + match variable.vartype() { netcdf::types::NcVariableType::Int(IntType::I8) => { - create_netcdf_ndarray!(variable, i8, I8, as_i8) + create_netcdf_ndarray!(variable, i8, I8, as_i8, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::I16) => { - create_netcdf_ndarray!(variable, i16, I16, as_i16) + create_netcdf_ndarray!(variable, i16, I16, as_i16, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::I32) => { - create_netcdf_ndarray!(variable, i32, I32, as_i32) + create_netcdf_ndarray!(variable, i32, I32, as_i32, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::I64) => { - create_netcdf_ndarray!(variable, i64, I64, as_i64) + create_netcdf_ndarray!(variable, i64, I64, as_i64, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::U8) => { - create_netcdf_ndarray!(variable, u8, U8, as_u8) + create_netcdf_ndarray!(variable, u8, U8, as_u8, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::U16) => { - create_netcdf_ndarray!(variable, u16, U16, as_u16) + create_netcdf_ndarray!(variable, u16, U16, as_u16, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::U32) => { - create_netcdf_ndarray!(variable, u32, U32, as_u32) + create_netcdf_ndarray!(variable, u32, U32, as_u32, Some(dims), extents) } netcdf::types::NcVariableType::Int(IntType::U64) => { - create_netcdf_ndarray!(variable, u64, U64, as_u64) + create_netcdf_ndarray!(variable, u64, U64, as_u64, Some(dims), extents) } netcdf::types::NcVariableType::Float(FloatType::F32) => { - create_netcdf_ndarray!(variable, f32, F32, as_f32) + create_netcdf_ndarray!(variable, f32, F32, as_f32, Some(dims), extents) } netcdf::types::NcVariableType::Float(FloatType::F64) => { - create_netcdf_ndarray!(variable, f64, F64, as_f64) + create_netcdf_ndarray!(variable, f64, F64, as_f64, Some(dims), extents) } netcdf::types::NcVariableType::Char => { // NcChar is both a value itself or possibly a fixed size string - let array = variable.get::(Extents::All)?; - //Get the last dimension of the variable if let Some(dim) = variable.dimensions().last() { //Check if the dimensions starts with string** or strlen** @@ -239,6 +304,25 @@ pub fn read_variable(variable: &Variable) -> NcResult { || dim_name.starts_with("strlen") || dim_name.starts_with("strnlen") { + // Append the dimension and extents + dims.push(Dimension { + name: dim.name().to_string(), + size: dim.len(), + }); + + extents_start_count.as_mut().map(|ex| { + ex.0.push(0); + ex.1.push(dim.len()); + }); + + extents = extents_start_count + .clone() + .map(|e| e.try_into()) + .unwrap_or(Ok(Extents::All)) + .unwrap(); + + let array = variable.get::(&extents)?; + let dims = variable .dimensions() .iter() @@ -262,7 +346,7 @@ pub fn read_variable(variable: &Variable) -> NcResult { }); } } - create_netcdf_ndarray!(variable, NcChar, Char, as_nc_char) + create_netcdf_ndarray!(variable, NcChar, Char, as_nc_char, Some(dims), extents) } nctype => Err(ArrowNetCDFError::UnsupportedNetCDFDataType(nctype)), } diff --git a/beacon-arrow-netcdf/test_files/gridded-example.nc b/beacon-arrow-netcdf/test_files/gridded-example.nc new file mode 100644 index 00000000..f1a04a19 Binary files /dev/null and b/beacon-arrow-netcdf/test_files/gridded-example.nc differ