servers/batcher/logical_table/

batch.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;
use std::time::Instant;

use catalog::CatalogManagerRef;
use common_batcher::flush_limiter::FlushLimiter;
use common_error::ext::BoxedError;
use common_query::{WriteProgress, prelude};
use common_telemetry::{debug, warn};
use operator::write_dispatcher::TrackedWrite;
use operator::write_dispatcher::WriteDispatcherRef;
use partition::manager::PartitionRuleManagerRef;
use session::context::QueryContextRef;
use snafu::{IntoError, OptionExt, ensure};

use crate::batcher::flow_notifier::FlowNotifier;
use crate::batcher::logical_table;
use crate::batcher::logical_table::batch_convert::{self, TableBatch};
use crate::batcher::logical_table::flow_notifier;
use crate::batcher::logical_table::pending_worker::FlushWaiter;
use crate::batcher::logical_table::region_write::{
    self, CatalogManagerPhysicalFlushAdapter, PartitionManagerPhysicalFlushAdapter,
    PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
    WriteDispatcherPhysicalFlushAdapter,
};
use crate::error;
use crate::error::Result;
use crate::metrics;

/// Rows and completion waiters owned by one physical flush.
pub(crate) struct Batch {
    pub(crate) table_batches: Vec<TableBatch>,
    pub(crate) total_row_count: usize,
    pub(crate) db_string: String,
    pub(crate) ctx: QueryContextRef,
    pub(crate) waiters: Vec<FlushWaiter>,
}

/// Spawns a physical flush with the configured concurrency limit.
pub(crate) async fn spawn_flush(
    flush: Batch,
    partition_manager: PartitionRuleManagerRef,
    write_dispatcher: WriteDispatcherRef,
    catalog_manager: CatalogManagerRef,
    flow_notification_tx: FlowNotifier,
    flush_limiter: FlushLimiter,
) {
    match flush_limiter.acquire().await {
        Ok(permit) => {
            tokio::spawn(async move {
                let _permit = permit;
                flush_batch_with_managers(
                    flush,
                    partition_manager,
                    write_dispatcher,
                    catalog_manager,
                    flow_notification_tx,
                )
                .await;
            });
        }
        Err(err) => {
            warn!(err; "Flush semaphore closed, flushing inline");
            flush_batch_with_managers(
                flush,
                partition_manager,
                write_dispatcher,
                catalog_manager,
                flow_notification_tx,
            )
            .await;
        }
    }
}

/// Connects shared routing and dispatch dependencies to a physical flush.
pub(crate) async fn flush_batch_with_managers(
    flush: Batch,
    partition_manager: PartitionRuleManagerRef,
    write_dispatcher: WriteDispatcherRef,
    catalog_manager: CatalogManagerRef,
    flow_notification_tx: FlowNotifier,
) {
    let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager };
    let write_dispatcher = WriteDispatcherPhysicalFlushAdapter { write_dispatcher };
    let catalog_provider = CatalogManagerPhysicalFlushAdapter { catalog_manager };
    flush_batch(
        flush,
        &partition_provider,
        &write_dispatcher,
        &catalog_provider,
        flow_notification_tx,
    )
    .await;
}

