diff --git a/src/common/datasource/src/file_format/parquet/packed_reader.rs b/src/common/datasource/src/file_format/parquet/packed_reader.rs index 936f87d784c..9136b1ff940 100644 --- a/src/common/datasource/src/file_format/parquet/packed_reader.rs +++ b/src/common/datasource/src/file_format/parquet/packed_reader.rs @@ -26,7 +26,7 @@ use parquet::errors::{ParquetError, Result}; use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; use tokio::sync::Mutex; -const WINDOW_SIZE: usize = 8 * 1024 * 1024; +pub const WINDOW_SIZE: usize = 8 * 1024 * 1024; const WINDOW_BUDGET: usize = 2 * WINDOW_SIZE; struct Window { @@ -70,8 +70,12 @@ impl PackReadWindows { } let required = usize::try_from(range.end - range.start) .map_err(|_| ParquetError::General("packed read length overflow".into()))?; - let fetch_len = - (object_length - range.start).min(required.max(WINDOW_SIZE) as u64) as usize; + if required > WINDOW_SIZE { + return Err(ParquetError::General( + "packed cache range exceeds window".into(), + )); + } + let fetch_len = (object_length - range.start).min(WINDOW_SIZE as u64) as usize; // Evict only unpinned allocations. Bytes slices held by a decoder still // count against the request budget, including their entire backing window. let mut retained: usize = windows.iter().map(|w| w.bytes.len()).sum(); @@ -112,13 +116,14 @@ impl PackReadWindows { } } -/// Translates offsets relative to one complete Parquet stream into object ranges. +/// Pins one complete small stream; larger streams use the streaming reader. pub struct PackedParquetReader { windows: Arc, path: String, object_length: u64, offset: u64, length: u64, + bytes: Option, } impl PackedParquetReader { @@ -130,6 +135,7 @@ impl PackedParquetReader { length: u64, ) -> Result { if length < 12 + || length > WINDOW_SIZE as u64 || offset .checked_add(length) .is_none_or(|end| end > object_length) @@ -142,6 +148,7 @@ impl PackedParquetReader { object_length, offset, length, + bytes: None, }) } } @@ -154,13 +161,22 @@ impl AsyncFileReader for PackedParquetReader { "read outside indexed Parquet stream".into(), )); } - self.windows - .read( - &self.path, - self.object_length, - self.offset + range.start..self.offset + range.end, - ) - .await + if self.bytes.is_none() { + self.bytes = Some( + self.windows + .read( + &self.path, + self.object_length, + self.offset..self.offset + self.length, + ) + .await?, + ); + } + Ok(self + .bytes + .as_ref() + .unwrap() + .slice(range.start as usize..range.end as usize)) }) } @@ -256,46 +272,69 @@ mod tests { } #[tokio::test] - async fn pinned_allocations_bound_fetches_and_release_for_reuse() { - let dir = common_test_util::temp_dir::create_temp_dir("packed-budget"); + async fn small_stream_pins_one_window_for_multiple_columns() { + use parquet::file::properties::WriterProperties; + let schema = Arc::new(Schema::new( + (0..3) + .map(|i| Field::new(format!("v{i}"), arrow_schema::DataType::Utf8, false)) + .collect::>(), + )); + let value = "x".repeat(1024 * 1024); + let arrays = (0..3) + .map(|_| Arc::new(StringArray::from(vec![value.as_str()])) as ArrayRef) + .collect(); + let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap(); + let mut writer = ArrowWriter::try_new( + Vec::new(), + schema, + Some( + WriterProperties::builder() + .set_dictionary_enabled(false) + .build(), + ), + ) + .unwrap(); + writer.write(&batch).unwrap(); + let bytes = writer.into_inner().unwrap(); + assert!(bytes.len() < WINDOW_SIZE); + let dir = common_test_util::temp_dir::create_temp_dir("packed-columns"); let store = crate::test_util::test_store(dir.path().to_str().unwrap()); - let length = (3 * WINDOW_SIZE) as u64; - store - .write("pack-0.bin", vec![1u8; length as usize]) - .await - .unwrap(); + let mut object = vec![0; WINDOW_SIZE - 1024]; + let offset = object.len() as u64; + object.extend_from_slice(&bytes); + let length = object.len() as u64; + store.write("pack-0.bin", object).await.unwrap(); let windows = PackReadWindows::new(store); - let first = windows.read("pack-0.bin", length, 0..1).await.unwrap(); - let second = windows - .read( - "pack-0.bin", - length, - WINDOW_SIZE as u64..WINDOW_SIZE as u64 + 1, - ) + let pinned = windows.read("pack-0.bin", length, 0..1).await.unwrap(); + let mut reader = PackedParquetReader::new( + windows.clone(), + "pack-0.bin".into(), + length, + offset, + bytes.len() as u64, + ) + .unwrap(); + let metadata = reader.get_metadata(None).await.unwrap(); + let mut columns = Vec::new(); + for column in metadata.row_group(0).columns() { + let (start, len) = column.byte_range(); + columns.push(reader.get_bytes(start..start + len).await.unwrap()); + } + let batches: Vec<_> = ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .build() + .unwrap() + .try_collect() .await .unwrap(); - assert!( - windows - .read("pack-0.bin", length, 2 * WINDOW_SIZE as u64..length) - .await - .is_err() - ); - assert_eq!( - windows - .windows - .lock() - .await - .iter() - .map(|w| w.bytes.len()) - .sum::(), - WINDOW_BUDGET - ); - drop(first); - windows - .read("pack-0.bin", length, 2 * WINDOW_SIZE as u64..length) - .await - .unwrap(); - assert_eq!(second[0], 1); + assert_eq!(batches, vec![batch]); + let retained = windows.windows.lock().await; + assert_eq!(retained.len(), 2); + assert!(retained.iter().all(|w| w.bytes.len() <= WINDOW_SIZE)); + assert!(retained.iter().map(|w| w.bytes.len()).sum::() <= WINDOW_BUDGET); + assert_eq!(columns.len(), 3); + assert_eq!(pinned[0], 0); } #[tokio::test] diff --git a/src/operator/src/statement/import_packed.rs b/src/operator/src/statement/import_packed.rs index dde155f7045..7f31804c70b 100644 --- a/src/operator/src/statement/import_packed.rs +++ b/src/operator/src/statement/import_packed.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use common_datasource::file_format::Format; use common_datasource::file_format::parquet::packed_reader::{ - PackReadWindows, PackedParquetReader, + PackReadWindows, PackedParquetReader, WINDOW_SIZE, }; use common_datasource::object_store::{BuiltBackend, build_backend_with_path}; use common_datasource::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex}; @@ -185,7 +185,7 @@ impl StatementExecutor { .build() })?; match object.kind { - ObjectKind::Pack => { + ObjectKind::Pack if entry.length <= WINDOW_SIZE as u64 => { let reader = PackedParquetReader::new( Arc::clone(&windows), object.path.clone(), @@ -204,7 +204,9 @@ impl StatementExecutor { ) .await? } - ObjectKind::Parquet => { + ObjectKind::Pack | ObjectKind::Parquet => { + // The virtual file keeps footer/column offsets relative to this + // stream. Parquet decoder allocations are separate from the cache. let reader = plan .backend .object_store @@ -212,7 +214,7 @@ impl StatementExecutor { .chunk(256 * 1024) .await .context(error::ReadObjectSnafu { path: &object.path })? - .into_futures_async_read(0..object.length) + .into_futures_async_read(entry.offset..entry.offset + entry.length) .await .context(error::ReadObjectSnafu { path: &object.path })? .compat(); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index bd30d73d419..f98470b1450 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1142,3 +1142,96 @@ async fn packed_copy_standalone_heterogeneous_streams() { ); } } + +#[tokio::test(flavor = "multi_thread")] +async fn packed_copy_large_column_in_middle_stream() { + use common_datasource::packed_snapshot::{ObjectKind, PackIndex, PackObject, PackTable}; + use datafusion::parquet::arrow::ArrowWriter; + use datafusion::parquet::basic::Compression; + use datafusion::parquet::file::properties::WriterProperties; + use datatypes::arrow::array::{ + ArrayRef, BooleanArray, Int64Array, StringArray, TimestampMillisecondArray, + }; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + let standalone = GreptimeDbStandaloneBuilder::new("packed_large_column") + .build() + .await; + let instance = standalone.fe_instance(); + let directory = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let large = "x".repeat(20 * 1024 * 1024); + let arrays: [ArrayRef; 3] = [ + Arc::new(Int64Array::from(vec![42])), + Arc::new(StringArray::from(vec![large.as_str()])), + Arc::new(BooleanArray::from(vec![true])), + ]; + let mut packed = Vec::new(); + let mut index = PackIndex { + version: 1, + objects: vec![], + tables: vec![], + }; + for (i, (array, typ)) in arrays + .into_iter() + .zip(["BIGINT", "STRING", "BOOLEAN"]) + .enumerate() + { + sql( + instance, + &format!("CREATE TABLE restored_{i} (ts TIMESTAMP TIME INDEX, val {typ})"), + ) + .await; + let timestamps: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![1])); + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", timestamps.data_type().clone(), false), + Field::new("val", array.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![timestamps, array]).unwrap(); + let properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .build(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + if i == 1 { + assert!(metadata.row_group(0).column(1).compressed_size() > 16 * 1024 * 1024); + assert!(!packed.is_empty()); + } + index.tables.push(PackTable { + table_name: format!("restored_{i}"), + object: "pack-0.bin".into(), + offset: packed.len() as u64, + length: bytes.len() as u64, + row_count: 1, + }); + packed.extend(bytes); + } + index.objects.push(PackObject { + path: "pack-0.bin".into(), + kind: ObjectKind::Pack, + length: packed.len() as u64, + }); + std::fs::write(directory.path().join("pack-0.bin"), packed).unwrap(); + std::fs::write( + directory.path().join("pack-index.json"), + serde_json::to_vec(&index).unwrap(), + ) + .unwrap(); + let output = sql(instance, &format!("COPY DATABASE public FROM '{}/' WITH (FORMAT='parquet', metric_data_layout='packed', parallelism=2)", directory.path().display())).await; + assert!(matches!(output.data, OutputData::AffectedRows(3))); + for (i, expected) in [ + datatypes::value::Value::Int64(42), + datatypes::value::Value::String(large.into()), + datatypes::value::Value::Boolean(true), + ] + .into_iter() + .enumerate() + { + assert_eq!( + values(instance, &format!("SELECT val FROM restored_{i}")).await, + vec![vec![expected]] + ); + } +}