diff --git a/Cargo.lock b/Cargo.lock index 095d6d4acec..61a9c246ed1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8702,6 +8702,7 @@ dependencies = [ "api", "aquamarine", "arrow-schema 59.2.0", + "arrow-select 59.2.0", "async-channel 1.9.0", "async-stream", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 549412f369a..f47907aa797 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -107,6 +107,7 @@ arrow-buffer = "=59.2.0" arrow-flight = "=59.2.0" arrow-ipc = { version = "=59.2.0", default-features = false, features = ["lz4", "zstd"] } arrow-schema = { version = "=59.2.0", features = ["serde"] } +arrow-select = "=59.2.0" async-stream = "0.3" async-trait = "0.1" # Remember to update axum-extra, axum-macros when updating axum diff --git a/docs/rfcs/2026-09-11-metric-export-import.md b/docs/rfcs/2026-09-11-metric-export-import.md index 8ddc31f40b0..57212ea53cb 100644 --- a/docs/rfcs/2026-09-11-metric-export-import.md +++ b/docs/rfcs/2026-09-11-metric-export-import.md @@ -31,13 +31,21 @@ includes reports of exports taking days for roughly 100,000 logical tables with only several GiB of data. Table count, physical schema width and data volume must be measured separately; total database bytes do not describe the transfer cost. -The [experiment conclusions](2026-09-11-metric-export-import/experiment-results.md) -support shared scans, batch DDL, explicit-file COPY and packed-object transfer. -Packing with a bounded range reader substantially reduces storage requests and -improves full export and restore for many small logical tables. Benefits depend -on request latency, and memory use increases. Remaining CREATE work makes batch -DDL part of the first release. These local experiments establish direction; -production V2 recovery and release performance still require acceptance tests. +Local standalone experiments with development builds, local files and Docker +MinIO support shared scans, batch DDL, explicit-file COPY and packed-object +transfer. Packing with a bounded range reader reduces storage requests and +improves full export and restore for many small logical tables; low-latency +export gains are smaller. CPU and request overhead fell while peak process +memory increased. Remaining CREATE work makes batch DDL part of the first +release. Explicit-file COPY removes repeated directory traversal while retaining +ordinary per-table insertion, independently of merged database writes. + +The exercised fixtures preserved schemas, DDL and typed values. Standalone-file +experiments used the existing importer; packing used an independent driver and +prototype range reader. These experiments establish design direction, not +production packed-format compatibility, V2 recovery safety or release performance. +The 8 MiB write/part limit was a design choice, not a measured optimum. Current +acceptance status is maintained in the [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120). # Terminology @@ -552,18 +560,39 @@ batching, still writing standalone table objects; (C) B plus packing and its reader, with identical DDL handling; (D) C plus batch DDL. If general restore request deduplication is available, include it in B as well. -The proposed release target is at least 3x faster full export and full restore -for D versus A on 10,000 and 100,000 small logical tables with MinIO request -latency increased by 5 ms. Fix this target during RFC review, before running -release measurements. The independent local C-versus-B experiment establishes -neither this comparison nor proof of that target. +The release target is at least 3x faster full export and full restore for D +versus A on both 10,000 and 100,000 small logical tables at both +50 ms and ++100 ms added MinIO request latency. Populate every selected table with the same +deterministic data across controls, initially 32 rows per table. Include schema +export, metadata discovery and CLI startup/shutdown in complete wall duration. +The independent local C-versus-B experiment establishes neither this comparison +nor proof of that target. Keep production improvements already shared by all +controls in A, and identify which increments can actually be isolated. -Use Linux release builds and at least three alternating rounds for primary -comparisons. Report medians/spread, schema/data/full durations, CPU, RSS, bytes, -objects, request counts and queue/cache peaks. Include zero-added-latency, -local-file, mixed small/large-table and wide-schema controls. Investigate and -remeasure full-duration regressions exceeding 5% versus A; do not generalize the -latency-injected speedup to these cases. Check values and schemas outside timing. +Start with native optimized release builds on macOS. Record hardware, OS, +MinIO placement, container/VM CPU and memory limits, and connection/storage +settings. Linux deployment-side performance confirmation can follow separately; +Linux root/non-root and Windows correctness checks remain distinct requirements. +Earlier development-build diagnostics, runs with paging interference and sparse +fixtures with only one populated table do not establish this acceptance. + +Calibrate one added delay per object-store HTTP request, including HEAD, GET +(and range GET), PUT, LIST, DELETE and each multipart operation. Do not delay +individual packets, streaming body chunks or database SQL/DDL. Route CLI and +server object access through the same controlled endpoint and report measured +per-method latency distributions. Synthetic delay is not a cloud provider profile. + +Use a small calibration pilot before each scale and at least three alternating +A/B/C/D rounds for primary comparisons, with fresh restore namespaces, equal +concurrency and comparable cache/compaction conditions. Report medians/spread, +schema/data/full durations, CPU, RSS, swap activity, bytes, objects, requests and +queue/cache peaks. Check RAM, disk and MinIO health before large runs; do not +compile concurrently. Fix interfering paging/load or report the scale unverified. +Retain local-file and MinIO 0/+5 ms controls, mixed small/large tables, wide +physical column unions and multiple physical groups. Add +20 ms only if needed +to explain sensitivity. Investigate and remeasure full-duration regressions +exceeding 5% versus A. Check table membership, schemas and typed values outside +timing, and report B/A, C/B and D/C independently where comparable. Validate actual write/part sizes at or below 8 MiB, aggregate concurrency across groups/chunks, and resource release after success, cancellation and failure. diff --git a/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md b/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md deleted file mode 100644 index 9cc7b46464d..00000000000 --- a/docs/rfcs/2026-09-11-metric-export-import/experiment-results.md +++ /dev/null @@ -1,34 +0,0 @@ -# Metric export/import experiment conclusions - -The experiments support shared physical scans, explicit-file COPY, batched -logical DDL, and packing small logical-table Parquet streams into shared objects. - -- **Shared scans reduce repeated query work.** Projecting each logical schema - before dictionary expansion preserves its columns and types, but does not - remove the memory cost of a wide physical schema or distributed ordering. -- **Explicit-file COPY removes repeated directory traversal.** This improvement - uses ordinary per-table insertion; it is independent of database merged writes. -- **Packing and its range reader reduce small-object overhead.** They improve - complete export and restore and sharply reduce object-store requests. Benefits - grow with table count and request latency; low-latency export gains are smaller. - Some reader request savings can also benefit standalone files. -- **Batch DDL remains necessary.** Once data restore is faster, logical-table - creation becomes the main remaining restore cost. The batch-DDL experiment used - an internal interface; production transport and retry integration remain work. -- **Bounded writes allow larger objects.** The packing experiment kept writes - and upload parts within 8 MiB without imposing that limit on whole objects or - splitting logical tables. This was the chosen size, not a measured optimum. -- **Packing has a memory trade-off.** CPU and request overhead fell, while peak - process memory increased. Encoding buffers, upload concurrency, range windows - and table/index metadata need separate accounting. - -The exercised fixtures preserved schemas, DDL and typed values after restore. -Earlier standalone-file experiments used the existing importer; the packing -experiment used an independent driver and prototype range reader. The latter -does not establish production packed-format compatibility or V2 resume safety. - -These were local standalone experiments using development builds, local files -and Docker MinIO. Release acceptance still needs the actual packed V2 path, -request-owned reader isolation, old-reader rejection, large-table throughput, -distributed ordering, cancellation and recovery. See the [RFC](../2026-09-11-metric-export-import.md) -and [tracking issue](https://github.com/GreptimeTeam/greptimedb/issues/9120). diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index 17008706066..f8dd22f0307 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -21,6 +21,7 @@ use std::collections::BTreeSet; use std::path::Component; use async_trait::async_trait; +use common_datasource::packed_writer::WRITE_BYTES; use futures::TryStreamExt; use object_store::services::{Azblob, Fs, Gcs, Oss, S3}; use object_store::util::{with_instrument_layers, with_retry_layers}; @@ -570,13 +571,31 @@ impl OpenDalStorage { /// Writes bytes to a file. async fn write_file(&self, path: &str, data: Vec) -> Result<()> { - self.object_store - .write(path, data) + let mut writer = self + .object_store + .writer_with(path) + .chunk(WRITE_BYTES) + .concurrent(1) .await - .map(|_| ()) - .context(StorageOperationSnafu { - operation: format!("write {}", path), - }) + .with_context(|_| StorageOperationSnafu { + operation: format!("open writer {}", path), + })?; + let data = object_store::Buffer::from(data); + let result = async { + for chunk in data.chunks(WRITE_BYTES) { + writer.write(chunk).await?; + } + writer.close().await.map(|_| ()) + } + .await; + if result.is_err() + && let Err(error) = writer.abort().await + { + common_telemetry::warn!(error; "Failed to abort snapshot metadata write"); + } + result.with_context(|_| StorageOperationSnafu { + operation: format!("write {}", path), + }) } /// Checks if a file exists using stat. @@ -837,6 +856,96 @@ mod tests { ) } + #[tokio::test] + async fn snapshot_metadata_submissions_are_bounded_and_invoke_abort_on_failure() { + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use object_store::layers::mock::{Metadata, ObservedWriter, WriterObservations, oio}; + + struct FailWriter { + inner: oio::Writer, + writes: usize, + fail: Arc, + } + impl oio::Write for FailWriter { + async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { + self.writes += 1; + if self.fail.load(Ordering::SeqCst) == 1 && self.writes == 2 { + return Err(object_store::Error::new( + ErrorKind::Unexpected, + "injected metadata write failure", + )); + } + self.inner.write(bytes).await + } + async fn close(&mut self) -> object_store::Result { + if self.fail.load(Ordering::SeqCst) == 2 { + return Err(object_store::Error::new( + ErrorKind::Unexpected, + "injected metadata close failure", + )); + } + self.inner.close().await + } + async fn abort(&mut self) -> object_store::Result<()> { + self.inner.abort().await + } + } + + let dir = tempdir().unwrap(); + let mut storage = make_storage_with_rooted_fs(dir.path()); + let observations = Arc::new(Mutex::new(WriterObservations::default())); + let fail = Arc::new(AtomicUsize::new(0)); + storage.object_store = storage.object_store.layer( + MockLayerBuilder::default() + .writer_factory(Arc::new({ + let observations = observations.clone(); + let fail = fail.clone(); + move |_, _, inner| { + Box::new(ObservedWriter::new( + Box::new(FailWriter { + inner, + writes: 0, + fail: fail.clone(), + }), + observations.clone(), + )) + } + })) + .build() + .unwrap(), + ); + for length in [2 * WRITE_BYTES + 17, 1, 0] { + *observations.lock().unwrap() = WriterObservations::default(); + let data = vec![42; length]; + storage.write_file("metadata", data.clone()).await.unwrap(); + assert_eq!(storage.read_file("metadata").await.unwrap(), data); + let observed = observations.lock().unwrap(); + assert_eq!(observed.write_sizes.iter().sum::(), length); + assert!(observed.write_sizes.iter().all(|size| *size <= WRITE_BYTES)); + assert_eq!(observed.close_calls, 1); + assert_eq!(observed.abort_calls, 0); + } + for failure in [1, 2] { + *observations.lock().unwrap() = WriterObservations::default(); + fail.store(failure, Ordering::SeqCst); + assert!( + storage + .write_file("failed", vec![42; 2 * WRITE_BYTES + 17]) + .await + .is_err() + ); + assert_eq!( + observations.lock().unwrap().write_sizes.len(), + if failure == 1 { 2 } else { 3 } + ); + let observed = observations.lock().unwrap(); + assert_eq!(observed.close_calls, usize::from(failure == 2)); + assert_eq!(observed.abort_calls, 1); + } + } + #[test] fn test_validate_uri_valid() { assert_eq!(validate_uri("s3://bucket/path").unwrap(), StorageScheme::S3); diff --git a/src/common/datasource/src/packed_writer.rs b/src/common/datasource/src/packed_writer.rs index 635a1a7a771..d0560fa06f7 100644 --- a/src/common/datasource/src/packed_writer.rs +++ b/src/common/datasource/src/packed_writer.rs @@ -410,7 +410,9 @@ mod tests { use arrow::array::{ArrayRef, Int64Array, StringArray}; use arrow::record_batch::RecordBatch; use futures::TryStreamExt; - use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; + use object_store::layers::mock::{ + Metadata, MockLayerBuilder, MockWriterFactory, ObservedWriter, WriterObservations, oio, + }; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; @@ -424,10 +426,10 @@ mod tests { let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() .build_operator(); - let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let observations = Arc::new(std::sync::Mutex::new(WriterObservations::default())); let factory: MockWriterFactory = Arc::new({ - let sizes = sizes.clone(); - move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + let observations = observations.clone(); + move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone())) }); let store = store.layer( MockLayerBuilder::default() @@ -483,7 +485,14 @@ mod tests { .validate_membership(["table0", "table1", "table2"]) .unwrap(); assert_eq!(inventory.len(), 3); - assert!(sizes.lock().unwrap().iter().all(|n| *n <= WRITE_BYTES)); + assert!( + observations + .lock() + .unwrap() + .write_sizes + .iter() + .all(|n| *n <= WRITE_BYTES) + ); assert_eq!( index .objects @@ -540,27 +549,13 @@ mod tests { } } - struct Observe(oio::Writer, Arc>>); - impl oio::Write for Observe { - async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { - self.1.lock().unwrap().push(bytes.len()); - self.0.write(bytes).await - } - async fn close(&mut self) -> object_store::Result { - self.0.close().await - } - async fn abort(&mut self) -> object_store::Result<()> { - self.0.abort().await - } - } - #[tokio::test] async fn rolls_at_table_boundary_and_aborts_standalone_at_part_limit() { let directory = common_test_util::temp_dir::create_temp_dir("packed-limit"); - let sizes = Arc::new(std::sync::Mutex::new(Vec::new())); + let observations = Arc::new(std::sync::Mutex::new(WriterObservations::default())); let factory: MockWriterFactory = Arc::new({ - let sizes = sizes.clone(); - move |_, _, writer| Box::new(Observe(writer, sizes.clone())) + let observations = observations.clone(); + move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone())) }); let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() @@ -605,7 +600,14 @@ mod tests { .iter() .all(|t| t.offset == 0 && t.length == 600) ); - assert!(sizes.lock().unwrap().iter().all(|n| *n <= 512)); + assert!( + observations + .lock() + .unwrap() + .write_sizes + .iter() + .all(|n| *n <= 512) + ); } struct FailClose(oio::Writer); diff --git a/src/common/datasource/src/parquet_writer.rs b/src/common/datasource/src/parquet_writer.rs index 243a46b4e1a..157733e856e 100644 --- a/src/common/datasource/src/parquet_writer.rs +++ b/src/common/datasource/src/parquet_writer.rs @@ -500,28 +500,17 @@ mod tests { async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() { use std::sync::Mutex; - use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio}; - struct ObservedWriter(oio::Writer, Arc>>); - impl oio::Write for ObservedWriter { - async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> { - self.1.lock().unwrap().push(bytes.len()); - self.0.write(bytes).await - } - async fn close(&mut self) -> object_store::Result { - self.0.close().await - } - async fn abort(&mut self) -> object_store::Result<()> { - self.0.abort().await - } - } + use object_store::layers::mock::{ + MockLayerBuilder, MockWriterFactory, ObservedWriter, WriterObservations, + }; let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet"); let store = object_store::secure_fs::SecureFsRoot::open(directory.path()) .unwrap() .build_operator(); - let sizes = Arc::new(Mutex::new(Vec::new())); + let observations = Arc::new(Mutex::new(WriterObservations::default())); let factory: MockWriterFactory = Arc::new({ - let sizes = sizes.clone(); - move |_, _, writer| Box::new(ObservedWriter(writer, sizes.clone())) + let observations = observations.clone(); + move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone())) }); let store = store.layer( MockLayerBuilder::default() @@ -548,7 +537,7 @@ mod tests { writer.sink = ParquetSink::Object(store.writer("large.parquet").await.unwrap()); writer.write(batch.clone(), None).await.unwrap(); writer.finish(None).await.unwrap(); - let sizes = sizes.lock().unwrap().clone(); + let sizes = observations.lock().unwrap().write_sizes.clone(); let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize; assert!(sizes.iter().sum::() > limit); assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}"); diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index 39dd96f687b..02c134e964c 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -19,6 +19,7 @@ workspace = true api.workspace = true aquamarine.workspace = true arrow-schema.workspace = true +arrow-select.workspace = true async-channel = "1.9" async-stream.workspace = true async-trait.workspace = true diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index f86b327735d..f55f084ebd7 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -4498,6 +4498,143 @@ async fn test_range_cache_key_separates_sequence_ranges() { ); } +#[tokio::test] +async fn test_single_source_ranges_preserve_batches() { + for (file_count, row_group_size) in [(1, 64), (1, 16), (2, 16)] { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("max_row_group_row_count", &row_group_size.to_string()) + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows_per_file = 64 / file_count; + for file in 0..file_count { + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: test_util::build_rows(file * rows_per_file, (file + 1) * rows_per_file), + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + let Scanner::Seq(scanner) = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + else { + panic!("expected a sequential scan"); + }; + assert_eq!(file_count, scanner.input().files.len()); + assert_eq!( + 64 / row_group_size, + scanner.properties().partitions[0].len() + ); + let stream = scanner + .scan_partition(&Default::default(), &ExecutionPlanMetricsSet::default(), 0) + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + assert_eq!( + vec![row_group_size; 64 / row_group_size], + batches + .iter() + .map(|batch| batch.num_rows()) + .collect::>() + ); + let mut values = scan_field_values(&batches); + values.sort_by(f64::total_cmp); + assert_eq!((0..64).map(|i| i as f64).collect::>(), values); + } +} + +#[tokio::test] +async fn test_short_series_level_zero_merge_preserves_batches_and_latest_values() { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: true, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .insert_option("max_row_group_row_count", "65536") + .build(); + let schema = test_util::rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + for (start, end, offset) in [(0, 12_000, 0), (6_000, 18_000, 100_000)] { + let rows = (start..end) + .map(|i| { + row(vec![ + ValueData::StringValue(format!("{i:05}")), + ValueData::F64Value((i + offset) as f64), + ValueData::TimestampMillisecondValue(0), + ]) + }) + .collect(); + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows, + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + let Scanner::Seq(scanner) = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap() + else { + panic!("expected a sequential scan"); + }; + assert_eq!(2, scanner.input().files.len()); + for file in &scanner.input().files { + assert_eq!(0, file.meta_ref().level); + assert_eq!(12_000, file.meta_ref().num_series); + } + let batches = scanner + .scan_partition(&Default::default(), &ExecutionPlanMetricsSet::default(), 0) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!( + batches.len() < 64, + "short series became {} batches", + batches.len() + ); + assert_eq!( + (0..18_000) + .map(|i| if i < 6_000 { + i as f64 + } else { + (i + 100_000) as f64 + }) + .collect::>(), + scan_field_values(&batches), + ); +} + #[tokio::test] async fn test_reset_state_discards_dynamic_field_pruning() { for (append_mode, expected_scanner_name) in [(false, "SeqScan"), (true, "UnorderedScan")] { diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 2abf6aa48af..1734ca364af 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -22,7 +22,7 @@ use common_recordbatch::error::{ ArrowComputeSnafu, DataTypesSnafu, ExternalSnafu, NewDfRecordBatchSnafu, }; use common_recordbatch::{DfRecordBatch, RecordBatch}; -use datatypes::arrow::array::Array; +use datatypes::arrow::array::{Array, AsArray}; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; @@ -320,6 +320,16 @@ impl FlatProjectionMapper { let mut arrays = Vec::with_capacity(self.output_schema.num_columns()); for (output_idx, index) in self.batch_indices.iter().enumerate() { let mut array = batch.column(*index).clone(); + if matches!( + self.output_schema.arrow_schema().fields()[output_idx].data_type(), + ArrowDataType::Dictionary(_, _) + ) && let Some(dictionary) = array.as_any_dictionary_opt() + && dictionary.values().len() > dictionary.len() + { + // Small scan slices must not retain an entire SST's tag dictionary in merges. + array = arrow_select::dictionary::garbage_collect_any_dictionary(dictionary) + .context(ArrowComputeSnafu)?; + } // Cast dictionary values to the target type. if let ArrowDataType::Dictionary(_key_type, value_type) = array.data_type() && !matches!( @@ -667,6 +677,63 @@ mod tests { ); } + #[test] + fn test_tag_projection_compacts_dictionary_slices() { + use datatypes::arrow::array::{ + DictionaryArray, StringArray, TimestampMillisecondArray, UInt32Array, + }; + use datatypes::arrow::datatypes::UInt32Type; + + let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mapper = FlatProjectionMapper::new(&metadata, [0, 1]) + .unwrap() + .with_pk_dictionary_encoding(); + let values = StringArray::from_iter( + (0..1024).map(|i| (i != 1).then(|| format!("{i}{}", "x".repeat(4096)))), + ); + let dictionary = DictionaryArray::::try_new( + UInt32Array::from(vec![Some(1023), None, Some(1), Some(1023)]), + Arc::new(values), + ) + .unwrap(); + let expected = datatypes::arrow::compute::cast(&dictionary, &ArrowDataType::Utf8).unwrap(); + let batch = DfRecordBatch::try_from_iter([ + ("tag", Arc::new(dictionary) as Arc), + ( + "ts", + Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])), + ), + ]) + .unwrap(); + let output = mapper.convert(&batch, &CacheStrategy::Disabled).unwrap(); + let actual = output.column(0); + assert!(actual.get_array_memory_size() < 64 * 1024); + assert_eq!(actual.data_type(), batch.column(0).data_type()); + assert_eq!( + datatypes::arrow::compute::cast(actual.as_ref(), &ArrowDataType::Utf8) + .unwrap() + .as_ref(), + expected.as_ref(), + ); + } + #[test] fn test_single_value_dictionary_preserves_string_type() { use datatypes::arrow::array::{DictionaryArray, TimestampMillisecondArray, UInt32Array}; diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index c5eed8853ae..7224143f2c4 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/src/mito2/src/read/scan_util.rs @@ -1234,7 +1234,15 @@ pub(crate) fn should_split_flat_batches_for_merge( let file = &stream_ctx.input.files[file_index]; let file_meta = file.meta_ref(); if file_meta.level == 0 { - // Always split level 0 files. + // Short series can leave too many tiny batches retained by the merge. + if file_meta.num_series > 0 { + if !can_split_series(file_meta.num_rows, file_meta.num_series) { + return None; + } + total_rows += file_meta.num_rows; + total_series += file_meta.num_series; + } + // Older L0 files may lack series statistics; keep their fallback behavior. num_files_to_split += 1; continue; } else if file_meta.num_rows < SPLIT_ROW_THRESHOLD || file_meta.num_series == 0 { @@ -1913,6 +1921,30 @@ mod tests { ); } + #[tokio::test] + async fn test_should_split_flat_batches_for_merge_uses_level_zero_series_stats() { + for (num_rows, num_series, expected) in [ + (70_144, 35_083, None), + (100_000, 1_000, Some(100)), + (100_000, 0, Some(DEFAULT_READ_BATCH_SIZE / 4)), + ] { + let mut meta = new_test_file(num_rows, num_series).meta_ref().clone(); + meta.level = 0; + let file = FileHandle::new(meta, Arc::new(NoopFilePurger)); + let stream_ctx = new_test_stream_ctx(vec![file], vec![]).await; + let range_meta = new_test_range_meta(smallvec![RowGroupIndex { + index: 0, + row_group_index: 0, + }]); + + assert_eq!( + expected, + should_split_flat_batches_for_merge(&stream_ctx, &range_meta), + "rows={num_rows}, series={num_series}", + ); + } + } + #[tokio::test] async fn test_should_split_flat_batches_for_merge_falls_back_to_memtables() { let stream_ctx = new_test_stream_ctx(vec![], vec![new_test_memtable(5_000, 100)]).await; diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 323704c31e2..bbf6ff1bf92 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -662,8 +662,6 @@ pub(crate) async fn build_flat_sources( return Ok(None); } - let split_batch_size = should_split_flat_batches_for_merge(stream_ctx, range_meta); - let should_split = split_batch_size.is_some(); sources.reserve(num_indices); let mut ordered_sources = Vec::with_capacity(num_indices); ordered_sources.resize_with(num_indices, || None); @@ -737,6 +735,16 @@ pub(crate) async fn build_flat_sources( } } + // Queries merge one range at a time; compaction merges sources from all ranges. + let needs_merge = + ordered_sources.iter().flatten().count() > 1 || (compaction && stream_ctx.ranges.len() > 1); + let split_batch_size = if needs_merge { + should_split_flat_batches_for_merge(stream_ctx, range_meta) + } else { + None + }; + let should_split = split_batch_size.is_some(); + for stream in ordered_sources.into_iter().flatten() { if should_split { sources.push(Box::pin(SplitRecordBatchStream::new(stream))); diff --git a/src/object-store/src/layers/mock.rs b/src/object-store/src/layers/mock.rs index 270f776f935..4807a401c95 100644 --- a/src/object-store/src/layers/mock.rs +++ b/src/object-store/src/layers/mock.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::fmt::Debug; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use derive_builder::Builder; pub use oio::*; @@ -225,6 +225,50 @@ impl oio::Read for MockReader { } } +/// Writer calls observed before forwarding, including calls that fail. +#[derive(Debug, Default)] +pub struct WriterObservations { + pub write_sizes: Vec, + pub close_calls: usize, + pub abort_calls: usize, +} + +/// Records calls without changing the inner writer's completion or abort semantics. +pub struct ObservedWriter { + inner: oio::Writer, + observations: Arc>, +} + +impl ObservedWriter { + pub fn new(inner: oio::Writer, observations: Arc>) -> Self { + Self { + inner, + observations, + } + } +} + +impl oio::Write for ObservedWriter { + async fn write(&mut self, bytes: Buffer) -> Result<()> { + self.observations + .lock() + .unwrap() + .write_sizes + .push(bytes.len()); + self.inner.write(bytes).await + } + + async fn close(&mut self) -> Result { + self.observations.lock().unwrap().close_calls += 1; + self.inner.close().await + } + + async fn abort(&mut self) -> Result<()> { + self.observations.lock().unwrap().abort_calls += 1; + self.inner.abort().await + } +} + pub struct MockWriter { inner: oio::Writer, } diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 368335f8450..8789e633995 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -41,7 +41,7 @@ use common_recordbatch::SendableRecordBatchStream; use common_time::range::TimestampRange; use datafusion::datasource::DefaultTableSource; use datafusion_common::TableReference as DfTableReference; -use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col}; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; use futures::StreamExt; use object_store::ObjectStore; use session::context::QueryContextRef; @@ -53,6 +53,7 @@ use store_api::metric_engine_consts::{ use table::TableRef; use table::metadata::TableId; use table::table::adapter::DfTableProviderAdapter; +use tokio::sync::OwnedSemaphorePermit; use tokio_util::sync::CancellationToken; use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; @@ -79,7 +80,7 @@ impl Default for LogicalTableExportLimits { input_batch_bytes: 64 * 1024 * 1024, conversion_bytes: 1024 * 1024, writer: ParquetWriterLimits { - row_group_rows: 8192, + row_group_rows: 1024 * 1024, flush_threshold_bytes: 8 * 1024 * 1024, max_row_groups: 4096, }, @@ -255,7 +256,7 @@ impl LogicalTableExport { fn build_plan(&self, time_range: Option<&TimestampRange>) -> Result { let info = self.physical_table.table_info(); - let filters = self + let mut filters = self .physical_table .schema() .timestamp_column() @@ -264,6 +265,11 @@ impl LogicalTableExport { }) .into_iter() .collect::>(); + if self.logical_tables.len() == 1 + && let Some((&id, _)) = self.logical_tables.first_key_value() + { + filters.push(col(TABLE_ID).eq(lit(id))); + } let source = Arc::new(DefaultTableSource::new(Arc::new( DfTableProviderAdapter::new(self.physical_table.clone()), ))); @@ -283,11 +289,13 @@ impl LogicalTableExport { .filter(filter) .context(error::BuildDfLogicalPlanSnafu)?; } - builder - .sort(vec![col(TABLE_ID).sort(true, false)]) - .context(error::BuildDfLogicalPlanSnafu)? - .build() - .context(error::BuildDfLogicalPlanSnafu) + // A single selected table has a constant routing key after filtering. + if self.logical_tables.len() > 1 { + builder = builder + .sort(vec![col(TABLE_ID).sort(true, false)]) + .context(error::BuildDfLogicalPlanSnafu)?; + } + builder.build().context(error::BuildDfLogicalPlanSnafu) } } @@ -613,6 +621,46 @@ async fn expand_bounded_slice( batch.num_columns(), requested, )?; + // Parquet accepts compatible string dictionaries without expanding their values. + let fields = schema + .fields() + .iter() + .zip(batch.columns()) + .map(|(field, array)| { + if let DataType::Dictionary(_, value) = array.data_type() + && value.as_ref() == field.data_type() + && matches!(value.as_ref(), DataType::Utf8 | DataType::LargeUtf8) + { + Arc::new( + field + .as_ref() + .clone() + .with_data_type(array.data_type().clone()), + ) + } else { + field.clone() + } + }) + .collect::>(); + let schema = Arc::new(arrow::datatypes::Schema::new_with_metadata( + fields, + schema.metadata().clone(), + )); + // Sharing compatible arrays needs no CPU conversion. Keep the size check on + // this task only when its dictionary inspection is bounded. + if start < end + && batch + .columns() + .iter() + .zip(schema.fields()) + .all(|(array, field)| array.data_type() == field.data_type()) + && let Some(estimated) = estimate_batch_bytes_upper_bound(&batch, end - start, &[])? + && estimated <= conversion + { + let reservation = retained.saturating_add(estimated.saturating_mul(4)); + let permit = budget.reserve(reservation, cancellation).await?; + return convert_slice_payload(batch, schema, start, end - start, reservation, permit); + } let input = batch.clone(); let (len, estimated) = common_runtime::spawn_blocking_global(move || { rows_within_budget(&input, start, end, conversion, &[]) @@ -622,26 +670,36 @@ async fn expand_bounded_slice( let reservation = retained.saturating_add(estimated.saturating_mul(4)); let permit = budget.reserve(reservation, cancellation).await?; common_runtime::spawn_blocking_global(move || { - let slice = batch.slice(start, len); - let expanded = expand_export_batch(&slice, schema)?; - ensure!( - expanded.get_array_memory_size() <= reservation, - LogicalTableExportResourceSnafu { - reason: "converted backing buffers exceed reservation" - } - ); - Ok(( - Payload { - batch: expanded, - permit, - }, - len, - )) + convert_slice_payload(batch, schema, start, len, reservation, permit) }) .await .context(error::JoinTaskSnafu)? } +fn convert_slice_payload( + batch: RecordBatch, + schema: SchemaRef, + start: usize, + len: usize, + reservation: usize, + permit: OwnedSemaphorePermit, +) -> Result<(Payload, usize)> { + let expanded = expand_export_batch(&batch.slice(start, len), schema)?; + ensure!( + expanded.get_array_memory_size() <= reservation, + LogicalTableExportResourceSnafu { + reason: "converted backing buffers exceed reservation" + } + ); + Ok(( + Payload { + batch: expanded, + permit, + }, + len, + )) +} + /// Limits nested casts to selected children: Arrow's List cast otherwise expands /// the entire values array, even when the parent has been sliced. pub(crate) fn expand_export_batch(batch: &RecordBatch, schema: SchemaRef) -> Result { @@ -788,6 +846,41 @@ fn estimate_value_size(array: &dyn Array, row: usize) -> Result { Ok(bytes.saturating_add(16)) } +/// Estimates a conservative byte upper bound with bounded work on an async task. +/// `None` means this path cannot estimate the batch, not that it exceeds a budget. +fn estimate_batch_bytes_upper_bound( + batch: &RecordBatch, + rows: usize, + json_columns: &[usize], +) -> Result> { + let mut dictionary_values = 1024usize; + let mut max_row_bytes = 0usize; + for (index, array) in batch.columns().iter().enumerate() { + let size = if let Some(width) = array.data_type().primitive_width() { + (width + 16).max(32) + } else if let Some(dictionary) = array.as_any_dictionary_opt() + && dictionary.values().len() < rows + && dictionary.values().len() <= dictionary_values + && matches!( + dictionary.values().data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) + { + let values = dictionary.values(); + dictionary_values -= values.len(); + let max = (0..values.len()).try_fold(0usize, |max, row| { + Ok::<_, error::Error>(max.max(estimate_value_size(values.as_ref(), row)?)) + })?; + max.saturating_add(16).max(32) + } else { + return Ok(None); + }; + let expansion = if json_columns.contains(&index) { 8 } else { 1 }; + max_row_bytes = max_row_bytes.saturating_add(size.saturating_mul(expansion)); + } + Ok(max_row_bytes.checked_mul(rows)) +} + pub(crate) fn rows_within_budget( batch: &RecordBatch, start: usize, @@ -795,6 +888,12 @@ pub(crate) fn rows_within_budget( budget: usize, json_columns: &[usize], ) -> Result<(usize, usize)> { + if start < end + && let Some(bytes) = estimate_batch_bytes_upper_bound(batch, end - start, json_columns)? + && bytes <= budget + { + return Ok((end - start, bytes)); + } let mut bytes = 0usize; let mut row = start; while row < end { diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 7aa6daab127..6513d7e43fd 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -127,6 +127,188 @@ async fn read(store: &ObjectStore, path: &str) -> (SchemaRef, Vec) ) } +#[tokio::test] +async fn single_table_plan_filters_other_routing_ids_before_export() { + let mut unit = unit(); + unit.logical_tables.retain(|id, _| *id == 1025); + let schema = unit.physical_table.schema(); + let input = expand_export_batch( + &batch( + vec![Some(1027), Some(1025), Some(1024), Some(1025), Some(1026)], + vec![ + Some("other"), + Some("selected-a"), + None, + Some("selected-b"), + None, + ], + ), + schema.arrow_schema().clone(), + ) + .unwrap(); + unit.physical_table = table::test_util::MemTable::table( + "phy", + GreptimeRecordBatch::from_df_record_batch(schema, input), + ); + let context = datafusion::prelude::SessionContext::new_with_config( + datafusion::prelude::SessionConfig::new().with_target_partitions(4), + ); + let batches = context + .execute_logical_plan(unit.build_plan(None).unwrap()) + .await + .unwrap() + .collect() + .await + .unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let summary = export_stream( + &unit, + stream(batches), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 2); + assert_eq!(summary.files, 1); + let (_, output) = read(&store, "cpu.v1.parquet").await; + let hosts = output + .iter() + .flat_map(|batch| batch.column(0).as_string::().iter()) + .collect::>(); + assert_eq!(hosts, vec![Some("selected-a"), Some("selected-b")]); +} + +#[tokio::test] +async fn dictionary_null_values_write_the_logical_parquet_schema() { + let input = batch(vec![Some(1025); 4], vec![Some("placeholder"); 4]); + let mut columns = input.columns().to_vec(); + columns[2] = Arc::new( + DictionaryArray::::try_new( + UInt32Array::from(vec![Some(0), Some(1), None, Some(0)]), + Arc::new(StringArray::from(vec![Some("host"), None])), + ) + .unwrap(), + ); + let input = RecordBatch::try_new(input.schema(), columns).unwrap(); + let unit = unit(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let summary = export_stream( + &unit, + stream(vec![input]), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 4); + let (schema, output) = read(&store, "cpu.v1.parquet").await; + assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); + let hosts = output + .iter() + .flat_map(|batch| batch.column(0).as_string::().iter()) + .collect::>(); + assert_eq!(hosts, vec![Some("host"), None, None, Some("host")]); +} + +#[tokio::test] +async fn dictionary_slices_share_backing_and_respect_budget() { + let input = batch(vec![Some(1025); 8192], vec![Some("host"); 8192]) + .project(&[2, 3, 1]) + .unwrap(); + let schema = unit().logical_tables[&1025].schema.clone(); + let budget = ExportWriteBudget::new(1); + let token = CancellationToken::new(); + let conversion = expand_bounded_slice( + input.clone(), + schema.clone(), + 7, + 8192, + 1024 * 1024, + &budget, + &token, + ); + tokio::pin!(conversion); + // Sharing an admitted batch completes without waiting for a conversion worker. + let std::task::Poll::Ready(result) = futures::poll!(&mut conversion) else { + panic!("compatible slice dispatched a conversion worker"); + }; + let (payload, rows) = result.unwrap(); + assert_eq!(rows, 8185); + let input_values = input.column(0).as_dictionary::().values(); + let output_values = payload + .batch + .column(0) + .as_dictionary::() + .values(); + assert!(Arc::ptr_eq(input_values, output_values)); + assert!(64 * 1024 * 1024 - budget.available().1 >= input.get_array_memory_size()); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + + // A batch which does not fit still takes the bounded conversion path. + let (payload, rows) = + expand_bounded_slice(input.clone(), schema.clone(), 7, 8192, 128, &budget, &token) + .await + .unwrap(); + assert_eq!(rows, 1); + drop(payload); + assert!( + expand_bounded_slice(input, schema, 7, 8192, 1, &budget, &token) + .await + .is_err() + ); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); +} + +#[test] +fn dictionary_size_bound_limits_inline_work() { + let values = (0..1025).map(|i| format!("host-{i}")).collect::>(); + let dictionary = DictionaryArray::::new( + UInt32Array::from(vec![0; 2048]), + Arc::new(StringArray::from_iter_values(values.iter())), + ); + let input = RecordBatch::try_from_iter([("host", Arc::new(dictionary) as ArrayRef)]).unwrap(); + assert!( + estimate_batch_bytes_upper_bound(&input, 2048, &[]) + .unwrap() + .is_none() + ); + assert_eq!( + rows_within_budget(&input, 0, 2048, 128, &[]).unwrap(), + (3, 114) + ); +} + +#[tokio::test] +async fn narrow_batches_share_row_groups() { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let batches = (0..8) + .map(|_| batch(vec![Some(1025); 8192], vec![Some("host"); 8192])) + .collect(); + let summary = export_stream( + &unit(), + stream(batches), + &store, + LogicalTableExportLimits::default(), + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 65536); + let bytes = store.read("cpu.v1.parquet").await.unwrap().to_bytes(); + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + assert_eq!(reader.metadata().num_row_groups(), 1); + let rows: usize = reader + .build() + .unwrap() + .map(|batch| batch.unwrap().num_rows()) + .sum(); + assert_eq!(rows, summary.rows); +} + #[tokio::test] async fn routes_across_batches_and_writes_empty_files() { let unit = unit(); @@ -249,6 +431,19 @@ async fn existing_outputs_are_not_overwritten() { ); } +#[test] +fn repeated_dictionary_budget_preserves_null_and_offset_charges() { + let host = "x".repeat(64); + let input = batch(vec![Some(1025); 8], vec![Some(host.as_str()); 8]) + .project(&[2, 4, 1]) + .unwrap(); + // Each row charges 96 bytes for the dictionary value, 32 for null, and 24 for time. + assert_eq!(rows_within_budget(&input, 0, 8, 1215, &[]).unwrap().0, 7); + let (rows, estimated) = rows_within_budget(&input, 0, 8, 1280, &[]).unwrap(); + assert_eq!(rows, 8); + assert!((1216..=1280).contains(&estimated)); +} + #[test] fn dictionary_and_nested_histogram_values_are_bounded_before_expansion() { let dictionary = DictionaryArray::::try_new( @@ -291,6 +486,25 @@ fn validates_selected_schemas_and_projects_only_selected_columns() { .unwrap(); assert_eq!(one.scan_projection, vec![0, 3]); assert_eq!(one.logical_tables[&1025].projection, vec![1]); + let LogicalPlan::Filter(filter) = one.build_plan(None).unwrap() else { + panic!("single-table export must filter the routing key without sorting"); + }; + let datafusion_expr::Expr::BinaryExpr(predicate) = filter.predicate else { + panic!("expected table-id equality"); + }; + assert_eq!(predicate.op, datafusion_expr::Operator::Eq); + assert!(matches!( + predicate.left.as_ref(), + datafusion_expr::Expr::Column(column) if column.name == TABLE_ID + )); + assert!(matches!( + predicate.right.as_ref(), + datafusion_expr::Expr::Literal(datafusion_common::ScalarValue::UInt32(Some(1025)), _) + )); + assert!(matches!( + unit.build_plan(None).unwrap(), + LogicalPlan::Sort(_) + )); assert!( LogicalTableExport::try_new(unit.physical_table.clone(), &[selected.clone(), selected]) .is_err() @@ -333,7 +547,7 @@ async fn cancellation_drops_input_and_aborts_active_upload() { } #[tokio::test] -async fn native_histogram_parquet_roundtrip() { +async fn supported_metric_types_parquet_roundtrip() { fn sample(data_type: &DataType) -> ArrayRef { match data_type { DataType::Struct(fields) => Arc::new(StructArray::new( @@ -359,44 +573,102 @@ async fn native_histogram_parquet_roundtrip() { other => panic!("unexpected histogram field {other}"), } } - let data_type = common_query::native_histogram::native_histogram_arrow_type(); - let field = Field::new("histogram", data_type.clone(), true); - let unit = LogicalTableExport::try_new( - table( - 1024, - "phy", - vec![Field::new(TABLE_ID, DataType::UInt32, false), field.clone()], - true, - ), - &[table(1025, "histogram", vec![field], false)], - ) - .unwrap(); - let batch = RecordBatch::try_from_iter_with_nullable([ - ( - TABLE_ID, - Arc::new(UInt32Array::from(vec![1025; 2])) as ArrayRef, - false, - ), - ("histogram", sample(&data_type), true), - ]) - .unwrap(); - let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); - let summary = export_stream( - &unit, - stream(vec![batch.clone()]), - &store, - export_limits(), - &CancellationToken::new(), - ) - .await - .unwrap(); - assert_eq!(summary.rows, 2); - let (schema, batches) = read(&store, "histogram.parquet").await; - assert_eq!(schema.fields(), unit.logical_tables[&1025].schema.fields()); - assert_eq!( - arrow::compute::concat_batches(&schema, &batches).unwrap(), - batch.project(&[1]).unwrap() - ); + use common_datasource::file_format::parquet::packed_reader::{ + PackReadWindows, PackedParquetReader, + }; + use common_datasource::packed_snapshot::{PACK_INDEX_FILE, PackIndex}; + use common_datasource::packed_writer::PackedWriter; + use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder; + use futures::TryStreamExt; + + let timestamps: [ArrayRef; 4] = [ + Arc::new(arrow::array::TimestampSecondArray::from(vec![4, 5])), + Arc::new(arrow::array::TimestampMillisecondArray::from(vec![4, 5])), + Arc::new(arrow::array::TimestampMicrosecondArray::from(vec![4, 5])), + Arc::new(arrow::array::TimestampNanosecondArray::from(vec![4, 5])), + ]; + let values = [ + Arc::new(Float64Array::from(vec![Some(3.5), None])) as ArrayRef, + sample(&common_query::native_histogram::native_histogram_arrow_type()), + ]; + for timestamp in timestamps { + for value in &values { + let logical = RecordBatch::try_from_iter_with_nullable([ + ("ts", timestamp.clone(), false), + ( + "host", + Arc::new(StringArray::from(vec![Some(""), None])) as ArrayRef, + true, + ), + ("value", value.clone(), true), + ]) + .unwrap(); + let mut fields = vec![Field::new(TABLE_ID, DataType::UInt32, false)]; + fields.extend(logical.schema().fields().iter().map(|f| f.as_ref().clone())); + let unit = LogicalTableExport::try_new( + table(1024, "phy", fields.clone(), true), + &[table(1025, "metric", fields[1..].to_vec(), false)], + ) + .unwrap(); + let mut columns = vec![Arc::new(UInt32Array::from(vec![1025; 2])) as ArrayRef]; + columns.extend_from_slice(logical.columns()); + let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + for packed in [false, true] { + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let writer = packed.then(|| PackedWriter::new(store.clone()).unwrap()); + let summary = export_stream_managed( + &unit, + stream(vec![input.clone()]), + &store, + export_limits(), + &CancellationToken::new(), + ExportWriteBudget::new(1), + writer.clone(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 2); + let batches = if let Some(writer) = writer { + writer + .lock() + .await + .finish(&CancellationToken::new()) + .await + .unwrap(); + let index: PackIndex = serde_json::from_slice( + &store.read(PACK_INDEX_FILE).await.unwrap().to_bytes(), + ) + .unwrap(); + index.validate_membership(["metric"]).unwrap(); + let entry = &index.tables[0]; + let object = &index.objects[0]; + let reader = PackedParquetReader::new( + PackReadWindows::new(store.clone()), + object.path.clone(), + object.length, + entry.offset, + entry.length, + ) + .unwrap(); + ParquetRecordBatchStreamBuilder::new(reader) + .await + .unwrap() + .build() + .unwrap() + .try_collect::>() + .await + .unwrap() + } else { + read(&store, "metric.parquet").await.1 + }; + assert_eq!(batches[0].schema().fields(), logical.schema().fields()); + assert_eq!( + arrow::compute::concat_batches(&logical.schema(), &batches).unwrap(), + logical + ); + } + } + } } struct PausedFileWriter { @@ -748,9 +1020,12 @@ async fn retained_backing_is_reserved_before_conversion_and_until_payload_drop() drop(blocker); let (payload, rows) = convert.await.unwrap(); assert_eq!(rows, 1); - assert!( - 64 * 1024 * 1024 - budget.available().1 >= full + payload.batch.get_array_memory_size() - ); + assert!(matches!( + payload.batch.column(0).data_type(), + DataType::Dictionary(_, _) + )); + // The retained dictionary shares its input buffers, so they are charged once. + assert!(64 * 1024 * 1024 - budget.available().1 >= full); drop(payload); assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); assert!(budget.reserve(64 * 1024 * 1024 + 1, &token).await.is_err()); diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 8e3057bb7aa..13f9384cb5c 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -1426,6 +1426,7 @@ struct ImportRequests { async fn import_http( instance: Arc, older: bool, + packed: bool, ) -> (String, tokio::task::JoinHandle<()>, Arc) { use axum::response::IntoResponse; use std::sync::atomic::Ordering; @@ -1446,6 +1447,12 @@ async fn import_http( if request.uri().path() == "/v1/capabilities" { observed.capabilities.fetch_add(1, Ordering::SeqCst); if older { + if packed { + return axum::Json(serde_json::json!({ + "metric_packed_import": 1 + })) + .into_response(); + } return axum::http::StatusCode::NOT_FOUND.into_response(); } } @@ -1528,6 +1535,11 @@ async fn packed_export_v2_cli_local_roundtrip() { metric_export_v2_cli_roundtrip(false, true, false).await; } +#[tokio::test(flavor = "multi_thread")] +async fn packed_import_v2_without_batch_ddl_roundtrip() { + metric_export_v2_cli_roundtrip(false, true, true).await; +} + #[tokio::test(flavor = "multi_thread")] async fn packed_export_v2_cli_s3_roundtrip() { if std::env::var("GT_S3_ENDPOINT_URL").is_ok_and(|e| !e.is_empty()) @@ -1606,7 +1618,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { &format!("CREATE VIEW \"{quoted_special}\" AS SELECT * FROM audit"), ) .await; - if packed && !s3 { + if packed && !s3 && !older { for id in 0..128 { let name = format!("batch_{id:03}"); sql(instance, &format!("CREATE TABLE {name} (host STRING PRIMARY KEY, val DOUBLE, ts TIMESTAMP TIME INDEX) ENGINE=metric WITH (on_physical_table='v2_a')")).await; @@ -1630,7 +1642,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { "CREATE TABLE bulk (host STRING PRIMARY KEY, payload STRING, ts TIMESTAMP TIME INDEX)", ) .await; - for batch in 0..64 { + for batch in 0..96 { let rows = (0..64) .map(|row| { let payload = (0..128) @@ -1813,7 +1825,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .await .unwrap() .content_length() - > 5 * 1024 * 1024 + > common_datasource::packed_writer::WRITE_BYTES as u64 ); } let mut preserved = Vec::new(); @@ -1923,7 +1935,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .await; } let (target_addr, target_server, requests) = - import_http(target.fe_instance().clone(), older).await; + import_http(target.fe_instance().clone(), older, packed).await; let state = destination.path().join("restore-state.json"); let mut import = vec![ "import-v2", @@ -1959,7 +1971,11 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { requests.capabilities.load(Ordering::SeqCst), if older { 1 } else { 2 } ); - assert_eq!(requests.batches.load(Ordering::SeqCst) > 1, !older); + if older { + assert_eq!(requests.batches.load(Ordering::SeqCst), 0); + } else { + assert!(requests.batches.load(Ordering::SeqCst) > 1); + } assert!(!state.exists()); for name in names { let source_table = table(instance, &name).await; @@ -2003,7 +2019,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .column_schemas() ); } - if packed { + if packed && !older { let schema_uri = format!("{uri}/schema-only"); let mut schema_args = vec![ "export-v2", @@ -2039,7 +2055,7 @@ async fn metric_export_v2_cli_roundtrip(s3: bool, packed: bool, older: bool) { .build() .await; let (schema_addr, schema_server, schema_requests) = - import_http(schema_target.fe_instance().clone(), false).await; + import_http(schema_target.fe_instance().clone(), false, false).await; schema_requests.fail_batch.store(false, Ordering::SeqCst); let schema_state = destination.path().join("schema-only-state.json"); let mut schema_import = vec![