/// Flushes a batch and reports completion to its waiters.
pub(crate) async fn flush_batch(
    flush: Batch,
    partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
    write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
    catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
    flow_notification_tx: FlowNotifier,
) {
    let Batch {
        table_batches,
        total_row_count,
        db_string,
        ctx,
        waiters,
    } = flush;
    let start = Instant::now();

    // Physical-table-level flush: transform all logical table batches
    // into physical format and write them together.
    let physical_table_name = ctx
        .extension(logical_table::PHYSICAL_TABLE_KEY)
        .unwrap_or(prelude::GREPTIME_PHYSICAL_TABLE)
        .to_string();
    let result = flush_batch_physical_tracked(
        &table_batches,
        &physical_table_name,
        &ctx,
        partition_manager,
        write_dispatcher,
        catalog_manager,
        waiters.len().max(
            table_batches
                .iter()
                .flat_map(|table| &table.batches)
                .map(|batch| batch.owner + 1)
                .max()
                .unwrap_or(0),
        ),
    )
    .await;

    let elapsed = start.elapsed().as_secs_f64();
    metrics::FLUSH_ELAPSED.observe(elapsed);

    debug!(
        "Pending rows batch flushed, total rows: {}, elapsed time: {}s",
        total_row_count, elapsed
    );

    match result {
        Ok(progress) => {
            if progress.iter().any(|result| result.error.is_some()) {
                metrics::FLUSH_FAILURES.inc();
            } else {
                metrics::FLUSH_TOTAL.inc();
                metrics::FLUSH_ROWS.observe(total_row_count as f64);
            }
            let affected_rows = progress.iter().map(|result| result.affected_rows).sum();
            metrics::FLUSH_DROPPED_ROWS
                .inc_by(total_row_count.saturating_sub(affected_rows) as u64);
            let successful = progress
                .iter()
                .enumerate()
                .filter_map(|(owner, result)| result.error.is_none().then_some(owner))
                .collect::<HashSet<_>>();
            for (waiter, result) in waiters.into_iter().zip(progress) {
                operator::metrics::DIST_INGEST_ROW_COUNT
                    .with_label_values(&[db_string.as_str()])
                    .inc_by(result.affected_rows as u64);
                let _ = waiter.response_tx.send(result);
            }
            let accepted = table_batches
                .into_iter()
                .filter_map(|mut table| {
                    table
                        .batches
                        .retain(|batch| successful.contains(&batch.owner));
                    table.row_count = table
                        .batches
                        .iter()
                        .map(|batch| batch.batch.num_rows())
                        .sum();
                    (!table.batches.is_empty()).then_some(table)
                })
                .collect();
            flow_notifier::enqueue_flow_notifications(accepted, &flow_notification_tx);
        }
        Err(err) => {
            metrics::FLUSH_FAILURES.inc();
            metrics::FLUSH_DROPPED_ROWS.inc_by(total_row_count as u64);
            notify_waiters(waiters, Err(err));
        }
    }
}

/// Flushes a batch of logical table rows by transforming them into the physical table format
/// and writing them to the appropriate datanode regions.
///
/// This function performs the end-to-end physical flush pipeline:
/// 1. Resolves the physical table metadata and column ID mapping.
/// 2. Fetches the physical table's partition rule.
/// 3. Transforms each logical table batch into the physical (sparse primary key) format.
/// 4. Concatenates all transformed batches into a single combined batch.
/// 5. Splits the combined batch by partition rule and sends region write requests
///    concurrently to the target datanodes.
pub async fn flush_batch_physical(
    table_batches: &[TableBatch],
    physical_table_name: &str,
    ctx: &QueryContextRef,
    partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
    write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
    catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
) -> Result<usize> {
    let inputs = table_batches
        .iter()
        .flat_map(|table| &table.batches)
        .map(|batch| batch.owner + 1)
        .max()
        .unwrap_or(1);
    let progress = flush_batch_physical_tracked(
        table_batches,
        physical_table_name,
        ctx,
        partition_manager,
        write_dispatcher,
        catalog_manager,
        inputs,
    )
    .await?;
    let mut result = WriteProgress::default();
    for part in progress {
        result.merge(part);
    }
    result
        .into_result()
        .map_err(|source| error::ExecuteGrpcQuerySnafu.into_error(source))
}

