From 24b484d77287cef8de084774b9abeb3bcf34846e Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:43:37 -0500 Subject: [PATCH 1/2] Keep current writer rows queryable across periodic flushes --- src/parseable/staging/writer.rs | 184 +++++++++++++++++++++++++++++++- src/parseable/streams.rs | 165 ++++++++++++++++++++++++++-- 2 files changed, 335 insertions(+), 14 deletions(-) diff --git a/src/parseable/staging/writer.rs b/src/parseable/staging/writer.rs index 919e78c23..04be9f1d8 100644 --- a/src/parseable/staging/writer.rs +++ b/src/parseable/staging/writer.rs @@ -32,7 +32,7 @@ use arrow_ipc::{ }; use arrow_schema::Schema; use arrow_select::concat::concat_batches; -use chrono::Utc; +use chrono::{DateTime, Utc}; use itertools::Itertools; use once_cell::sync::Lazy; use rand::distributions::{Alphanumeric, DistString}; @@ -112,18 +112,34 @@ impl Writer { } } - /// Detaches completed disk writers, or all writers when `forced` is true. + /// Detaches completed disk writers and evicts their memory segments. + /// Forced flushes also detach current writers; ownerless memory always expires. pub fn take_flushable_disk(&mut self, forced: bool) -> HashMap { + self.take_flushable_disk_at(forced, Utc::now()) + } + + pub(crate) fn take_flushable_disk_at( + &mut self, + forced: bool, + now: DateTime, + ) -> HashMap { let mut flushable_disk = HashMap::new(); let old_disk = std::mem::take(&mut self.disk); for (filename, writer) in old_disk { - if !forced && writer.is_current() { + if !forced && writer.range.contains(now) { self.disk.insert(filename, writer); } else { flushable_disk.insert(filename, writer); } } + if let Some(mem) = self.mem.as_mut() { + // Use the selected writer keys, not a second clock/range check. + // Query-only segments have no disk owner and expire at each flush. + mem.clear_unsegmented(); + mem.segments + .retain(|filename, _| self.disk.contains_key(filename)); + } flushable_disk } } @@ -217,6 +233,9 @@ impl Drop for DiskWriter { /// Recordbatches are pushed to mutable buffer first and then concated together and pushed to read buffer #[derive(Debug)] pub struct MemWriter { + // Private file-owned buffers preserve the existing public MemWriter shape. + // Stream query callers use recordbatch_cloned to read every segment. + segments: HashMap>, schema: Schema, // for checking uniqueness of schema schema_map: HashSet, @@ -227,6 +246,7 @@ pub struct MemWriter { impl Default for MemWriter { fn default() -> Self { Self { + segments: HashMap::new(), schema: Schema::empty(), schema_map: HashSet::default(), read_buffer: Vec::default(), @@ -236,6 +256,18 @@ impl Default for MemWriter { } impl MemWriter { + pub(crate) fn push_for_file( + &mut self, + filename: String, + schema_key: &str, + rb: &RecordBatch, + ) -> Result<(), StagingError> { + self.segments + .entry(filename) + .or_default() + .push(schema_key, rb) + } + pub fn push(&mut self, schema_key: &str, rb: &RecordBatch) -> Result<(), StagingError> { if !self.schema_map.contains(schema_key) { self.schema_map.insert(schema_key.to_owned()); @@ -250,6 +282,11 @@ impl MemWriter { } pub fn clear(&mut self) { + self.segments.clear(); + self.clear_unsegmented(); + } + + fn clear_unsegmented(&mut self) { self.schema = Schema::empty(); self.schema_map.clear(); self.read_buffer.clear(); @@ -266,10 +303,14 @@ impl MemWriter { read_buffer.push(rb) } - Ok(read_buffer + let mut batches = read_buffer .into_iter() .map(|rb| adapt_batch(schema.clone(), &rb)) - .collect()) + .collect::>(); + for segment in self.segments.values() { + batches.extend(segment.recordbatch_cloned(schema)?); + } + Ok(batches) } } @@ -331,6 +372,139 @@ mod tests { use super::{DiskWriter, Writer}; + #[test] + fn current_minute_query_visibility_segments_follow_disk_ownership() { + let dir = TempDir::new().unwrap(); + let now = chrono::DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") + .unwrap() + .with_timezone(&Utc); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let extended = Arc::new(Schema::new(vec![ + Field::new("extra", DataType::Int32, true), + Field::new("id", DataType::Int32, false), + ])); + let query_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("extra", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![10, 11]))], + ) + .unwrap(); + let extended_batch = RecordBatch::try_new( + extended, + vec![ + Arc::new(Int32Array::from(vec![99])), + Arc::new(Int32Array::from(vec![12])), + ], + ) + .unwrap(); + let mut writer = Writer::default(); + let current_range = TimeRange { + start: now - chrono::TimeDelta::hours(1), + end: now + chrono::TimeDelta::hours(1), + }; + for (name, range, rb, key) in [ + ( + "past", + TimeRange { + start: now - chrono::TimeDelta::hours(3), + end: now - chrono::TimeDelta::hours(2), + }, + &batch, + "id", + ), + ("current-a", current_range.clone(), &batch, "id"), + ("current-b", current_range, &extended_batch, "extended"), + ( + "future", + TimeRange { + start: now + chrono::TimeDelta::hours(2), + end: now + chrono::TimeDelta::hours(3), + }, + &batch, + "id", + ), + ] { + let filename = format!("{name}.date=2020-01-02.data.arrows"); + writer + .push_disk(filename.clone(), rb, dir.path().join(&filename), range) + .unwrap(); + writer + .mem + .as_mut() + .unwrap() + .segments + .entry(filename) + .or_default() + .push(key, rb) + .unwrap(); + } + // Query-mode rows have no disk owner, so retain the existing flush behavior. + writer + .mem + .as_mut() + .unwrap() + .segments + .entry("query-only".to_owned()) + .or_default() + .push("id", &batch) + .unwrap(); + let key = "current-a.date=2020-01-02.data.arrows"; + assert!( + Arc::ptr_eq( + writer.mem.as_ref().unwrap().segments[key] + .mutable_buffer + .inner[0] + .column(0), + batch.column(0) + ), + "ingestion must keep the existing Arrow allocation" + ); + let detached = writer.take_flushable_disk_at(false, now); + assert_eq!(detached.len(), 2); + assert_eq!(writer.disk.len(), 2); + let mem = &writer.mem.as_ref().unwrap().segments; + assert_eq!(mem.len(), 2); + eprintln!( + "memory segment fixed header: {} bytes; filename String header: {} bytes (heap buffers/hash buckets excluded)", + std::mem::size_of::>(), + std::mem::size_of::() + ); + assert!( + Arc::ptr_eq(mem[key].mutable_buffer.inner[0].column(0), batch.column(0)), + "flushing must retain the same Arrow allocation" + ); + let mut ids = Vec::new(); + for segment in mem.values() { + for rb in segment.recordbatch_cloned(&query_schema).unwrap() { + assert_eq!(rb.schema(), query_schema); + ids.extend_from_slice( + rb.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + ); + let extras = rb.column(1).as_any().downcast_ref::().unwrap(); + if rb.num_rows() == 1 { + assert_eq!(extras.value(0), 99); + } else { + assert_eq!(arrow_array::Array::null_count(extras), 2); + } + } + } + ids.sort_unstable(); + assert_eq!(ids, vec![10, 11, 12]); + drop(detached); + assert!(writer.take_flushable_disk_at(false, now).is_empty()); + assert_eq!(writer.mem.as_ref().unwrap().segments.len(), 2); + drop(writer.take_flushable_disk_at(true, now)); + assert!(writer.disk.is_empty()); + assert!(writer.mem.as_ref().unwrap().segments.is_empty()); + } + #[test] fn disk_writer_flushes_each_complete_batch_before_drop() { let dir = TempDir::new().unwrap(); diff --git a/src/parseable/streams.rs b/src/parseable/streams.rs index 6832bd7f6..6569ba06f 100644 --- a/src/parseable/streams.rs +++ b/src/parseable/streams.rs @@ -23,7 +23,7 @@ use arrow::{ }; use arrow_array::{ArrayRef, RecordBatch, UInt32Array}; use arrow_schema::{ArrowError, Field, Fields, Schema, SortOptions}; -use chrono::{NaiveDate, NaiveDateTime, Timelike, Utc}; +use chrono::{DateTime, NaiveDate, NaiveDateTime, Timelike, Utc}; use derive_more::derive::{Deref, DerefMut}; use itertools::Itertools; use once_cell::sync::{Lazy, OnceCell}; @@ -724,20 +724,20 @@ impl Stream { } } }; + let filename = + self.filename_by_partition(schema_key, parsed_timestamp, custom_partition_values); if self.options.mode != Mode::Query || stream_type == StreamType::Internal { - let filename = - self.filename_by_partition(schema_key, parsed_timestamp, custom_partition_values); let range = TimeRange::granularity_range( parsed_timestamp.and_local_timezone(Utc).unwrap(), OBJECT_STORE_DATA_GRANULARITY, ); let file_path = self.data_path.join(&filename); - guard.push_disk(filename, record, file_path, range)?; + guard.push_disk(filename.clone(), record, file_path, range)?; } if let Some(mem) = guard.mem.as_mut() { - mem.push(schema_key, record)?; + mem.push_for_file(filename, schema_key, record)?; } Ok(()) @@ -1166,6 +1166,14 @@ impl Stream { #[cfg_attr(feature = "hotpath", hotpath::measure)] pub fn flush(&self, forced: bool) -> Result<(), StagingError> { + self.flush_with_clock(forced, Utc::now) + } + + fn flush_with_clock( + &self, + forced: bool, + clock: impl FnOnce() -> DateTime, + ) -> Result<(), StagingError> { let _span = info_span!("flush", stream_name = %self.stream_name, forced).entered(); // Swap out stale writers under the lock, drop them after releasing it. // DiskWriter::Drop does I/O (IPC finish + file rename) so dropping @@ -1178,10 +1186,7 @@ impl Stream { ))) })?; - if let Some(mem) = writer.mem.as_mut() { - mem.clear(); - } - writer.take_flushable_disk(forced) + writer.take_flushable_disk_at(forced, clock()) }; // DiskWriter::Drop I/O happens here, outside the lock drop(stale_writers); @@ -2967,6 +2972,148 @@ mod tests { use super::*; + #[test] + fn current_minute_query_visibility_query_mode_flushes_ownerless_memory() { + let dir = TempDir::new().unwrap(); + let staging = Stream::new( + Arc::new(Options { + local_staging_path: dir.path().to_path_buf(), + mode: Mode::Query, + ..Default::default() + }), + "query-memory", + LogStreamMetadata::default(), + None, + &None, + ); + let now = DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") + .unwrap() + .with_timezone(&Utc); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + staging + .push( + "id-schema", + &batch, + now.naive_utc(), + &HashMap::new(), + StreamType::UserDefined, + ) + .unwrap(); + assert!(staging.writer.lock().unwrap().disk.is_empty()); + assert_eq!( + staging.recordbatches_cloned(&schema).unwrap()[0].num_rows(), + 1 + ); + staging.flush_with_clock(false, || now).unwrap(); + assert!(staging.recordbatches_cloned(&schema).unwrap().is_empty()); + } + + #[test] + fn current_minute_query_visibility_survives_nonforced_flush() { + // A fixed reference clock tests the production push/flush/query path without + // relying on a wall-clock minute boundary or retrying the test. + let now = DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") + .unwrap() + .with_timezone(&Utc); + let dir = TempDir::new().unwrap(); + let staging = Stream::new( + Arc::new(Options { + local_staging_path: dir.path().to_path_buf(), + ..Default::default() + }), + "visibility", + LogStreamMetadata::default(), + None, + &None, + ); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + for (id, timestamp, partition) in [ + (1, now - TimeDelta::minutes(2), "past"), + (2, now, "a"), + (3, now, "b"), + (4, now + TimeDelta::minutes(2), "future"), + ] { + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![id]))]) + .unwrap(); + staging + .push( + "id-schema", + &batch, + timestamp.naive_utc(), + &HashMap::from([("tenant".to_owned(), partition.to_owned())]), + StreamType::UserDefined, + ) + .unwrap(); + } + let ids = || { + let mut ids = staging + .recordbatches_cloned(&schema) + .unwrap() + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + ids.sort_unstable(); + ids + }; + assert_eq!(ids(), vec![1, 2, 3, 4]); + staging.flush_with_clock(false, || now).unwrap(); + assert_eq!(staging.writer.lock().unwrap().disk.len(), 2); + assert_eq!( + ids(), + vec![2, 3], + "attached current writers must remain queryable" + ); + staging.flush_with_clock(false, || now).unwrap(); + assert_eq!( + ids(), + vec![2, 3], + "a second flush must neither lose nor duplicate rows" + ); + staging.clear().unwrap(); + assert!(ids().is_empty()); + assert_eq!(staging.writer.lock().unwrap().disk.len(), 2); + let next = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![5]))]) + .unwrap(); + staging + .push( + "id-schema", + &next, + now.naive_utc(), + &HashMap::from([("tenant".to_owned(), "a".to_owned())]), + StreamType::UserDefined, + ) + .unwrap(); + assert_eq!(ids(), vec![5]); + staging.flush_with_clock(true, || now).unwrap(); + assert!(ids().is_empty()); + assert!(staging.writer.lock().unwrap().disk.is_empty()); + // A late append for a previously detached file starts a fresh segment. + staging + .push( + "id-schema", + &next, + (now - TimeDelta::minutes(2)).naive_utc(), + &HashMap::from([("tenant".to_owned(), "past".to_owned())]), + StreamType::UserDefined, + ) + .unwrap(); + assert_eq!(ids(), vec![5]); + staging.flush_with_clock(false, || now).unwrap(); + assert!(ids().is_empty()); + assert!(staging.writer.lock().unwrap().disk.is_empty()); + } + #[test] fn parallel_metric_writer_creates_valid_ordered_row_groups() { let temp_dir = TempDir::new().unwrap(); From 4ce678c3d383f37eb2b9c6b02360239403c85432 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:56:58 -0500 Subject: [PATCH 2/2] Reduce staging segment capacity and filename allocation overhead --- src/parseable/staging/writer.rs | 132 ++++++++++++++++++++++++++------ src/parseable/streams.rs | 105 +++++++++++++++++++++++-- 2 files changed, 204 insertions(+), 33 deletions(-) diff --git a/src/parseable/staging/writer.rs b/src/parseable/staging/writer.rs index 04be9f1d8..1f3f47888 100644 --- a/src/parseable/staging/writer.rs +++ b/src/parseable/staging/writer.rs @@ -90,7 +90,6 @@ impl Default for Writer { impl Writer { /// Appends a batch to the writer for its schema/minute file. - #[cfg_attr(feature = "hotpath", hotpath::measure)] pub fn push_disk( &mut self, filename: String, @@ -98,7 +97,18 @@ impl Writer { file_path: PathBuf, range: TimeRange, ) -> Result<(), StagingError> { - match self.disk.get_mut(&filename) { + self.push_disk_ref(&filename, rb, file_path, range) + } + + #[cfg_attr(feature = "hotpath", hotpath::measure)] + pub(crate) fn push_disk_ref( + &mut self, + filename: &str, + rb: &RecordBatch, + file_path: PathBuf, + range: TimeRange, + ) -> Result<(), StagingError> { + match self.disk.get_mut(filename) { Some(writer) => writer.write(rb), None => { if let Some(parent) = file_path.parent() { @@ -106,7 +116,7 @@ impl Writer { } let mut writer = DiskWriter::try_new(file_path, rb.schema().as_ref(), range)?; writer.write(rb)?; - self.disk.insert(filename, writer); + self.disk.insert(filename.to_owned(), writer); Ok(()) } } @@ -139,6 +149,15 @@ impl Writer { mem.clear_unsegmented(); mem.segments .retain(|filename, _| self.disk.contains_key(filename)); + if mem.segments.is_empty() { + // clear/retain alone keep the burst's bucket allocation alive. + mem.segments = HashMap::new(); + } else if mem.segments.capacity() + > mem.segments.len().saturating_mul(2).saturating_add(16) + { + // Rehash only sparse maps; move segment owners without copying Arrow data. + mem.segments.shrink_to_fit(); + } } flushable_disk } @@ -258,14 +277,17 @@ impl Default for MemWriter { impl MemWriter { pub(crate) fn push_for_file( &mut self, - filename: String, + filename: &str, schema_key: &str, rb: &RecordBatch, ) -> Result<(), StagingError> { - self.segments - .entry(filename) - .or_default() - .push(schema_key, rb) + if let Some(segment) = self.segments.get_mut(filename) { + return segment.push(schema_key, rb); + } + let mut segment = Self::default(); + segment.push(schema_key, rb)?; + self.segments.insert(filename.to_owned(), segment); + Ok(()) } pub fn push(&mut self, schema_key: &str, rb: &RecordBatch) -> Result<(), StagingError> { @@ -282,7 +304,7 @@ impl MemWriter { } pub fn clear(&mut self) { - self.segments.clear(); + self.segments = HashMap::new(); self.clear_unsegmented(); } @@ -370,7 +392,74 @@ mod tests { use crate::{OBJECT_STORE_DATA_GRANULARITY, utils::time::TimeRange}; - use super::{DiskWriter, Writer}; + use super::{DiskWriter, MemWriter, Writer}; + + #[test] + fn current_minute_query_visibility_releases_burst_capacity() { + let dir = TempDir::new().unwrap(); + let now = chrono::DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") + .unwrap() + .with_timezone(&Utc); + let current = TimeRange::granularity_range(now, OBJECT_STORE_DATA_GRANULARITY); + let expired = TimeRange { + start: current.start - chrono::TimeDelta::minutes(2), + end: current.start, + }; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batches = (0..128) + .map(|id| { + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![id]))]) + .unwrap() + }) + .collect::>(); + let mut writer = Writer { + mem: Some(MemWriter::default()), + disk: Default::default(), + }; + for (id, batch) in batches.iter().enumerate() { + let filename = format!("{id}.date=2020-01-02.data.arrows"); + let range = if id < 2 { ¤t } else { &expired }; + writer + .push_disk_ref(&filename, batch, dir.path().join(&filename), range.clone()) + .unwrap(); + writer + .mem + .as_mut() + .unwrap() + .push_for_file(&filename, "id", batch) + .unwrap(); + } + let peak_capacity = writer.mem.as_ref().unwrap().segments.capacity(); + assert!(peak_capacity >= 128); + drop(writer.take_flushable_disk_at(false, now)); + let mem = writer.mem.as_ref().unwrap(); + assert_eq!(mem.segments.len(), 2); + assert!(mem.segments.capacity() <= 2 * mem.segments.len() + 16); + assert!(mem.segments.capacity() < peak_capacity); + for (id, batch) in batches.iter().take(2).enumerate() { + let key = format!("{id}.date=2020-01-02.data.arrows"); + let retained = &mem.segments[&key].mutable_buffer.inner; + assert_eq!(retained.len(), 1); + assert_eq!(retained[0].num_rows(), 1); + assert!(Arc::ptr_eq(retained[0].column(0), batch.column(0))); + } + assert_eq!( + mem.recordbatch_cloned(&schema) + .unwrap() + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); + drop(writer.take_flushable_disk_at(false, current.end)); + assert!(writer.disk.is_empty()); + assert_eq!(writer.mem.as_ref().unwrap().segments.capacity(), 0); + let mem = writer.mem.as_mut().unwrap(); + mem.push_for_file("clear-test", "id", &batches[0]).unwrap(); + assert!(mem.segments.capacity() > 0); + mem.clear(); + assert_eq!(mem.segments.capacity(), 0); + } #[test] fn current_minute_query_visibility_segments_follow_disk_ownership() { @@ -400,7 +489,10 @@ mod tests { ], ) .unwrap(); - let mut writer = Writer::default(); + let mut writer = Writer { + mem: Some(MemWriter::default()), + disk: Default::default(), + }; let current_range = TimeRange { start: now - chrono::TimeDelta::hours(1), end: now + chrono::TimeDelta::hours(1), @@ -435,22 +527,11 @@ mod tests { .mem .as_mut() .unwrap() - .segments - .entry(filename) - .or_default() - .push(key, rb) + .push_for_file(&filename, key, rb) .unwrap(); } - // Query-mode rows have no disk owner, so retain the existing flush behavior. - writer - .mem - .as_mut() - .unwrap() - .segments - .entry("query-only".to_owned()) - .or_default() - .push("id", &batch) - .unwrap(); + // Query-mode user rows have no disk owner and use the public direct buffer. + writer.mem.as_mut().unwrap().push("id", &batch).unwrap(); let key = "current-a.date=2020-01-02.data.arrows"; assert!( Arc::ptr_eq( @@ -503,6 +584,7 @@ mod tests { drop(writer.take_flushable_disk_at(true, now)); assert!(writer.disk.is_empty()); assert!(writer.mem.as_ref().unwrap().segments.is_empty()); + assert_eq!(writer.mem.as_ref().unwrap().segments.capacity(), 0); } #[test] diff --git a/src/parseable/streams.rs b/src/parseable/streams.rs index 6569ba06f..495946cc9 100644 --- a/src/parseable/streams.rs +++ b/src/parseable/streams.rs @@ -724,20 +724,21 @@ impl Stream { } } }; - let filename = - self.filename_by_partition(schema_key, parsed_timestamp, custom_partition_values); if self.options.mode != Mode::Query || stream_type == StreamType::Internal { + let filename = + self.filename_by_partition(schema_key, parsed_timestamp, custom_partition_values); let range = TimeRange::granularity_range( parsed_timestamp.and_local_timezone(Utc).unwrap(), OBJECT_STORE_DATA_GRANULARITY, ); let file_path = self.data_path.join(&filename); - guard.push_disk(filename.clone(), record, file_path, range)?; - } - - if let Some(mem) = guard.mem.as_mut() { - mem.push_for_file(filename, schema_key, record)?; + guard.push_disk_ref(&filename, record, file_path, range)?; + if let Some(mem) = guard.mem.as_mut() { + mem.push_for_file(&filename, schema_key, record)?; + } + } else if let Some(mem) = guard.mem.as_mut() { + mem.push(schema_key, record)?; } Ok(()) @@ -2986,6 +2987,8 @@ mod tests { None, &None, ); + // Keep memory coverage independent of the process-wide staging setting. + staging.writer.lock().unwrap().mem = Some(Default::default()); let now = DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") .unwrap() .with_timezone(&Utc); @@ -3001,7 +3004,17 @@ mod tests { StreamType::UserDefined, ) .unwrap(); - assert!(staging.writer.lock().unwrap().disk.is_empty()); + { + let writer = staging.writer.lock().unwrap(); + assert!(writer.disk.is_empty()); + let mem = writer.mem.as_ref().unwrap(); + assert!(mem.read_buffer.is_empty()); + assert_eq!(mem.mutable_buffer.inner.len(), 1); + assert!(Arc::ptr_eq( + mem.mutable_buffer.inner[0].column(0), + batch.column(0) + )); + } assert_eq!( staging.recordbatches_cloned(&schema).unwrap()[0].num_rows(), 1 @@ -3010,6 +3023,81 @@ mod tests { assert!(staging.recordbatches_cloned(&schema).unwrap().is_empty()); } + #[test] + fn current_minute_query_visibility_query_internal_exact_minute_end() { + let dir = TempDir::new().unwrap(); + let staging = Stream::new( + Arc::new(Options { + local_staging_path: dir.path().to_path_buf(), + mode: Mode::Query, + ..Default::default() + }), + "query-internal", + LogStreamMetadata { + stream_type: StreamType::Internal, + ..Default::default() + }, + None, + &None, + ); + staging.writer.lock().unwrap().mem = Some(Default::default()); + let now = DateTime::parse_from_rfc3339("2020-01-02T12:34:30.500Z") + .unwrap() + .with_timezone(&Utc); + let end = TimeRange::granularity_range(now, OBJECT_STORE_DATA_GRANULARITY).end; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![7]))]) + .unwrap(); + for _ in 0..2 { + staging + .push( + "id-schema", + &batch, + now.naive_utc(), + &HashMap::new(), + StreamType::Internal, + ) + .unwrap(); + } + staging.flush_with_clock(false, || now).unwrap(); + { + let writer = staging.writer.lock().unwrap(); + assert_eq!(writer.disk.len(), 1); + let mem = writer.mem.as_ref().unwrap(); + assert!(mem.read_buffer.is_empty()); + assert!(mem.mutable_buffer.inner.is_empty()); + } + assert_eq!( + staging.recordbatches_cloned(&schema).unwrap()[0].num_rows(), + 2 + ); + let next = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![8]))]) + .unwrap(); + staging + .push( + "id-schema", + &next, + end.naive_utc(), + &HashMap::new(), + StreamType::Internal, + ) + .unwrap(); + assert_eq!(staging.writer.lock().unwrap().disk.len(), 2); + staging.flush_with_clock(false, || end).unwrap(); + let expected = staging.filename_by_partition("id-schema", end.naive_utc(), &HashMap::new()); + { + let writer = staging.writer.lock().unwrap(); + assert_eq!(writer.disk.len(), 1); + assert!(writer.disk.contains_key(&expected)); + } + let retained = staging.recordbatches_cloned(&schema).unwrap(); + assert_eq!(retained.len(), 1); + assert_eq!(retained[0], next); + staging.flush_with_clock(true, || end).unwrap(); + assert!(staging.writer.lock().unwrap().disk.is_empty()); + assert!(staging.recordbatches_cloned(&schema).unwrap().is_empty()); + } + #[test] fn current_minute_query_visibility_survives_nonforced_flush() { // A fixed reference clock tests the production push/flush/query path without @@ -3028,6 +3116,7 @@ mod tests { None, &None, ); + staging.writer.lock().unwrap().mem = Some(Default::default()); let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); for (id, timestamp, partition) in [ (1, now - TimeDelta::minutes(2), "past"),