rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-12 15:02:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

test: cover bounded Metric export conversion and mixed roundtrips

Browse Source 
Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
2026-09-22 11:04:29 +08:00
parent 7bff9a9fef
commit c8244ab9c7
5 changed files with 189 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/operator/src/statement/copy_table_to.rs
+10 -7
View File
@@ -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<usize> {
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)
src/operator/src/statement/export_logical_tables.rs
+70 -11
View File
@@ -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::<Result<Vec<_>>>()?;
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<RecordBatch> {
let arrays = batch
.columns()
.iter()
.zip(schema.fields())
.map(|(array, field)| expand_export_array(array, field.data_type()))
.collect::<Result<Vec<_>>>()?;
RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)
}
fn expand_export_array(array: &ArrayRef, target: &DataType) -> Result<ArrayRef> {
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::<i32>();
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::<Vec<_>>()
.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::<Result<Vec<_>>>()?;
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<usize> {
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() {
src/operator/src/statement/export_logical_tables/tests.rs
+79
View File
@@ -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::<UInt32Type>::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::<i32>();
assert_eq!(list.value_offsets(), &[0, 1]);
assert_eq!(list.values().as_string::<i32>().value(0), "x".repeat(4096));
drop(payload);
assert_eq!(budget.available(), (1, 64 * 1024 * 1024));
let values = Arc::new(DictionaryArray::<UInt32Type>::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));
}
}
src/operator/src/statement/export_logical_tables/writers.rs
+6 -1
View File
@@ -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<u32> {
self.current.as_ref().map(|(id, _)| *id)
}
tests-integration/tests/export_logical_tables.rs
+24 -17
View File
@@ -228,7 +228,7 @@ fn database_export_request(directory: &std::path::Path) -> table::requests::Copy
}
}
async fn database_export_roundtrip(instance: &Arc<Instance>) {
async fn database_export_roundtrip(instance: &Arc<Instance>, 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<Instance>) {
];
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<Instance>) {
#[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")]