async fn flush_batch_physical_tracked(
    table_batches: &[TableBatch],
    physical_table_name: &str,
    ctx: &QueryContextRef,
    partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized),
    write_dispatcher: &(impl PhysicalFlushWriteDispatcher + ?Sized),
    catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized),
    inputs: usize,
) -> Result<Vec<WriteProgress>> {
    // 1. Resolve the physical table and get column ID mapping
    let physical_table = {
        let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
            .with_label_values(&["flush_physical_resolve_table"])
            .start_timer();
        catalog_manager
            .physical_table(
                ctx.current_catalog(),
                &ctx.current_schema(),
                physical_table_name,
                ctx.as_ref(),
            )
            .await?
            .with_context(|| error::InternalSnafu {
                err_msg: format!(
                    "Physical table '{}' not found during pending flush",
                    physical_table_name
                ),
            })?
    };

    let physical_table_info = physical_table.table_info;
    let name_to_ids = physical_table
        .col_name_to_ids
        .with_context(|| error::InternalSnafu {
            err_msg: format!(
                "Physical table '{}' has no column IDs for pending flush",
                physical_table_name
            ),
        })?;

    // 2. Get the physical table's partition rule (one lookup instead of N)
    let (partition_rule, partition_versions) = {
        let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
            .with_label_values(&["flush_physical_fetch_partition_rule"])
            .start_timer();
        partition_manager
            .find_table_partition_rule(physical_table_info.as_ref())
            .await?
    };
    let partition_columns = partition_rule.partition_columns();
    let partition_columns_set: HashSet<&str> =
        partition_columns.iter().map(String::as_str).collect();

    // 3. Transform each logical table batch into physical format
    let modified_batches = batch_convert::transform_logical_batches_to_physical(
        table_batches,
        &name_to_ids,
        &partition_columns_set,
    )?;

    // 4. Concatenate all modified batches (all share the same physical schema)
    let combined_batch = batch_convert::concat_modified_batches(&modified_batches)?;
    // 5. Split by physical partition rule and send to regions
    let physical_table_id = physical_table_info.table_id();
    let row_count = table_batches
        .iter()
        .flat_map(|table| &table.batches)
        .map(|batch| batch.batch.num_rows())
        .sum::<usize>();
    ensure!(
        row_count == combined_batch.num_rows(),
        error::InternalSnafu {
            err_msg: "Physical transformation changed row ownership"
        }
    );
    let masks =
        region_write::split_combined_batch_by_region(&combined_batch, partition_rule.as_ref())?;
    let mut writes = Vec::new();
    for (region, mask) in masks {
        let mut selected = Vec::with_capacity(mask.selected_rows());
        let bits = mask.array().values();
        let mut offset = 0;
        // Filtering preserves batch order, and each original batch has one owner.
        for batch in table_batches.iter().flat_map(|table| &table.batches) {
            let num_rows = batch.batch.num_rows();
            let selected_rows = if num_rows == 1 {
                // Avoid bitmap range setup for single-row submissions.
                usize::from(bits.value(offset))
            } else {
                // `offset` is mask-relative; the raw buffer also needs the mask's slice offset.
                bits.inner()
                    .count_set_bits_offset(bits.offset() + offset, num_rows)
            };
            selected.extend(std::iter::repeat_n(batch.owner, selected_rows));
            offset += num_rows;
        }
        if let Some(batch) =
            region_write::plan_region_batch(&combined_batch, physical_table_id, region, &mask)?
        {
            let request = region_write::encode_region_write_request(
                batch,
                ctx.skip_wal(),
                &partition_versions,
            )?;
            writes.push(TrackedWrite {
                request,
                owners: selected,
            });
        }
    }
    let progress = write_dispatcher.dispatch_tracked(writes, inputs).await;
    // A submission can contribute several logical batches; credit its counter once.
    let counters = table_batches
        .iter()
        .flat_map(|table| &table.batches)
        .filter_map(|batch| {
            batch
                .progress
                .as_ref()
                .map(|counter| (batch.owner, counter))
        })
        .collect::<std::collections::HashMap<_, _>>();
    for (owner, counter) in counters {
        counter.add_written(progress[owner].affected_rows);
    }
    Ok(progress)
}

