From e3eea4794ac3b2447dfd472556afa8df68830dab Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 21 Oct 2025 23:08:50 +0200 Subject: [PATCH] Keep writer state across batches --- beacon-arrow-odv/src/writer.rs | 156 +++++++++++++++++++++++++++++++-- 1 file changed, 148 insertions(+), 8 deletions(-) diff --git a/beacon-arrow-odv/src/writer.rs b/beacon-arrow-odv/src/writer.rs index d0cd9b1c..ebe2e433 100644 --- a/beacon-arrow-odv/src/writer.rs +++ b/beacon-arrow-odv/src/writer.rs @@ -1,4 +1,6 @@ +use core::time; use std::{ + collections::HashMap, fs::File, io::{Read, Seek, Write}, path::{Path, PathBuf}, @@ -7,8 +9,8 @@ use std::{ use arrow::{ array::{ - Array, ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray, RecordBatch, StringArray, - StringBuilder, + Array, ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray, RecordBatch, Scalar, + StringArray, StringBuilder, }, datatypes::{DataType, Field, Schema, SchemaRef}, }; @@ -302,6 +304,93 @@ pub struct AsyncOdvWriter { trajectory_file: OdvFile, zip_file_writer: ZipFileWriter, + + /// Key entry cache + trajectory_profile_key_entry_state: Option, + trajectory_time_series_key_entry_state: Option, +} + +struct FileWriterState { + entry_key: Scalar, + key_state: KeyEntryState, +} + +impl FileWriterState { + fn is_equal_entry_key(&self, other_key: &Scalar) -> bool { + let eq_array = arrow::compute::kernels::cmp::eq(&self.entry_key, other_key).unwrap(); + eq_array + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + } + + fn is_equal_key_state(&self, other_state: &KeyEntryState) -> bool { + self.key_state == *other_state + } +} + +#[derive(Default, Clone)] +struct KeyEntryState { + key: Vec>, + count: usize, +} + +impl KeyEntryState { + fn new(key: Vec>) -> Self { + Self { key, count: 0 } + } +} + +impl PartialEq for KeyEntryState { + fn eq(&self, other: &Self) -> bool { + if self.key.len() != other.key.len() { + return false; + } + + for (a, b) in self.key.iter().zip(other.key.iter()) { + let eq_array = arrow::compute::kernels::cmp::eq(a, b).unwrap(); + if eq_array + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + == false + { + return false; + } + } + + true + } +} + +impl PartialOrd for KeyEntryState { + fn partial_cmp(&self, other: &Self) -> Option { + for (a, b) in self.key.iter().zip(other.key.iter()) { + let lt_array = arrow::compute::kernels::cmp::lt(a, b).unwrap(); + if lt_array + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + { + return Some(std::cmp::Ordering::Less); + } + + let gt_array = arrow::compute::kernels::cmp::gt(a, b).unwrap(); + if gt_array + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + { + return Some(std::cmp::Ordering::Greater); + } + } + + Some(std::cmp::Ordering::Equal) + } } impl AsyncOdvWriter { @@ -344,6 +433,8 @@ impl AsyncOdvWriter { trajectory_file, zip_file_writer, compression, + trajectory_profile_key_entry_state: None, + trajectory_time_series_key_entry_state: None, }) } @@ -425,7 +516,12 @@ impl AsyncOdvWriter { let key_batches = Self::key_batches(record_batch, &self.options)?; let mut classified_batches = vec![]; for batch in key_batches { - let batch_type = Self::classify_batch(batch, &self.options)?; + let batch_type = Self::classify_batch( + batch, + &self.options, + &mut self.trajectory_profile_key_entry_state, + &mut self.trajectory_time_series_key_entry_state, + )?; classified_batches.push(batch_type); } @@ -524,6 +620,8 @@ impl AsyncOdvWriter { fn classify_batch( batch: RecordBatch, odv_options: &OdvOptions, + profile_file_state: &mut Option, + time_series_file_state: &mut Option, ) -> anyhow::Result { let lon_col = batch .column_by_name(&odv_options.longitude_column.column_name) @@ -566,6 +664,8 @@ impl AsyncOdvWriter { &odv_options.latitude_column.column_name, &odv_options.time_column.column_name, &odv_options.key_column, + profile_file_state, + time_series_file_state, )? { return Ok(feature_type); } @@ -599,8 +699,10 @@ impl AsyncOdvWriter { feature_type_columns: &str, longitude_column: &str, latitude_column: &str, - time_column: &str, + time_column_name: &str, key_column: &str, + profile_file_state: &mut Option, + time_series_file_state: &mut Option, ) -> anyhow::Result> { let feature_array = batch.column_by_name(feature_type_columns).ok_or_else(|| { anyhow::anyhow!( @@ -638,7 +740,7 @@ impl AsyncOdvWriter { .expect("Latitude column not found") .clone(); let time_column = batch - .column_by_name(time_column) + .column_by_name(time_column_name) .expect("Time column not found") .clone(); @@ -663,9 +765,47 @@ impl AsyncOdvWriter { .expect("Key column not found") .as_string::(); - let updated_key_column = - Arc::new(Self::append_suffix(key_column_array, &format!("[+{}]", i))); + let timestamp_scalar = Scalar::new( + profile + .column_by_name(time_column_name) + .unwrap() + .slice(0, 1), + ); + let current_entry_key_state = + KeyEntryState::new(vec![timestamp_scalar.clone()]); + let entry_key_scalar = + Scalar::new(profile.column_by_name(key_column).unwrap().slice(0, 1)); + + if profile_file_state.is_none() { + *profile_file_state = Some(FileWriterState { + entry_key: entry_key_scalar.clone(), + key_state: current_entry_key_state.clone(), + }); + } + + let count = match profile_file_state.as_mut() { + Some(state) => { + if state.is_equal_entry_key(&entry_key_scalar) { + if state.is_equal_key_state(¤t_entry_key_state) { + state.key_state.count + } else { + state.key_state.count += 1; + state.key_state.count + } + } else { + state.entry_key = entry_key_scalar; + state.key_state = + KeyEntryState::new(vec![timestamp_scalar.clone()]); + state.key_state.count + } + } + None => 0, + }; + let updated_key_column = Arc::new(Self::append_suffix( + key_column_array, + &format!("[+{}]", count), + )); // Update the record batch key column let profile = Self::replace_column(&profile, key_column, updated_key_column).unwrap(); @@ -677,7 +817,7 @@ impl AsyncOdvWriter { } "timeSeriesProfile" => { let time_column = batch - .column_by_name(time_column) + .column_by_name(time_column_name) .expect("Time column not found") .clone();