diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index a415a7c3904..5f5b73003db 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -47,7 +47,7 @@ use crate::error::{self, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, Result}; use crate::statement::StatementExecutor; use crate::statement::export_logical_tables::writers::ExportWriteBudget; use crate::statement::export_logical_tables::{ - LogicalTableExportLimits, map_writer_error, rows_within_budget, + LogicalTableExportLimits, expand_export_batch, map_writer_error, rows_within_budget, }; // The buffer size should be greater than 5MB (minimum multipart upload size). @@ -242,9 +242,7 @@ pub(crate) async fn stream_to_managed_parquet( budget: &ExportWriteBudget, token: &CancellationToken, ) -> Result { - use common_recordbatch::{ - RecordBatch, map_dictionary_to_values, map_dictionary_to_values_schema, - }; + use common_recordbatch::{RecordBatch, map_dictionary_to_values_schema}; let original = stream.schema(); let (expanded_schema, expand) = map_dictionary_to_values_schema(original.clone()); let (mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone()); @@ -284,7 +282,7 @@ pub(crate) async fn stream_to_managed_parquet( }) .await .context(error::JoinTaskSnafu)??; - let reservation = retained.saturating_add(estimated.saturating_mul(2 * expansion)); + let reservation = retained.saturating_add(estimated.saturating_mul(4 * expansion)); let permit = budget.reserve(reservation, token).await?; let batch = batch.clone(); let (original, expanded_schema, output_schema) = ( @@ -298,8 +296,13 @@ pub(crate) async fn stream_to_managed_parquet( batch.slice(offset, len), ); if expand { - batch = map_dictionary_to_values(batch, &original, &expanded_schema) - .context(error::BuildRecordBatchSnafu)?; + batch = RecordBatch::from_df_record_batch( + expanded_schema.clone(), + expand_export_batch( + &batch.into_df_record_batch(), + expanded_schema.arrow_schema().clone(), + )?, + ); } if json { batch = map_json_type_to_string(batch, &expanded_schema, &output_schema) diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index a3f025d3627..1b8d012fc65 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -24,8 +24,9 @@ pub(crate) mod writers; use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; -use arrow::array::{Array, AsArray, UInt32Array}; -use arrow::compute::cast; +use arrow::array::{Array, ArrayRef, AsArray, ListArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{cast, take}; use arrow::datatypes::{DataType, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; @@ -586,17 +587,11 @@ async fn expand_bounded_slice( }) .await .context(error::JoinTaskSnafu)??; - let reservation = retained.saturating_add(estimated.saturating_mul(2)); + 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 arrays = slice - .columns() - .iter() - .zip(schema.fields()) - .map(|(array, field)| cast(array, field.data_type()).context(error::ComputeArrowSnafu)) - .collect::>>()?; - let expanded = RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)?; + let expanded = expand_export_batch(&slice, schema)?; ensure!( expanded.get_array_memory_size() <= reservation, LogicalTableExportResourceSnafu { @@ -615,6 +610,70 @@ async fn expand_bounded_slice( .context(error::JoinTaskSnafu)? } +/// 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 { + let arrays = batch + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu) +} + +fn expand_export_array(array: &ArrayRef, target: &DataType) -> Result { + if array.data_type() == target { + return Ok(array.clone()); + } + match (array.data_type(), target) { + (DataType::Dictionary(_, _), _) => { + downcast_dictionary_array! { + array => { + // Select dictionary entries before recursively converting nested values. + let selected = take(array.values().as_ref(), array.keys(), None) + .context(error::ComputeArrowSnafu)?; + expand_export_array(&selected, target) + }, + _ => error::UnexpectedSnafu { violated: "invalid dictionary array" }.fail(), + } + } + (DataType::List(_), DataType::List(field)) => { + let list = array.as_list::(); + let offsets = list.value_offsets(); + let start = offsets[0]; + let end = offsets[list.len()]; + let values = list.values().slice(start as usize, (end - start) as usize); + let values = expand_export_array(&values, field.data_type())?; + let offsets = OffsetBuffer::new( + offsets + .iter() + .map(|offset| offset - start) + .collect::>() + .into(), + ); + Ok(Arc::new( + ListArray::try_new(field.clone(), offsets, values, list.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + (DataType::Struct(_), DataType::Struct(fields)) => { + let array = array.as_struct(); + let columns = array + .columns() + .iter() + .zip(fields) + .map(|(array, field)| expand_export_array(array, field.data_type())) + .collect::>>()?; + Ok(Arc::new( + StructArray::try_new(fields.clone(), columns, array.nulls().cloned()) + .context(error::ComputeArrowSnafu)?, + )) + } + _ => cast(array, target).context(error::ComputeArrowSnafu), + } +} + pub(crate) fn map_writer_error( source: common_datasource::error::Error, path: &str, @@ -639,7 +698,7 @@ pub(crate) fn map_writer_error( // Count only selected logical values before dictionary expansion, including nested // histogram lists/structs. Offset and validity overhead is charged per value. fn estimate_value_size(array: &dyn Array, row: usize) -> Result { - if array.is_null(row) { + if array.is_null(row) && !matches!(array.data_type(), DataType::Struct(_) | DataType::List(_)) { return Ok(32); } let bytes = match array.data_type() { diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index 74311963828..dc0e71c9cb8 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -799,3 +799,82 @@ async fn groups_and_ordinary_files_share_writer_admission_and_drain() { } } } + +#[tokio::test] +async fn nested_dictionary_conversion_limits_child_ranges_and_charges_null_parents() { + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0; 1024]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let list = ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0i32, 1023, 1024].into()), + values, + None, + ); + let input = RecordBatch::try_from_iter([("list", Arc::new(list) as ArrayRef)]).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "list", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + )])); + let budget = ExportWriteBudget::new(1); + let (payload, rows) = expand_bounded_slice( + input, + schema, + 1, + 2, + 8192, + &budget, + &CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(rows, 1); + assert!(payload.batch.get_array_memory_size() < 16384); + let list = payload.batch.column(0).as_list::(); + assert_eq!(list.value_offsets(), &[0, 1]); + assert_eq!(list.values().as_string::().value(0), "x".repeat(4096)); + drop(payload); + assert_eq!(budget.available(), (1, 64 * 1024 * 1024)); + + let values = Arc::new(DictionaryArray::::new( + UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["x".repeat(4096)])), + )) as ArrayRef; + let structure = StructArray::new( + vec![Arc::new(Field::new( + "child", + values.data_type().clone(), + true, + ))] + .into(), + vec![values], + Some(arrow::buffer::NullBuffer::from(vec![false])), + ); + let input = RecordBatch::try_from_iter([("struct", Arc::new(structure) as ArrayRef)]).unwrap(); + assert!(rows_within_budget(&input, 0, 1, 128).is_err()); +} + +#[tokio::test] +async fn completed_table_workers_are_reaped_during_admission() { + for parallelism in [1, 4] { + let mut unit = unit(); + let file = unit.logical_tables.get_mut(&1025).unwrap(); + let store = ObjectStore::new(object_store::services::Memory::default()).unwrap(); + let budget = ExportWriteBudget::new(parallelism); + let token = CancellationToken::new(); + let mut writers = TableWriters::new(budget.clone()); + for id in 0..64 { + file.output.path = format!("empty-{id}.parquet"); + writers + .open(id, file, &store, export_limits(), &token) + .await + .unwrap(); + assert!(writers.pending_tasks() <= parallelism + 1); + } + writers.drain(Ok(()), &token).await.unwrap(); + assert_eq!(writers.pending_tasks(), 0); + assert_eq!(budget.available(), (parallelism, 64 * 1024 * 1024)); + } +} diff --git a/src/operator/src/statement/export_logical_tables/writers.rs b/src/operator/src/statement/export_logical_tables/writers.rs index c19ee4824fd..0c818b9119c 100644 --- a/src/operator/src/statement/export_logical_tables/writers.rs +++ b/src/operator/src/statement/export_logical_tables/writers.rs @@ -90,7 +90,7 @@ impl ExportWriteBudget { let available = PAYLOAD_BYTES .saturating_sub(backing) .saturating_sub(overhead) - / 2; + / 4; let conversion = requested.min(available); ensure!( conversion > 0, @@ -122,6 +122,11 @@ impl TableWriters { } } + #[cfg(test)] + pub(crate) fn pending_tasks(&self) -> usize { + self.tasks.len() + } + pub(crate) fn table_id(&self) -> Option { self.current.as_ref().map(|(id, _)| *id) } diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 92888cced3c..7930abb7c39 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -228,7 +228,7 @@ fn database_export_request(directory: &std::path::Path) -> table::requests::Copy } } -async fn database_export_roundtrip(instance: &Arc) { +async fn database_export_roundtrip(instance: &Arc, parallelism: usize) { let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); let (first_logical_table_names, _, renamed_physical_table) = create_metric_export_source_tables(instance, "db_a", "dense").await; @@ -253,7 +253,9 @@ async fn database_export_roundtrip(instance: &Arc) { ]; let mut names = selected.clone(); names.extend([renamed_physical_table, "dashboard".into()]); - let req = database_export_request(&destination.path().join("data")); + let mut req = database_export_request(&destination.path().join("data")); + req.with + .insert("parallelism".into(), parallelism.to_string()); let executor = instance.statement_executor(); let captured = executor .capture_database_export_tables(&req, None, &QueryContext::arc()) @@ -445,26 +447,31 @@ async fn database_export_roundtrip(instance: &Arc) { #[tokio::test(flavor = "multi_thread")] async fn database_export_standalone_roundtrip() { - let standalone = GreptimeDbStandaloneBuilder::new("database_export") - .build() - .await; - database_export_roundtrip(standalone.fe_instance()).await; + for parallelism in [1, 4] { + let standalone = + GreptimeDbStandaloneBuilder::new(&format!("database_export_{parallelism}")) + .build() + .await; + database_export_roundtrip(standalone.fe_instance(), parallelism).await; + } } #[tokio::test(flavor = "multi_thread")] async fn database_export_distributed_roundtrip() { - let cluster = GreptimeDbClusterBuilder::new("database_export") - .await - .with_datanodes(2) - .with_local_file_access( - common_datasource::object_store::LocalFileAccess::sandboxed( - common_test_util::find_workspace_path("."), + for parallelism in [1, 4] { + let cluster = GreptimeDbClusterBuilder::new(&format!("database_export_{parallelism}")) + .await + .with_datanodes(2) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed( + common_test_util::find_workspace_path("."), + ) + .unwrap(), ) - .unwrap(), - ) - .build(false) - .await; - database_export_roundtrip(cluster.fe_instance()).await; + .build(false) + .await; + database_export_roundtrip(cluster.fe_instance(), parallelism).await; + } } #[tokio::test(flavor = "multi_thread")]