pub(in crate::batcher::logical_table) fn notify_waiters(
    waiters: Vec<FlushWaiter>,
    result: Result<()>,
) {
    let progress = match result {
        Ok(()) => WriteProgress::default(),
        Err(error) => WriteProgress::failed(BoxedError::new(error)),
    };
    for waiter in waiters {
        let _ = waiter.response_tx.send(progress.clone());
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::future::poll_fn;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::task::Poll;
    use std::time::Duration;

    use api::v1::meta::Peer;
    use api::v1::region::RegionRequest;
    use arrow::record_batch::RecordBatch;
    use async_trait::async_trait;
    use catalog::error::Result as CatalogResult;
    use common_batcher::request_limiter::RequestLimiter;
    use common_meta::cache::TableFlownodeSetCacheRef;
    use common_meta::node_manager::NodeManagerRef;
    use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema};
    use partition::error::Result as PartitionResult;
    use partition::partition::{PartitionRule, PartitionRuleRef, RegionMask};
    use table::metadata::TableId;
    use table::test_util::table_info::test_table_info;
    use tokio::sync::{mpsc, oneshot};

    use crate::batcher::flow_notifier::{FlowNotifier, start_flow_notification_worker};
    use crate::batcher::logical_table::batch::{
        Batch, flush_batch, flush_batch_physical, notify_waiters,
    };
    use crate::batcher::logical_table::batch_convert::TableBatch;
    use crate::batcher::logical_table::pending_worker::FlushWaiter;
    use crate::batcher::logical_table::region_write::{
        PhysicalFlushCatalogProvider, PhysicalFlushPartitionProvider, PhysicalFlushWriteDispatcher,
        PhysicalTableMetadata,
    };
    use crate::batcher::logical_table::test_util::{
        FlowNotificationMockNodeManager, RecordingFlownode, mock_aligned_tag_batch,
    };
    use crate::batcher::test_util::mock_table_flownode_cache;
    use crate::error;
    use crate::error::Error;
    use crate::metrics::{self, FLOW_NOTIFICATION_DROPPED};

    #[tokio::test]
    async fn test_flush_batch_notifies_flownode_after_successful_physical_write() {
        let table_id = 42;
        let peer = Peer {
            id: 7,
            addr: "flow-7".to_string(),
        };
        let cache = mock_table_flownode_cache(table_id, vec![(0, peer.clone()), (1, peer)]).await;
        let (requests_tx, mut requests_rx) = mpsc::unbounded_channel();
        let _requests_tx = requests_tx.clone();
        let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager {
            flownode: Arc::new(RecordingFlownode { requests_tx }),
        });
        let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone());
        let ctx = session::context::QueryContext::arc();
        let table_batches = vec![TableBatch {
            table_name: "cpu".to_string(),
            table_id,
            batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)],
            row_count: 1,
        }];
        let writes = Arc::new(AtomicUsize::new(0));

        flush_batch(
            Batch {
                table_batches,
                total_row_count: 1,
                db_string: ctx.get_db_string(),
                ctx,
                waiters: Vec::new(),
            },
            &MockFlushPartitionProvider {
                partition_rule_calls: Arc::new(AtomicUsize::new(0)),
            },
            &MockFlushWriteDispatcher {
                writes: writes.clone(),
                fail: false,
            },
            &MockFlushCatalogProvider {
                table: Some(mock_physical_table_metadata(1024)),
            },
            flow_notification_tx,
        )
        .await;

        assert_eq!(1, writes.load(Ordering::SeqCst));
        let requests = tokio::time::timeout(Duration::from_secs(1), requests_rx.recv())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(
            vec![api::v1::flow::DirtyWindowRequest {
                table_id,
                timestamps: vec![1000],
                time_ranges: Vec::new(),
            }],
            requests.requests
        );
    }

    #[tokio::test]
    async fn test_flush_batch_does_not_notify_flownode_after_physical_write_error() {
        let table_id = 42;
        let peer = Peer {
            id: 7,
            addr: "flow-7".to_string(),
        };
        let cache = mock_table_flownode_cache(table_id, vec![(0, peer.clone()), (1, peer)]).await;
        let (requests_tx, mut requests_rx) = mpsc::unbounded_channel();
        let _requests_tx = requests_tx.clone();
        let flow_node_manager: NodeManagerRef = Arc::new(FlowNotificationMockNodeManager {
            flownode: Arc::new(RecordingFlownode { requests_tx }),
        });
        let flow_notification_tx = mock_flow_notification_sender(cache, flow_node_manager.clone());
        let ctx = session::context::QueryContext::arc();
        let table_batches = vec![TableBatch {
            table_name: "cpu".to_string(),
            table_id,
            batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)],
            row_count: 1,
        }];
        let writes = Arc::new(AtomicUsize::new(0));

        flush_batch(
            Batch {
                table_batches,
                total_row_count: 1,
                db_string: ctx.get_db_string(),
                ctx,
                waiters: Vec::new(),
            },
            &MockFlushPartitionProvider {
                partition_rule_calls: Arc::new(AtomicUsize::new(0)),
            },
            &MockFlushWriteDispatcher {
                writes: writes.clone(),
                fail: true,
            },
            &MockFlushCatalogProvider {
                table: Some(mock_physical_table_metadata(1024)),
            },
            flow_notification_tx,
        )
        .await;

        assert_eq!(1, writes.load(Ordering::SeqCst));
        assert!(
            tokio::time::timeout(Duration::from_millis(50), requests_rx.recv())
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_cancelled_waiter_retains_request_slot_until_notification() {
        let limiter = RequestLimiter::try_new(1).unwrap();
        let (response_tx, response_rx) = oneshot::channel();
        let waiter = FlushWaiter {
            _completion: None,
            response_tx,
            _permit: limiter.acquire().await.unwrap(),
        };
        drop(response_rx);
        let next = limiter.acquire();
        tokio::pin!(next);
        assert!(poll_fn(|cx| Poll::Ready(next.as_mut().poll(cx).is_pending())).await);
        notify_waiters(vec![waiter], Ok(()));
        let _permit = next.await.unwrap();
    }

    #[tokio::test]
    async fn test_flush_batch_physical_uses_mockable_trait_dependencies() {
        let table_batches = vec![TableBatch {
            table_name: "t1".to_string(),
            table_id: 11,
            batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)],
            row_count: 1,
        }];
        let partition_calls = Arc::new(AtomicUsize::new(0));
        let dispatcher = MockFlushWriteDispatcher::default();
        let ctx = session::context::QueryContext::arc();

        flush_batch_physical(
            &table_batches,
            "phy",
            &ctx,
            &MockFlushPartitionProvider {
                partition_rule_calls: partition_calls.clone(),
            },
            &dispatcher,
            &MockFlushCatalogProvider {
                table: Some(mock_physical_table_metadata(1024)),
            },
        )
        .await
        .unwrap();

        assert_eq!(1, partition_calls.load(Ordering::SeqCst));
        assert_eq!(1, dispatcher.writes.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_flush_batch_physical_rejects_unattributable_affected_rows() {
        let table_batches = vec![TableBatch {
            table_name: "t1".to_string(),
            table_id: 11,
            batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)],
            row_count: 1,
        }];
        let ctx = session::context::QueryContext::arc();

        let error = flush_batch_physical(
            &table_batches,
            "phy",
            &ctx,
            &MockFlushPartitionProvider {
                partition_rule_calls: Arc::new(AtomicUsize::new(0)),
            },
            &AffectedRowsFlushWriteDispatcher { affected_rows: 7 },
            &MockFlushCatalogProvider {
                table: Some(mock_physical_table_metadata(1024)),
            },
        )
        .await
        .unwrap_err();

        assert!(format!("{error:?}").contains("Unattributable"));
    }

    #[tokio::test]
    async fn test_exemplar_coalesced_logical_flush_retains_each_requests_progress() {
        use common_error::ext::BoxedError;
        use common_query::WriteProgress;
        use operator::write_dispatcher::TrackedWrite;

        struct ProgressDispatcher([(usize, bool); 2]);
        #[async_trait]
        impl PhysicalFlushWriteDispatcher for ProgressDispatcher {
            async fn dispatch(&self, _: Vec<RegionRequest>) -> error::Result<usize> {
                unreachable!()
            }
            async fn dispatch_tracked(
                &self,
                requests: Vec<TrackedWrite>,
                inputs: usize,
            ) -> Vec<WriteProgress> {
                assert_eq!(inputs, 2);
                assert_eq!(
                    requests
                        .iter()
                        .flat_map(|r| &r.owners)
                        .copied()
                        .collect::<Vec<_>>(),
                    vec![0, 1, 1]
                );
                self.0
                    .iter()
                    .map(|&(affected_rows, failed)| WriteProgress {
                        affected_rows,
                        error: failed.then(|| {
                            Arc::new(BoxedError::new(
                                error::InternalSnafu {
                                    err_msg: "lost acknowledgement",
                                }
                                .build(),
                            ))
                        }),
                    })
                    .collect()
            }
        }
        let ctx = session::context::QueryContext::arc();
        let ingested = operator::metrics::DIST_INGEST_ROW_COUNT
            .with_label_values(&[ctx.get_db_string().as_str()]);
        for (expected_progress, successful_flush, expected_timestamps) in [
            ([(1, false), (2, false)], true, vec![1000; 3]),
            ([(1, false), (1, true)], false, vec![1000]),
            ([(0, true), (0, true)], false, vec![]),
        ] {
            let limiter = RequestLimiter::try_new(2).unwrap();
            let mut waiters = Vec::new();
            let mut receivers = Vec::new();
            let counters =
                [0, 1].map(|_| Arc::new(session::write_progress::WriteProgress::default()));
            for counter in &counters {
                let (response_tx, response_rx) = oneshot::channel();
                waiters.push(FlushWaiter {
                    _completion: Some(counter.start()),
                    response_tx,
                    _permit: limiter.acquire().await.unwrap(),
                });
                receivers.push(response_rx);
            }
            let batches = [0, 1, 1]
                .into_iter()
                .map(|owner| {
                    let mut batch = mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0);
                    batch.owner = owner;
                    batch.progress = Some(counters[owner].clone());
                    batch
                })
                .collect();
            let (notifier, mut notifications) =
                FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
            let successes = metrics::FLUSH_TOTAL.get();
            let failures = metrics::FLUSH_FAILURES.get();
            let samples = metrics::FLUSH_ROWS.get_sample_count();
            let sum = metrics::FLUSH_ROWS.get_sample_sum();
            let dropped = metrics::FLUSH_DROPPED_ROWS.get();
            let written = ingested.get();
            flush_batch(
                Batch {
                    table_batches: vec![TableBatch {
                        table_name: "cpu".into(),
                        table_id: 42,
                        batches,
                        row_count: 3,
                    }],
                    total_row_count: 3,
                    db_string: ctx.get_db_string(),
                    ctx: ctx.clone(),
                    waiters,
                },
                &MockFlushPartitionProvider {
                    partition_rule_calls: Arc::new(AtomicUsize::new(0)),
                },
                &ProgressDispatcher(expected_progress),
                &MockFlushCatalogProvider {
                    table: Some(mock_physical_table_metadata(1024)),
                },
                notifier,
            )
            .await;
            let expected_written: usize = expected_progress.iter().map(|(rows, _)| rows).sum();
            assert_eq!(
                metrics::FLUSH_TOTAL.get(),
                successes + u64::from(successful_flush)
            );
            assert_eq!(
                metrics::FLUSH_FAILURES.get(),
                failures + u64::from(!successful_flush)
            );
            assert_eq!(
                metrics::FLUSH_ROWS.get_sample_count(),
                samples + u64::from(successful_flush)
            );
            assert_eq!(
                metrics::FLUSH_ROWS.get_sample_sum(),
                sum + if successful_flush { 3.0 } else { 0.0 }
            );
            assert_eq!(
                metrics::FLUSH_DROPPED_ROWS.get(),
                dropped + (3 - expected_written) as u64
            );
            assert_eq!(ingested.get(), written + expected_written as u64);
            for (counter, (affected_rows, _)) in counters.iter().zip(expected_progress) {
                counter.wait().await;
                assert_eq!(counter.written(), affected_rows as u64);
            }
            for (receiver, (affected_rows, failed)) in receivers.into_iter().zip(expected_progress)
            {
                let result = receiver.await.unwrap();
                assert_eq!(result.affected_rows, affected_rows);
                assert_eq!(result.error.is_some(), failed);
            }
            if !expected_timestamps.is_empty() {
                let notification = notifications.try_recv().unwrap();
                assert_eq!(notification.table_id, 42);
                assert_eq!(notification.timestamps, expected_timestamps);
            }
            assert!(notifications.try_recv().is_err());
        }
    }

    #[tokio::test]
    async fn test_flush_batch_physical_preserves_owners_with_sliced_masks() {
        use std::sync::Mutex;

        use api::v1::region::{bulk_insert_request, region_request};
        use arrow::array::{BooleanArray, Float64Array, StringArray, TimestampMillisecondArray};
        use common_grpc::flight::FlightDecoder;
        use common_query::WriteProgress;
        use operator::write_dispatcher::TrackedWrite;

        #[derive(Clone)]
        struct MaskPartitions {
            offset: usize,
            pattern: usize,
        }

        impl PartitionRule for MaskPartitions {
            fn as_any(&self) -> &dyn std::any::Any {
                self
            }

            fn partition_columns(&self) -> &[String] {
                &[]
            }

            fn find_region(&self, _: &[datatypes::prelude::Value]) -> PartitionResult<u32> {
                unreachable!()
            }

            fn split_record_batch(
                &self,
                batch: &RecordBatch,
            ) -> PartitionResult<HashMap<u32, RegionMask>> {
                Ok([false, true]
                    .into_iter()
                    .enumerate()
                    .map(|(region, selected)| {
                        let values = std::iter::repeat_n(true, self.offset)
                            .chain((0..batch.num_rows()).map(|row| {
                                let matches = match self.pattern {
                                    0 => false,
                                    1 => true,
                                    _ => row % 3 != 0,
                                };
                                matches == selected
                            }))
                            .chain([true; 8])
                            .collect::<Vec<_>>();
                        let mask = BooleanArray::from(values).slice(self.offset, batch.num_rows());
                        (region as u32, RegionMask::from(mask))
                    })
                    .collect())
            }
        }

        #[async_trait]
        impl PhysicalFlushPartitionProvider for MaskPartitions {
            async fn find_table_partition_rule(
                &self,
                _: &table::metadata::TableInfo,
            ) -> PartitionResult<(PartitionRuleRef, HashMap<u32, Option<u64>>)> {
                Ok((Arc::new(self.clone()), HashMap::new()))
            }
        }

        #[derive(Default)]
        struct RecordingDispatcher(Mutex<Vec<TrackedWrite>>);

        #[async_trait]
        impl PhysicalFlushWriteDispatcher for RecordingDispatcher {
            async fn dispatch(&self, _: Vec<RegionRequest>) -> error::Result<usize> {
                unreachable!()
            }

            async fn dispatch_tracked(
                &self,
                requests: Vec<TrackedWrite>,
                inputs: usize,
            ) -> Vec<WriteProgress> {
                let mut progress = vec![WriteProgress::default(); inputs];
                for owner in requests.iter().flat_map(|request| &request.owners) {
                    progress[*owner].affected_rows += 1;
                }
                *self.0.lock().unwrap() = requests;
                progress
            }
        }

        let ctx = session::context::QueryContext::arc();
        let catalog = MockFlushCatalogProvider {
            table: Some(mock_physical_table_metadata(1024)),
        };
        for lengths in [[0; 7], [1; 7], [0, 3, 65, 1, 127, 0, 64]] {
            let mut expected = Vec::new();
            let tables = lengths
                .chunks(4)
                .enumerate()
                .map(|(table, lengths)| {
                    let batches = lengths
                        .iter()
                        .enumerate()
                        .map(|(index, &rows)| {
                            let owner = [2, 0, 1, 0, 2, 1, 0][table * 4 + index];
                            let start = expected.len();
                            expected.extend(std::iter::repeat_n(owner, rows));
                            let mut batch = mock_aligned_tag_batch("tag1", "host-1", 0, 1.0);
                            batch.batch = RecordBatch::try_new(
                                batch.batch.schema(),
                                vec![
                                    Arc::new(TimestampMillisecondArray::from_iter_values(
                                        (start..start + rows).map(|row| row as i64),
                                    )),
                                    Arc::new(Float64Array::from(vec![1.0; rows])),
                                    Arc::new(StringArray::from(vec!["host-1"; rows])),
                                ],
                            )
                            .unwrap();
                            batch.owner = owner;
                            batch
                        })
                        .collect();
                    TableBatch {
                        table_name: format!("logical_{table}"),
                        table_id: 100 + table as u32,
                        batches,
                        row_count: lengths.iter().sum(),
                    }
                })
                .collect::<Vec<_>>();
            for offset in [0, 1, 7, 65] {
                for pattern in 0..3 {
                    let dispatcher = RecordingDispatcher::default();
                    let affected = flush_batch_physical(
                        &tables,
                        "phy",
                        &ctx,
                        &MaskPartitions { offset, pattern },
                        &dispatcher,
                        &catalog,
                    )
                    .await
                    .unwrap();
                    assert_eq!(affected, expected.len());
                    let mut actual = Vec::new();
                    for write in dispatcher.0.into_inner().unwrap() {
                        let Some(region_request::Body::BulkInsert(request)) = write.request.body
                        else {
                            panic!("expected bulk request")
                        };
                        let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
                            panic!("expected Arrow IPC")
                        };
                        let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
                            .unwrap()
                            .try_decode_record_batch(&ipc.data_header, &ipc.payload)
                            .unwrap();
                        assert_eq!(batch.num_rows(), write.owners.len());
                        let timestamps = batch
                            .column_by_name("greptime_timestamp")
                            .unwrap()
                            .as_any()
                            .downcast_ref::<TimestampMillisecondArray>()
                            .unwrap();
                        assert!(timestamps.values().windows(2).all(|pair| pair[0] < pair[1]));
                        actual.extend(timestamps.values().iter().copied().zip(write.owners));
                    }
                    actual.sort_unstable();
                    assert_eq!(
                        actual,
                        expected
                            .iter()
                            .enumerate()
                            .map(|(row, owner)| (row as i64, *owner))
                            .collect::<Vec<_>>(),
                    );
                }
            }
        }
    }

    #[tokio::test]
    async fn test_flush_batch_physical_stops_before_partition_and_dispatch_when_table_missing() {
        let table_batches = vec![TableBatch {
            table_name: "t1".to_string(),
            table_id: 11,
            batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)],
            row_count: 1,
        }];
        let partition_calls = Arc::new(AtomicUsize::new(0));
        let dispatcher = MockFlushWriteDispatcher::default();
        let ctx = session::context::QueryContext::arc();

        let err = flush_batch_physical(
            &table_batches,
            "missing_phy",
            &ctx,
            &MockFlushPartitionProvider {
                partition_rule_calls: partition_calls.clone(),
            },
            &dispatcher,
            &MockFlushCatalogProvider { table: None },
        )
        .await
        .unwrap_err();

        assert!(
            err.to_string()
                .contains("Physical table 'missing_phy' not found")
        );
        assert_eq!(0, partition_calls.load(Ordering::SeqCst));
        assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_flush_batch_physical_aborts_immediately_on_transform_error() {
        let table_batches = vec![
            TableBatch {
                table_name: "broken".to_string(),
                table_id: 11,
                batches: vec![mock_aligned_tag_batch("unknown_tag", "host-1", 1000, 1.0)],
                row_count: 1,
            },
            TableBatch {
                table_name: "healthy".to_string(),
                table_id: 12,
                batches: vec![mock_aligned_tag_batch("tag1", "host-2", 2000, 2.0)],
                row_count: 1,
            },
        ];
        let partition_calls = Arc::new(AtomicUsize::new(0));
        let dispatcher = MockFlushWriteDispatcher::default();
        let ctx = session::context::QueryContext::arc();

        let err = flush_batch_physical(
            &table_batches,
            "phy",
            &ctx,
            &MockFlushPartitionProvider {
                partition_rule_calls: partition_calls.clone(),
            },
            &dispatcher,
            &MockFlushCatalogProvider {
                table: Some(mock_physical_table_metadata(1024)),
            },
        )
        .await
        .unwrap_err();

        assert!(err.to_string().contains("unknown_tag"));
        assert_eq!(1, partition_calls.load(Ordering::SeqCst));
        assert_eq!(0, dispatcher.writes.load(Ordering::SeqCst));
    }

    fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata {
        let schema = Arc::new(
            DtSchema::try_new(vec![
                DtColumnSchema::new(
                    "__primary_key",
                    datatypes::prelude::ConcreteDataType::binary_datatype(),
                    false,
                ),
                DtColumnSchema::new(
                    "greptime_timestamp",
                    datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype(),
                    false,
                ),
                DtColumnSchema::new(
                    "greptime_value",
                    datatypes::prelude::ConcreteDataType::float64_datatype(),
                    true,
                ),
                DtColumnSchema::new(
                    "tag1",
                    datatypes::prelude::ConcreteDataType::string_datatype(),
                    true,
                ),
            ])
            .unwrap(),
        );
        let mut table_info = test_table_info(table_id, "phy", "public", "greptime", schema);
        table_info.meta.column_ids = vec![0, 1, 2, 3];

        PhysicalTableMetadata {
            table_info: Arc::new(table_info),
            col_name_to_ids: Some(HashMap::from([("tag1".to_string(), 3)])),
        }
    }

    struct MockFlushCatalogProvider {
        table: Option<PhysicalTableMetadata>,
    }

    #[async_trait]
    impl PhysicalFlushCatalogProvider for MockFlushCatalogProvider {
        async fn physical_table(
            &self,
            _catalog: &str,
            _schema: &str,
            _table_name: &str,
            _query_ctx: &session::context::QueryContext,
        ) -> CatalogResult<Option<PhysicalTableMetadata>> {
            Ok(self.table.clone())
        }
    }

    struct SingleRegionPartitionRule;

    impl PartitionRule for SingleRegionPartitionRule {
        fn as_any(&self) -> &dyn std::any::Any {
            self
        }

        fn partition_columns(&self) -> &[String] {
            &[]
        }

        fn find_region(
            &self,
            _values: &[datatypes::prelude::Value],
        ) -> partition::error::Result<store_api::storage::RegionNumber> {
            unimplemented!()
        }

        fn split_record_batch(
            &self,
            record_batch: &RecordBatch,
        ) -> partition::error::Result<HashMap<store_api::storage::RegionNumber, RegionMask>>
        {
            Ok(HashMap::from([(
                1,
                RegionMask::new(
                    arrow::array::BooleanArray::from(vec![true; record_batch.num_rows()]),
                    record_batch.num_rows(),
                ),
            )]))
        }
    }

    struct MockFlushPartitionProvider {
        partition_rule_calls: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl PhysicalFlushPartitionProvider for MockFlushPartitionProvider {
        async fn find_table_partition_rule(
            &self,
            _table_info: &table::metadata::TableInfo,
        ) -> PartitionResult<(PartitionRuleRef, HashMap<u32, Option<u64>>)> {
            self.partition_rule_calls.fetch_add(1, Ordering::SeqCst);
            Ok((Arc::new(SingleRegionPartitionRule), HashMap::new()))
        }
    }

    #[derive(Default)]
    struct MockFlushWriteDispatcher {
        writes: Arc<AtomicUsize>,
        fail: bool,
    }

    #[async_trait]
    impl PhysicalFlushWriteDispatcher for MockFlushWriteDispatcher {
        async fn dispatch(&self, _requests: Vec<RegionRequest>) -> error::Result<usize> {
            self.writes.fetch_add(1, Ordering::SeqCst);
            if self.fail {
                return Err(Error::Internal {
                    err_msg: "physical write failed".to_string(),
                });
            }
            Ok(1)
        }
    }

    fn mock_flow_notification_sender(
        cache: TableFlownodeSetCacheRef,
        node_manager: NodeManagerRef,
    ) -> FlowNotifier {
        let (tx, rx) = FlowNotifier::try_new(16, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
        start_flow_notification_worker(rx, cache, node_manager);
        tx
    }

    #[derive(Default)]
    struct AffectedRowsFlushWriteDispatcher {
        affected_rows: usize,
    }

    #[async_trait]
    impl PhysicalFlushWriteDispatcher for AffectedRowsFlushWriteDispatcher {
        async fn dispatch(&self, _requests: Vec<RegionRequest>) -> error::Result<usize> {
            Ok(self.affected_rows)
        }
    }
}