diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 83f9dad01f8..f0369356b2a 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -13,11 +13,16 @@ // limitations under the License. use std::any::Any; +#[cfg(test)] +use std::cell::Cell; use std::sync::{Arc, Mutex}; use std::time::Duration; use ahash::{HashMap, HashSet}; -use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, SortOptions}; +use arrow_schema::{ + ArrowError, DataType, DataType as ArrowDataType, Field, Schema as ArrowSchema, + SchemaRef as ArrowSchemaRef, SortOptions, +}; use async_stream::stream; use common_catalog::parse_catalog_and_schema_from_db_string; use common_plugins::GREPTIME_EXEC_READ_COST; @@ -78,6 +83,117 @@ fn remote_dyn_filter_enabled(query_ctx: &QueryContextRef) -> Result { .map_err(|err| DataFusionError::External(Box::new(err))) } +fn remote_schema_mismatch(message: impl Into) -> DataFusionError { + DataFusionError::ArrowError(Box::new(ArrowError::SchemaError(message.into())), None) +} + +fn record_merge_scan_schema_error() { + MERGE_SCAN_ERRORS_TOTAL.inc(); + + #[cfg(test)] + TEST_MERGE_SCAN_SCHEMA_ERRORS.with(|count| count.set(count.get() + 1)); +} + +#[cfg(test)] +thread_local! { + // The Prometheus counter is process-global and tests run concurrently. Keep + // a thread-local companion at the production increment site for stable + // execution-path assertions. + static TEST_MERGE_SCAN_SCHEMA_ERRORS: Cell = const { Cell::new(0) }; +} + +#[cfg(test)] +fn merge_scan_schema_error_count_for_test() -> u64 { + TEST_MERGE_SCAN_SCHEMA_ERRORS.with(Cell::get) +} + +/// Returns true when the field is a JSON column, identified either by its +/// Arrow extension type or by its `greptime:type=Json` marker. +fn is_json_field(field: &Field) -> bool { + is_json_extension_type(field) + || field + .metadata() + .get(datatypes::schema::TYPE_KEY) + .map(String::as_ref) + == Some("Json") +} + +/// Returns true when two fields have the same semantic identity. Field +/// metadata is auxiliary at the MergeScan boundary and is intentionally +/// excluded from this comparison. +fn fields_semantically_equal(expected_field: &Field, actual_field: &Field) -> bool { + expected_field.name() == actual_field.name() + && expected_field.data_type() == actual_field.data_type() + && expected_field.is_nullable() == actual_field.is_nullable() +} + +fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool { + is_json_field(expected_field) + && is_json_field(actual_field) + && expected_field.name() == actual_field.name() + && expected_field.is_nullable() == actual_field.is_nullable() + && expected_field.metadata().get(datatypes::schema::TYPE_KEY) + == actual_field.metadata().get(datatypes::schema::TYPE_KEY) + && expected_field + .metadata() + .get(arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY) + == actual_field + .metadata() + .get(arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY) +} + +fn validate_remote_schema( + expected: &ArrowSchema, + actual: &ArrowSchema, + source: &str, +) -> Result<()> { + if expected.fields().len() != actual.fields().len() { + return Err(remote_schema_mismatch(format!( + "MergeScan {source} schema field count mismatch: expected {}, actual {}", + expected.fields().len(), + actual.fields().len() + ))); + } + + for (index, (expected_field, actual_field)) in expected + .fields() + .iter() + .zip(actual.fields().iter()) + .enumerate() + { + // Field metadata is auxiliary information (for example, index and + // encoding hints), so it does not participate in field identity. + if fields_semantically_equal(expected_field, actual_field) { + continue; + } + + // JSON may use Binary on the wire and a decoded Struct locally. This + // physical-type difference is allowed only when both fields retain + // the same JSON identity and JSON2 extension settings. + if json_fields_compatible(expected_field, actual_field) { + continue; + } + + let timezone_only_difference = matches!( + (expected_field.data_type(), actual_field.data_type()), + ( + ArrowDataType::Timestamp(expected_unit, expected_timezone), + ArrowDataType::Timestamp(actual_unit, actual_timezone), + ) if expected_unit == actual_unit + && expected_timezone != actual_timezone + && expected_field.name() == actual_field.name() + && expected_field.is_nullable() == actual_field.is_nullable() + ); + if !timezone_only_difference { + return Err(remote_schema_mismatch(format!( + "MergeScan {source} schema field mismatch at position {index}: expected {:?}, actual {:?}", + expected_field, actual_field + ))); + } + } + Ok(()) +} + fn acquire_remote_dyn_filter_registry_lease( context: &TaskContext, query_ctx: &QueryContextRef, @@ -434,6 +550,13 @@ impl MergeScanExec { MERGE_SCAN_ERRORS_TOTAL.inc(); DataFusionError::External(Box::new(e)) })?; + let mut advertised_schema = stream.schema().arrow_schema().clone(); + validate_remote_schema( + arrow_schema.as_ref(), + advertised_schema.as_ref(), + "advertised remote stream", + ) + .inspect_err(|_| record_merge_scan_schema_error())?; let do_get_cost = do_get_start.elapsed(); if let Some(remote_dyn_filter_registry_lease) = @@ -478,10 +601,18 @@ impl MergeScanExec { poll_duration += poll_elapsed; let batch = batch.map_err(|e| DataFusionError::External(Box::new(e)))?; - let batch = patch_batch_timezone( - arrow_schema.clone(), - batch.into_df_record_batch().columns().to_vec(), - )?; + let df_batch = batch.into_df_record_batch(); + if !Arc::ptr_eq(&advertised_schema, df_batch.schema_ref()) { + validate_remote_schema( + arrow_schema.as_ref(), + df_batch.schema_ref().as_ref(), + "remote record batch", + ) + .inspect_err(|_| record_merge_scan_schema_error())?; + advertised_schema = df_batch.schema_ref().clone(); + } + let batch = + patch_batch_timezone(arrow_schema.clone(), df_batch.columns().to_vec())?; metric.record_output_batch_rows(batch.num_rows()); if let Some(mut first_consume_timer) = first_consume_timer.take() { first_consume_timer.stop(); @@ -1089,11 +1220,15 @@ impl MergeScanMetric { #[cfg(test)] mod tests { - use std::collections::BTreeSet; + use std::collections::{BTreeSet, HashMap as StdHashMap}; + use std::pin::Pin; + use std::task::{Context, Poll}; + use arrow_schema::{DataType as TestArrowDataType, Field}; use async_trait::async_trait; use common_query::request::INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY; use common_recordbatch::adapter::{PlanMetrics, RecordBatchMetrics}; + use common_recordbatch::{DfRecordBatch, RecordBatch, RecordBatchStream}; use datafusion::config::ConfigOptions; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; @@ -1103,6 +1238,11 @@ mod tests { use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, lit as physical_lit, }; + use datatypes::arrow::array::TimestampMillisecondArray; + use datatypes::prelude::{ConcreteDataType, VectorRef}; + use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{Int64Vector, StringVector}; + use futures_util::{Stream, TryStreamExt}; use session::ReadPreference; use session::context::QueryContext; use session::query_id::QueryId; @@ -1114,6 +1254,524 @@ mod tests { use crate::dist_plan::{DynFilterRegistryManager, Subscriber}; use crate::region_query::RegionQueryHandler; + fn int64_schema(columns: &[&str]) -> Arc { + Arc::new(Schema::new( + columns + .iter() + .map(|name| ColumnSchema::new(*name, ConcreteDataType::int64_datatype(), false)) + .collect(), + )) + } + + fn record_batch(schema: Arc, columns: Vec) -> RecordBatch { + RecordBatch::new(schema, columns).expect("test record batch must match its schema") + } + + fn expected_int64_schema() -> ArrowSchema { + int64_schema(&["a", "b"]).arrow_schema().as_ref().clone() + } + + fn merge_scan_exec_with_handler( + regions: Vec, + expected_schema: ArrowSchema, + handler: Arc, + ) -> MergeScanExec { + let plan = LogicalPlanBuilder::empty(true).build().unwrap(); + MergeScanExec::new( + &SessionStateBuilder::new().build(), + TableName::new("catalog", "schema", "table"), + regions, + plan, + &expected_schema, + handler, + QueryContext::arc(), + 1, + AliasMapping::new(), + None, + false, + ) + .unwrap() + } + + async fn collect_merge_scan( + exec: MergeScanExec, + ) -> datafusion_common::Result> { + exec.execute(0, Arc::new(TaskContext::default()))? + .try_collect() + .await + } + + fn json_field_metadata(wire_form: bool, settings: &str) -> StdHashMap { + let mut metadata = StdHashMap::from([ + (datatypes::schema::TYPE_KEY.to_string(), "Json".to_string()), + ( + arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY.to_string(), + settings.to_string(), + ), + ]); + if wire_form { + metadata.insert( + arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(), + "greptime.json".to_string(), + ); + } + metadata + } + + #[test] + fn merge_scan_validates_remote_schema_semantics() { + let expected = expected_int64_schema(); + let metadata_mismatch = ArrowSchema::new_with_metadata( + vec![ + expected + .field(0) + .as_ref() + .clone() + .with_metadata(StdHashMap::from([( + "remote".to_string(), + "different".to_string(), + )])), + expected.field(1).as_ref().clone(), + ], + expected.metadata().clone(), + ); + assert!(validate_remote_schema(&expected, &metadata_mismatch, "test").is_ok()); + let wrong_type = ArrowSchema::new(vec![ + Field::new("a", TestArrowDataType::Utf8, false), + expected.field(1).as_ref().clone(), + ]); + assert!(validate_remote_schema(&expected, &wrong_type, "test").is_err()); + } + + #[test] + fn merge_scan_validates_json_binary_and_struct_compatibility() { + let settings = r#"{"json_settings":{"type_hints":[]}}"#; + let expected = ArrowSchema::new(vec![ + Field::new("j", TestArrowDataType::Binary, true) + .with_metadata(json_field_metadata(true, settings)), + ]); + let actual = + ArrowSchema::new(vec![ + Field::new( + "j", + TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new( + Field::new("a", TestArrowDataType::Utf8View, true), + )])), + true, + ) + .with_metadata(json_field_metadata(false, settings)), + ]); + assert!(validate_remote_schema(&expected, &actual, "test").is_ok()); + + // JSON extension identity metadata is semantic only for the + // Binary-vs-Struct compatibility exemption. When the physical type + // already matches, all field metadata remains non-semantic, just as + // it is for ordinary fields. + let plain_binary = ArrowSchema::new(vec![Field::new("j", TestArrowDataType::Binary, true)]); + assert!(validate_remote_schema(&expected, &plain_binary, "test").is_ok()); + + let changed = ArrowSchema::new(vec![ + Field::new( + "j", + TestArrowDataType::Struct(arrow_schema::Fields::empty()), + true, + ) + .with_metadata(json_field_metadata( + false, + r#"{"json_settings":{"type_hints":[["a",{"JsonType":"Int64"}]]}}"#, + )), + ]); + assert!(validate_remote_schema(&expected, &changed, "test").is_err()); + } + + #[test] + fn merge_scan_allows_top_level_schema_metadata_mismatch() { + let fields = expected_int64_schema().fields().clone(); + let expected = ArrowSchema::new_with_metadata( + fields.clone(), + StdHashMap::from([("version".to_string(), "1".to_string())]), + ); + let actual = ArrowSchema::new_with_metadata( + fields, + StdHashMap::from([("version".to_string(), "2".to_string())]), + ); + assert!(validate_remote_schema(&expected, &actual, "test").is_ok()); + } + + #[tokio::test] + async fn merge_scan_execution_rejects_reordered_advertised_columns_before_repair() { + let region_id = RegionId::new(1024, 1); + let schema = int64_schema(&["b", "a"]); + let batch = record_batch( + schema.clone(), + vec![ + Arc::new(Int64Vector::from_slice([2])) as _, + Arc::new(Int64Vector::from_slice([1])) as _, + ], + ); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + schema, + vec![batch], + )])); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_err() + ); + } + + #[tokio::test] + async fn merge_scan_execution_rejects_missing_and_extra_advertised_columns() { + for columns in [["a"].as_slice(), ["a", "b", "extra"].as_slice()] { + let region_id = RegionId::new(1024, columns.len() as u32); + let schema = int64_schema(columns); + let vectors = columns + .iter() + .map(|_| Arc::new(Int64Vector::from_slice([1])) as VectorRef) + .collect(); + let batch = record_batch(schema.clone(), vectors); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + schema, + vec![batch], + )])); + let before = merge_scan_schema_error_count_for_test(); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_err() + ); + assert_eq!(merge_scan_schema_error_count_for_test(), before + 1); + } + } + + #[tokio::test] + async fn merge_scan_execution_rejects_advertised_schema_and_records_metric() { + let region_id = RegionId::new(1024, 10); + let advertised = int64_schema(&["a", "b", "extra"]); + let batch = record_batch( + advertised.clone(), + vec![ + Arc::new(Int64Vector::from_slice([1])) as _, + Arc::new(Int64Vector::from_slice([2])) as _, + Arc::new(Int64Vector::from_slice([3])) as _, + ], + ); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + let before = merge_scan_schema_error_count_for_test(); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_err() + ); + assert_eq!(merge_scan_schema_error_count_for_test(), before + 1); + } + + #[tokio::test] + async fn merge_scan_execution_rejects_advertised_schema_inner_batch_mismatch() { + let region_id = RegionId::new(1024, 1); + let advertised = int64_schema(&["a", "b"]); + let inner = record_batch( + int64_schema(&["b", "a"]), + vec![ + Arc::new(Int64Vector::from_slice([2])) as _, + Arc::new(Int64Vector::from_slice([1])) as _, + ], + ) + .into_df_record_batch(); + let batch = RecordBatch::from_df_record_batch(advertised.clone(), inner); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + let before = merge_scan_schema_error_count_for_test(); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_err() + ); + assert_eq!(merge_scan_schema_error_count_for_test(), before + 1); + } + + #[tokio::test] + async fn merge_scan_execution_rejects_non_timezone_type_mismatch() { + let region_id = RegionId::new(1024, 11); + let expected = ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Int64, false)]); + let advertised = Arc::new( + Schema::try_from(ArrowSchema::new(vec![Field::new( + "value", + TestArrowDataType::Utf8, + false, + )])) + .unwrap(), + ); + let batch = record_batch( + advertised.clone(), + vec![Arc::new(StringVector::from_slice(&["wrong"])) as _], + ); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected, + handler, + )) + .await + .is_err() + ); + } + + #[tokio::test] + async fn merge_scan_execution_rejects_nullability_mismatch() { + let region_id = RegionId::new(1024, 12); + let expected = ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Int64, false)]); + let advertised = Arc::new( + Schema::try_from(ArrowSchema::new(vec![Field::new( + "value", + TestArrowDataType::Int64, + true, + )])) + .unwrap(), + ); + let batch = record_batch( + advertised.clone(), + vec![Arc::new(Int64Vector::from_slice([1])) as _], + ); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected, + handler, + )) + .await + .is_err() + ); + } + + #[tokio::test] + async fn merge_scan_execution_accepts_timezone_difference_and_repairs_output_schema() { + let region_id = RegionId::new(1024, 13); + let expected = Arc::new(ArrowSchema::new(vec![Field::new( + "ts", + TestArrowDataType::Timestamp( + arrow_schema::TimeUnit::Millisecond, + Some("Asia/Shanghai".into()), + ), + false, + )])); + let advertised = Arc::new( + Schema::try_from(Arc::new(ArrowSchema::new(vec![Field::new( + "ts", + TestArrowDataType::Timestamp( + arrow_schema::TimeUnit::Millisecond, + Some("UTC".into()), + ), + false, + )]))) + .unwrap(), + ); + let inner_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "ts", + TestArrowDataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + false, + )])); + let inner = DfRecordBatch::try_new( + inner_schema, + vec![Arc::new(TimestampMillisecondArray::from_iter_values([1])) as _], + ) + .unwrap(); + let batch = RecordBatch::from_df_record_batch(advertised.clone(), inner); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + let output = collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected.as_ref().clone(), + handler, + )) + .await + .unwrap(); + assert_eq!(output.len(), 1); + assert_eq!(output[0].schema_ref().as_ref(), expected.as_ref()); + } + + #[tokio::test] + async fn merge_scan_execution_rejects_hidden_extra_inner_column() { + let region_id = RegionId::new(1024, 1); + let advertised = int64_schema(&["a", "b"]); + let inner = record_batch( + int64_schema(&["a", "b", "extra"]), + vec![ + Arc::new(Int64Vector::from_slice([1])) as _, + Arc::new(Int64Vector::from_slice([2])) as _, + Arc::new(Int64Vector::from_slice([3])) as _, + ], + ) + .into_df_record_batch(); + let batch = RecordBatch::from_df_record_batch(advertised.clone(), inner); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_err() + ); + } + + #[tokio::test] + async fn merge_scan_execution_accepts_structurally_equal_distinct_schema_arcs() { + let region_id = RegionId::new(1024, 1); + let advertised = int64_schema(&["a", "b"]); + let distinct = Arc::new( + Schema::try_from(Arc::new(advertised.arrow_schema().as_ref().clone())).unwrap(), + ); + assert_eq!( + advertised.arrow_schema().as_ref(), + distinct.arrow_schema().as_ref() + ); + assert!(!Arc::ptr_eq(&advertised, &distinct)); + let batch = record_batch( + distinct, + vec![ + Arc::new(Int64Vector::from_slice([1])) as _, + Arc::new(Int64Vector::from_slice([2])) as _, + ], + ); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + advertised, + vec![batch], + )])); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_ok() + ); + } + + #[tokio::test] + async fn merge_scan_execution_ignores_skipping_and_inverted_index_metadata() { + let region_id = RegionId::new(1024, 1); + let fields = vec![ + Field::new("a", TestArrowDataType::Int64, false).with_metadata(StdHashMap::from([ + ( + datatypes::schema::SKIPPING_INDEX_KEY.to_string(), + "bloom".to_string(), + ), + ( + datatypes::schema::INVERTED_INDEX_KEY.to_string(), + "true".to_string(), + ), + ])), + Field::new("b", TestArrowDataType::Int64, false), + ]; + let schema = Arc::new(Schema::try_from(Arc::new(ArrowSchema::new(fields))).unwrap()); + let batch = record_batch( + schema.clone(), + vec![ + Arc::new(Int64Vector::from_slice([1])) as _, + Arc::new(Int64Vector::from_slice([2])) as _, + ], + ); + let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( + region_id, + schema, + vec![batch], + )])); + assert!( + collect_merge_scan(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + )) + .await + .is_ok() + ); + } + + #[test] + fn merge_scan_schema_validation_rejects_nullability_and_type_differences_but_accepts_timezone_only_difference() + { + let expected = ArrowSchema::new(vec![ + Field::new("value", TestArrowDataType::Int64, false), + Field::new( + "ts", + TestArrowDataType::Timestamp( + arrow_schema::TimeUnit::Millisecond, + Some("UTC".into()), + ), + false, + ), + ]); + let nullable = ArrowSchema::new(vec![ + Field::new("value", TestArrowDataType::Int64, true), + expected.field(1).as_ref().clone(), + ]); + let wrong_type = ArrowSchema::new(vec![ + Field::new("value", TestArrowDataType::Utf8, false), + expected.field(1).as_ref().clone(), + ]); + assert!(validate_remote_schema(&expected, &nullable, "test").is_err()); + assert!(validate_remote_schema(&expected, &wrong_type, "test").is_err()); + + let timezone = ArrowSchema::new(vec![ + expected.field(0).as_ref().clone(), + Field::new( + "ts", + TestArrowDataType::Timestamp( + arrow_schema::TimeUnit::Millisecond, + Some("Asia/Shanghai".into()), + ), + false, + ), + ]); + assert!(validate_remote_schema(&expected, &timezone, "test").is_ok()); + } + fn test_query_id(value: u128) -> QueryId { QueryId::from(Uuid::from_u128(value)) } @@ -1143,7 +1801,7 @@ mod tests { regions, plan, &schema, - Arc::new(TestRegionQueryHandler), + Arc::new(TestRegionQueryHandler::default()), QueryContext::arc(), target_partition, AliasMapping::new(), @@ -1286,16 +1944,87 @@ mod tests { assert_eq!(registry_manager.registry_count(), 0); } - struct TestRegionQueryHandler; + #[derive(Clone)] + struct TestRegionResponse { + advertised_schema: Arc, + batches: Vec, + } + + #[derive(Default)] + struct TestRegionQueryHandler { + responses: StdHashMap, + } + + impl TestRegionQueryHandler { + fn with_responses( + responses: impl IntoIterator, Vec)>, + ) -> Self { + Self { + responses: responses + .into_iter() + .map(|(region_id, advertised_schema, batches)| { + ( + region_id, + TestRegionResponse { + advertised_schema, + batches, + }, + ) + }) + .collect(), + } + } + } + + struct TestRecordBatchStream { + schema: Arc, + batches: Vec, + index: usize, + } + + impl Stream for TestRecordBatchStream { + type Item = common_recordbatch::error::Result; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + if let Some(batch) = self.batches.get(self.index).cloned() { + self.index += 1; + Poll::Ready(Some(Ok(batch))) + } else { + Poll::Ready(None) + } + } + } + + impl RecordBatchStream for TestRecordBatchStream { + fn schema(&self) -> Arc { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> { + None + } + + fn metrics(&self) -> Option { + None + } + } #[async_trait] impl RegionQueryHandler for TestRegionQueryHandler { async fn do_get( &self, _read_preference: ReadPreference, - _request: common_query::request::QueryRequest, + request: common_query::request::QueryRequest, ) -> crate::error::Result { - unimplemented!("test only") + let response = self + .responses + .get(&request.region_id) + .expect("test handler needs a response for every requested region"); + Ok(Box::pin(TestRecordBatchStream { + schema: response.advertised_schema.clone(), + batches: response.batches.clone(), + index: 0, + })) } async fn handle_remote_dyn_filter_update( @@ -1342,7 +2071,7 @@ mod tests { let session_state = SessionStateBuilder::new().build(); - let handler = Arc::new(TestRegionQueryHandler); + let handler = Arc::new(TestRegionQueryHandler::default()); let target_partition = 2; let exec = MergeScanExec::new( @@ -1398,7 +2127,7 @@ mod tests { let regions = vec![RegionId::new(1024, 1)]; let query_ctx = QueryContext::arc(); let session_state = SessionStateBuilder::new().build(); - let handler = Arc::new(TestRegionQueryHandler); + let handler = Arc::new(TestRegionQueryHandler::default()); let exec = MergeScanExec::new( &session_state, table, @@ -1511,7 +2240,7 @@ mod tests { vec![region_id], plan, &schema, - Arc::new(TestRegionQueryHandler), + Arc::new(TestRegionQueryHandler::default()), QueryContext::arc(), 1, AliasMapping::new(), @@ -1580,7 +2309,7 @@ mod tests { vec![logical_region_id], plan, &schema, - Arc::new(TestRegionQueryHandler), + Arc::new(TestRegionQueryHandler::default()), QueryContext::arc(), 1, AliasMapping::new(),