diff --git a/src/datafusion-local/src/tests/batch_size_alignment.rs b/src/datafusion-local/src/tests/batch_size_alignment.rs new file mode 100644 index 00000000..5c2e29b3 --- /dev/null +++ b/src/datafusion-local/src/tests/batch_size_alignment.rs @@ -0,0 +1,253 @@ +//! Regression tests for batch-size alignment: the reader indexes the cache by batch +//! id, and the parquet fallback turns that id back into rows using the *cache* batch +//! size. Reading at the session batch size (`datafusion.execution.batch_size`) +//! instead addresses the wrong rows whenever the two differ — the scan either runs +//! off the end of the row group and fails, or silently returns rows from the wrong +//! offsets. +//! +//! `LiquidCacheLocalBuilder::build` pins `execution.batch_size` to the cache batch +//! size, which is why the default configuration never hit this. A host that sets a +//! per-query batch size afterwards — as `SET datafusion.execution.batch_size` does +//! here — used to break the alignment. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Array, Int64Array}; +use arrow::record_batch::RecordBatch; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::error::Result; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::WriterProperties; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// One row group of `rows` rows, `id` ascending from 0. +fn write_single_row_group(path: &Path, rows: i64) { + write_parquet(path, rows, None); +} + +/// `rows` rows of ascending `id`, optionally split into row groups of +/// `row_group_rows`. +fn write_parquet(path: &Path, rows: i64, row_group_rows: Option) { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from((0..rows).collect::>()))], + ) + .unwrap(); + let mut props = WriterProperties::builder(); + if let Some(n) = row_group_rows { + props = props.set_max_row_group_row_count(Some(n)); + } + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, Some(props.build())).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); +} + +/// Cache batch size stays at the 8192 default; the session batch size is then set +/// to `session_batch_size`, as a host sizing batches by row width would do. +async fn ctx_with_session_batch_size( + cache_dir: &Path, + parquet_path: &Path, + session_batch_size: usize, +) -> Result { + ctx_with_session_batch_size_and_partitions(cache_dir, parquet_path, session_batch_size, 1).await +} + +async fn ctx_with_session_batch_size_and_partitions( + cache_dir: &Path, + parquet_path: &Path, + session_batch_size: usize, + target_partitions: usize, +) -> Result { + std::fs::create_dir_all(cache_dir)?; + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = target_partitions; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir.to_path_buf()) + .build(config) + .await?; + assert_eq!(cache.batch_size(), 8192, "cache batch size should be 8192"); + + ctx.sql(&format!( + "SET datafusion.execution.batch_size = {session_batch_size}" + )) + .await? + .collect() + .await?; + + ctx.register_parquet( + "t", + parquet_path.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await?; + Ok(ctx) +} + +async fn collect_batches(ctx: &SessionContext, sql: &str) -> Vec { + ctx.sql(sql).await.unwrap().collect().await.unwrap() +} + +fn ids_of(batches: &[RecordBatch]) -> Vec { + let mut ids = Vec::new(); + for batch in batches { + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + ids.extend((0..col.len()).map(|i| col.value(i))); + } + ids +} + +async fn collect_ids(ctx: &SessionContext, sql: &str) -> Vec { + ids_of(&collect_batches(ctx, sql).await) +} + +/// The reported failure: a 12000-row row group (fewer than `2 * 8192` rows) read +/// with a session batch size of 2145 used to drive the fallback past the end of +/// the row group on query batch 2, panicking with "parquet fallback ended before +/// batch 2". +#[tokio::test] +async fn narrowed_session_batch_size_reads_the_whole_row_group() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..12000).collect::>()); + } +} + +/// Silent corruption, not an error: a LIMIT stopped the scan before the fallback +/// ran out of rows, so the query succeeded and returned rows 0-2144 followed by +/// rows 8192-10336 in place of rows 2145-4289. +#[tokio::test] +async fn narrowed_session_batch_size_with_limit_returns_aligned_rows() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + let ids = collect_ids(&ctx, "SELECT id FROM t LIMIT 4290").await; + assert_eq!(ids, (0..4290).collect::>()); +} + +/// The fully-cached path misaligns as well. The cold run +/// populates cache batches 0 and 1; the warm run reads only from the cache and +/// used to emit rows 8192.. for query batch 1, because `read_arrow_array` applied +/// the 2145-bit selection to the 8192-row stored chunk and `arrow::compute::filter` +/// only rejects a mask *longer* than its target. +#[tokio::test] +async fn narrowed_session_batch_size_stays_aligned_on_the_fully_cached_path() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + let sql = "SELECT id FROM t LIMIT 4290"; + let _cold = collect_ids(&ctx, sql).await; + let warm = collect_ids(&ctx, sql).await; + assert_eq!(warm, (0..4290).collect::>()); +} + +/// A session batch size *larger* than the cache batch size misaligned in the other +/// direction: the selection mask outran the stored chunk, and +/// `arrow::compute::filter` rejected it with "Filter predicate of length .. is +/// larger than target array of length ..". +#[tokio::test] +async fn widened_session_batch_size_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 20000) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..12000).collect::>()); + } +} + +/// The scan now reads in cache-sized batches, but the caller must still receive +/// session-sized ones: `DataSourceExec::execute` wraps every source stream in a +/// `BatchSplitStream` sized by the session config. This is what makes ignoring the +/// session batch size inside the reader safe. +#[tokio::test] +async fn session_batch_size_still_bounds_the_batches_the_caller_receives() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_single_row_group(&parquet_path, 12000); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let batches = collect_batches(&ctx, "SELECT id FROM t").await; + let oversized: Vec<_> = batches + .iter() + .map(|b| b.num_rows()) + .filter(|rows| *rows > 2145) + .collect(); + assert!( + oversized.is_empty(), + "batches exceeded the session batch size: {oversized:?}" + ); + assert_eq!(ids_of(&batches), (0..12000).collect::>()); + } +} + +/// `current_batch_id` restarts at 0 for each row group, so a file with several of +/// them exercises the mapping repeatedly rather than once. +#[tokio::test] +async fn multi_row_group_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000)); + let ctx = ctx_with_session_batch_size(&tmp.path().join("cache"), &parquet_path, 2145) + .await + .unwrap(); + + for _ in 0..2 { + let ids = collect_ids(&ctx, "SELECT id FROM t").await; + assert_eq!(ids, (0..30000).collect::>()); + } +} + +/// Several partitions read the same file concurrently, each with its own reader +/// and its own `current_batch_id` sequence. +#[tokio::test] +async fn multi_partition_scan_stays_aligned() { + let tmp = TempDir::new().unwrap(); + let parquet_path = tmp.path().join("t.parquet"); + write_parquet(&parquet_path, 30000, Some(5000)); + let ctx = ctx_with_session_batch_size_and_partitions( + &tmp.path().join("cache"), + &parquet_path, + 2145, + 4, + ) + .await + .unwrap(); + + for _ in 0..2 { + let mut ids = collect_ids(&ctx, "SELECT id FROM t").await; + ids.sort_unstable(); + assert_eq!(ids, (0..30000).collect::>()); + } +} diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index a08b76c5..3e6a90cd 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -19,6 +19,7 @@ use datafusion::{ }; use crate::LiquidCacheLocalBuilder; +mod batch_size_alignment; mod constant_conjunct; mod date_optimizer; mod filter_limit; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 7e75f1b0..6e9ee546 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -183,7 +183,6 @@ impl Morselizer for LiquidMorselizer { metrics, file_pruner, reader_factory, - batch_size: self.batch_size, logical_file_schema, output_schema, projection, @@ -234,7 +233,6 @@ struct PreparedLiquidOpen { metrics: LiquidFileMetrics, file_pruner: Option, reader_factory: Arc, - batch_size: usize, logical_file_schema: SchemaRef, output_schema: SchemaRef, projection: ProjectionExprs, @@ -616,7 +614,6 @@ fn plan_row_group_morsels(planned: PlannedRowGroups) -> Result, pub(crate) cached_row_group: CachedRowGroupRef, @@ -92,13 +91,17 @@ pub(crate) struct ParquetFallback { impl LiquidCacheReader { pub(crate) fn new(config: LiquidCacheReaderConfig) -> Self { - debug_assert_eq!( - config.batch_size, - config.cached_row_group.batch_size(), - "DataFusion and LiquidCache batch sizes must agree" - ); + // The selection is walked in cache-sized windows, not session-sized + // ones. `current_batch_id` indexes stored chunks, and both the cache + // read and the parquet fallback turn that id back into rows with the + // cache batch size; windowing at `datafusion.execution.batch_size` + // would address different rows than the id names whenever a caller + // sets it to anything other than the cache's own size. The caller is + // unaffected either way: `DataSourceExec` re-splits every source + // stream to the session batch size. + let batch_size = config.cached_row_group.batch_size(); let inner = LiquidCacheReaderInner::new( - config.batch_size, + batch_size, config.selection, config.cached_row_group, config.projection_columns, @@ -526,7 +529,6 @@ mod tests { use std::sync::Arc; struct TestRowGroup { - batch_size: usize, row_group: CachedRowGroupRef, schema: SchemaRef, fallback: ParquetFallbackConfig, @@ -543,7 +545,6 @@ mod tests { impl TestRowGroup { fn reader(&self, request: ReaderRequest) -> LiquidCacheReader { LiquidCacheReader::new(LiquidCacheReaderConfig { - batch_size: self.batch_size, selection: request.selection, row_filter: request.row_filter, cached_row_group: Arc::clone(&self.row_group), @@ -622,7 +623,6 @@ mod tests { } TestRowGroup { - batch_size, row_group, schema, fallback: ParquetFallbackConfig { diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs index 271f53ff..035a0a4e 100644 --- a/src/datafusion/src/reader/runtime/morsel.rs +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -31,7 +31,6 @@ pub(crate) struct LiquidRowGroupPlanner { pub(crate) row_filter: Option, pub(crate) cached_file: CachedFileRef, pub(crate) projection: ProjectionMask, - pub(crate) batch_size: usize, pub(crate) stream_schema: SchemaRef, pub(crate) output_schema: SchemaRef, pub(crate) projector: Arc, @@ -151,7 +150,6 @@ impl LiquidRowGroupPlanner { Some(LiquidRowGroupMorsel { config: LiquidCacheReaderConfig { - batch_size: self.batch_size, selection, row_filter: self.row_filter.clone(), cached_row_group,