fix(query): validate merge scan remote schema (#8579)

Reconstructed remote advertised and batch schema validation from 7bf8ba624f against beta2 merge-scan APIs; excludes #8615 dynamic-filter lifecycle changes.

Backport-of: 7bf8ba624f
Signed-off-by: discord9 <discord9@163.com>
(cherry picked from commit 8f6710ba395270f66e4f616baf1f091324f42e2d)

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<bool> {
         .map_err(|err| DataFusionError::External(Box::new(err)))
 }
 
+fn remote_schema_mismatch(message: impl Into<String>) -> 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<u64> = 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<Schema> {
+        Arc::new(Schema::new(
+            columns
+                .iter()
+                .map(|name| ColumnSchema::new(*name, ConcreteDataType::int64_datatype(), false))
+                .collect(),
+        ))
+    }
+
+    fn record_batch(schema: Arc<Schema>, columns: Vec<VectorRef>) -> 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<RegionId>,
+        expected_schema: ArrowSchema,
+        handler: Arc<TestRegionQueryHandler>,
+    ) -> 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<Vec<DfRecordBatch>> {
+        exec.execute(0, Arc::new(TaskContext::default()))?
+            .try_collect()
+            .await
+    }
+
+    fn json_field_metadata(wire_form: bool, settings: &str) -> StdHashMap<String, String> {
+        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<Schema>,
+        batches: Vec<RecordBatch>,
+    }
+
+    #[derive(Default)]
+    struct TestRegionQueryHandler {
+        responses: StdHashMap<RegionId, TestRegionResponse>,
+    }
+
+    impl TestRegionQueryHandler {
+        fn with_responses(
+            responses: impl IntoIterator<Item = (RegionId, Arc<Schema>, Vec<RecordBatch>)>,
+        ) -> Self {
+            Self {
+                responses: responses
+                    .into_iter()
+                    .map(|(region_id, advertised_schema, batches)| {
+                        (
+                            region_id,
+                            TestRegionResponse {
+                                advertised_schema,
+                                batches,
+                            },
+                        )
+                    })
+                    .collect(),
+            }
+        }
+    }
+
+    struct TestRecordBatchStream {
+        schema: Arc<Schema>,
+        batches: Vec<RecordBatch>,
+        index: usize,
+    }
+
+    impl Stream for TestRecordBatchStream {
+        type Item = common_recordbatch::error::Result<RecordBatch>;
+
+        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            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<Schema> {
+            self.schema.clone()
+        }
+
+        fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> {
+            None
+        }
+
+        fn metrics(&self) -> Option<RecordBatchMetrics> {
+            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<common_recordbatch::SendableRecordBatchStream> {
-            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(),