servers/batcher/logical_table/
region_write.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16
17use api::v1::region::{
18    BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request,
19};
20use api::v1::{ArrowIpc, PartitionExprVersion};
21use arrow::compute;
22use arrow::record_batch::RecordBatch;
23use async_trait::async_trait;
24use catalog::CatalogManagerRef;
25use common_error::ext::BoxedError;
26use common_grpc::error::Error as GrpcError;
27use common_grpc::flight;
28use common_query::WriteProgress;
29use common_telemetry::tracing_context::TracingContext;
30use operator::write_dispatcher::TrackedWrite;
31use operator::write_dispatcher::WriteDispatcherRef;
32use partition::manager::PartitionRuleManagerRef;
33use partition::partition::PartitionRuleRef;
34use snafu::ResultExt;
35use store_api::storage::RegionId;
36use table::metadata::{TableId, TableInfo, TableInfoRef};
37
38use crate::error;
39use crate::error::{Error, Result};
40use crate::metrics;
41
42#[derive(Clone)]
43pub struct PhysicalTableMetadata {
44    pub table_info: TableInfoRef,
45    /// Mapping from column name to column id
46    pub col_name_to_ids: Option<HashMap<String, u32>>,
47}
48
49#[async_trait]
50pub trait PhysicalFlushCatalogProvider: Send + Sync {
51    async fn physical_table(
52        &self,
53        catalog: &str,
54        schema: &str,
55        table_name: &str,
56        query_ctx: &session::context::QueryContext,
57    ) -> catalog::error::Result<Option<PhysicalTableMetadata>>;
58}
59
60/// Loads the physical partition rule and its expression versions.
61#[async_trait]
62pub trait PhysicalFlushPartitionProvider: Send + Sync {
63    /// Loads a routing snapshot for the physical table.
64    async fn find_table_partition_rule(
65        &self,
66        table_info: &TableInfo,
67    ) -> partition::error::Result<(PartitionRuleRef, HashMap<u32, Option<u64>>)>;
68}
69
70/// Dispatches the region requests produced by a physical flush.
71#[async_trait]
72pub trait PhysicalFlushWriteDispatcher: Send + Sync {
73    /// Dispatches one physical flush together, sharing routing refreshes across its writes.
74    async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<usize>;
75
76    async fn dispatch_tracked(
77        &self,
78        requests: Vec<TrackedWrite>,
79        inputs: usize,
80    ) -> Vec<WriteProgress> {
81        let mut progress = vec![WriteProgress::default(); inputs];
82        let total = requests.iter().map(|r| r.owners.len()).sum::<usize>();
83        let owners = requests
84            .iter()
85            .flat_map(|r| r.owners.iter().copied())
86            .collect::<Vec<_>>();
87        match self
88            .dispatch(requests.into_iter().map(|r| r.request).collect())
89            .await
90        {
91            Ok(rows) if rows == total => {
92                for owner in owners {
93                    progress[owner].affected_rows += 1;
94                }
95            }
96            Ok(_) => {
97                return vec![
98                    WriteProgress::failed(BoxedError::new(
99                        error::InternalSnafu {
100                            err_msg: "Unattributable physical write acknowledgement"
101                        }
102                        .build()
103                    ));
104                    inputs
105                ];
106            }
107            Err(error) => return vec![WriteProgress::failed(BoxedError::new(error)); inputs],
108        }
109        progress
110    }
111}
112
113#[derive(Clone)]
114pub(in crate::batcher::logical_table) struct CatalogManagerPhysicalFlushAdapter {
115    pub(in crate::batcher::logical_table) catalog_manager: CatalogManagerRef,
116}
117
118#[async_trait]
119impl PhysicalFlushCatalogProvider for CatalogManagerPhysicalFlushAdapter {
120    async fn physical_table(
121        &self,
122        catalog: &str,
123        schema: &str,
124        table_name: &str,
125        query_ctx: &session::context::QueryContext,
126    ) -> catalog::error::Result<Option<PhysicalTableMetadata>> {
127        self.catalog_manager
128            .table(catalog, schema, table_name, Some(query_ctx))
129            .await
130            .map(|table| {
131                table.map(|table| {
132                    let table_info = table.table_info();
133                    let name_to_ids = table_info.name_to_ids();
134                    PhysicalTableMetadata {
135                        table_info,
136                        col_name_to_ids: name_to_ids,
137                    }
138                })
139            })
140    }
141}
142
143#[derive(Clone)]
144pub(in crate::batcher::logical_table) struct PartitionManagerPhysicalFlushAdapter {
145    pub(in crate::batcher::logical_table) partition_manager: PartitionRuleManagerRef,
146}
147
148#[async_trait]
149impl PhysicalFlushPartitionProvider for PartitionManagerPhysicalFlushAdapter {
150    async fn find_table_partition_rule(
151        &self,
152        table_info: &TableInfo,
153    ) -> partition::error::Result<(PartitionRuleRef, HashMap<u32, Option<u64>>)> {
154        self.partition_manager
155            .find_table_partition_rule(table_info)
156            .await
157    }
158}
159
160/// Sends a physical flush through the shared write dispatcher.
161#[derive(Clone)]
162pub(crate) struct WriteDispatcherPhysicalFlushAdapter {
163    pub(crate) write_dispatcher: WriteDispatcherRef,
164}
165
166#[async_trait]
167impl PhysicalFlushWriteDispatcher for WriteDispatcherPhysicalFlushAdapter {
168    async fn dispatch_tracked(
169        &self,
170        requests: Vec<TrackedWrite>,
171        inputs: usize,
172    ) -> Vec<WriteProgress> {
173        self.write_dispatcher
174            .dispatch_tracked(
175                requests,
176                inputs,
177                Some(
178                    &metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
179                        .with_label_values(&["flush_write_region"]),
180                ),
181            )
182            .await
183    }
184    async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<usize> {
185        self.write_dispatcher
186            .dispatch_with_metrics(
187                requests,
188                metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
189                    .with_label_values(&["flush_write_region"]),
190            )
191            .await
192            .context(error::CommonMetaSnafu)
193            .map(|response| response.affected_rows)
194    }
195}
196
197/// Arrow rows selected for one physical region.
198pub(crate) struct PlannedRegionBatch {
199    pub(crate) region_id: RegionId,
200    pub(crate) batch: RecordBatch,
201}
202
203#[cfg(test)]
204impl PlannedRegionBatch {
205    pub(crate) fn num_rows(&self) -> usize {
206        self.batch.num_rows()
207    }
208}
209
210pub(in crate::batcher::logical_table) fn split_combined_batch_by_region(
211    combined_batch: &RecordBatch,
212    partition_rule: &dyn partition::partition::PartitionRule,
213) -> Result<HashMap<u32, partition::partition::RegionMask>> {
214    let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
215        .with_label_values(&["flush_physical_split_record_batch"])
216        .start_timer();
217    let map = partition_rule.split_record_batch(combined_batch)?;
218    Ok(map)
219}
220
221/// Selects the rows assigned to one physical region.
222pub(crate) fn plan_region_batch(
223    combined_batch: &RecordBatch,
224    physical_table_id: TableId,
225    region_number: u32,
226    mask: &partition::partition::RegionMask,
227) -> Result<Option<PlannedRegionBatch>> {
228    if mask.select_none() {
229        return Ok(None);
230    }
231
232    let region_batch = if mask.select_all() {
233        combined_batch.clone()
234    } else {
235        let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
236            .with_label_values(&["flush_physical_filter_record_batch"])
237            .start_timer();
238        compute::filter_record_batch(combined_batch, mask.array()).context(error::ArrowSnafu)?
239    };
240
241    let row_count = region_batch.num_rows();
242    if row_count == 0 {
243        return Ok(None);
244    }
245
246    Ok(Some(PlannedRegionBatch {
247        region_id: RegionId::new(physical_table_id, region_number),
248        batch: region_batch,
249    }))
250}
251
252/// Partitions a physical flush while retaining columns needed for rerouting.
253#[cfg(test)]
254pub(crate) fn plan_region_batches(
255    combined_batch: RecordBatch,
256    physical_table_id: TableId,
257    partition_rule: &dyn partition::partition::PartitionRule,
258) -> Result<Vec<PlannedRegionBatch>> {
259    let region_masks = split_combined_batch_by_region(&combined_batch, partition_rule)?;
260
261    // Retain partition columns in the RPC payload so routing can be rebuilt after repartition.
262    // Mito's sparse bulk conversion projects the stored fields and encoded primary key.
263    let mut planned_batches = Vec::new();
264    for (region_number, mask) in region_masks {
265        if let Some(planned_batch) =
266            plan_region_batch(&combined_batch, physical_table_id, region_number, &mask)?
267        {
268            planned_batches.push(planned_batch);
269        }
270    }
271
272    Ok(planned_batches)
273}
274
275/// Encodes a region batch with the partition version used to select its rows.
276pub(crate) fn encode_region_write_request(
277    planned: PlannedRegionBatch,
278    skip_wal: bool,
279    partition_versions: &HashMap<u32, Option<u64>>,
280) -> Result<RegionRequest> {
281    let region_id = planned.region_id;
282    let (schema_bytes, data_header, payload) = {
283        let _timer = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
284            .with_label_values(&["flush_physical_encode_ipc"])
285            .start_timer();
286        flight::record_batch_to_ipc(planned.batch).map_err(map_ipc_error)?
287    };
288
289    Ok(RegionRequest {
290        header: Some(RegionRequestHeader {
291            tracing_context: TracingContext::from_current_span().to_w3c(),
292            ..Default::default()
293        }),
294        body: Some(region_request::Body::BulkInsert(BulkInsertRequest {
295            skip_wal,
296            region_id: region_id.as_u64(),
297            partition_expr_version: partition_versions
298                .get(&region_id.region_number())
299                .copied()
300                .flatten()
301                .map(|value| PartitionExprVersion { value }),
302            // Set aligned_schema_version to None so that datanode will check the batch schema again to see if any
303            // column is missing.
304            aligned_schema_version: None,
305            body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
306                schema: schema_bytes,
307                data_header,
308                payload,
309            })),
310        })),
311    })
312}
313
314pub(in crate::batcher::logical_table) fn map_ipc_error(error: GrpcError) -> Error {
315    match error {
316        GrpcError::NotSupported { feat } => Error::NotSupported { feat },
317        GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason },
318        error => Error::Internal {
319            err_msg: error.to_string(),
320        },
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use std::collections::{HashMap, HashSet};
327    use std::sync::Arc;
328
329    use api::region::RegionResponse;
330    use api::v1::PartitionExprVersion;
331    use api::v1::region::{
332        InsertRequest, InsertRequests, RegionRequest, bulk_insert_request, region_request,
333    };
334    use arrow::array::{Array, BinaryArray, BooleanArray, StringArray, TimestampMillisecondArray};
335    use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
336    use arrow::record_batch::RecordBatch;
337    use common_grpc::flight::FlightDecoder;
338    use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher;
339    use common_meta::test_util::MockDatanodeManager;
340    use operator::test_util::{create_partition_rule_manager, prepare_mocked_backend};
341    use operator::write_dispatcher::{TrackedWrite, WriteDispatcher};
342    use partition::partition::{PartitionRule, RegionMask};
343    use store_api::storage::RegionId;
344
345    use crate::batcher::logical_table::batch_convert::{self, RecordBatchWithTsIdx, TableBatch};
346    use crate::batcher::logical_table::region_write::{
347        PhysicalFlushWriteDispatcher, PlannedRegionBatch, WriteDispatcherPhysicalFlushAdapter,
348        encode_region_write_request, plan_region_batches,
349    };
350    use crate::metrics;
351
352    #[tokio::test]
353    async fn test_tracked_dispatch_records_flush_write_region_metric() {
354        let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await;
355        let (sender, _receiver) = tokio::sync::mpsc::channel(1);
356        let handler = DatanodeWatcher::new(sender).with_handler(|_, _| Ok(RegionResponse::new(1)));
357        let dispatcher = WriteDispatcherPhysicalFlushAdapter {
358            write_dispatcher: Arc::new(WriteDispatcher::new(
359                Arc::new(MockDatanodeManager::new(handler)),
360                partitions,
361            )),
362        };
363        let histogram = metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
364            .with_label_values(&["flush_write_region"]);
365        let before = histogram.get_sample_count();
366        let progress = dispatcher
367            .dispatch_tracked(
368                vec![TrackedWrite {
369                    request: RegionRequest {
370                        body: Some(region_request::Body::Inserts(InsertRequests {
371                            requests: vec![InsertRequest {
372                                region_id: RegionId::new(1, 1).as_u64(),
373                                ..Default::default()
374                            }],
375                        })),
376                        ..Default::default()
377                    },
378                    owners: vec![0],
379                }],
380                1,
381            )
382            .await;
383
384        assert_eq!(progress[0].affected_rows, 1);
385        assert!(progress[0].error.is_none());
386        assert_eq!(histogram.get_sample_count(), before + 1);
387    }
388
389    #[test]
390    fn test_plan_region_batches_preserves_partition_columns_through_encoding() {
391        let logical_batch = RecordBatch::try_new(
392            Arc::new(ArrowSchema::new(vec![
393                Field::new(
394                    "greptime_timestamp",
395                    ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
396                    false,
397                ),
398                Field::new("greptime_value", ArrowDataType::Float64, true),
399                Field::new("host", ArrowDataType::Utf8, true),
400                Field::new("zone", ArrowDataType::Utf8, true),
401            ])),
402            vec![
403                Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000_i64])),
404                Arc::new(arrow::array::Float64Array::from(vec![1.0_f64, 2.0_f64])),
405                Arc::new(StringArray::from(vec!["node-1", "node-2"])),
406                Arc::new(StringArray::from(vec!["east", "west"])),
407            ],
408        )
409        .unwrap();
410        let table_batches = [TableBatch {
411            table_name: "metric".to_string(),
412            table_id: 2048,
413            batches: vec![RecordBatchWithTsIdx::try_new(logical_batch, 0).unwrap()],
414            row_count: 2,
415        }];
416        let modified = batch_convert::transform_logical_batches_to_physical(
417            &table_batches,
418            &HashMap::from([("host".to_string(), 1), ("zone".to_string(), 2)]),
419            &HashSet::from(["host"]),
420        )
421        .unwrap();
422        let combined_batch = batch_convert::concat_modified_batches(&modified).unwrap();
423        assert!(combined_batch.column_by_name("host").is_some());
424        assert!(combined_batch.column_by_name("zone").is_none());
425        let expected = combined_batch.clone();
426        let mut planned_batches = plan_region_batches(
427            combined_batch,
428            1024,
429            &TwoRegionPartitionRule {
430                partition_columns: vec!["host".to_string()],
431            },
432        )
433        .unwrap();
434        planned_batches.sort_by_key(|planned| planned.region_id.region_number());
435
436        assert_eq!(2, planned_batches.len());
437        assert_eq!(RegionId::new(1024, 1), planned_batches[0].region_id);
438        assert_eq!(1, planned_batches[0].num_rows());
439        assert_eq!(4, planned_batches[0].batch.num_columns());
440        assert_eq!(RegionId::new(1024, 2), planned_batches[1].region_id);
441        assert_eq!(1, planned_batches[1].num_rows());
442        assert_eq!(4, planned_batches[1].batch.num_columns());
443
444        let partition_versions = HashMap::from([(1, Some(7)), (2, Some(8))]);
445        for (index, planned) in planned_batches.into_iter().enumerate() {
446            let write = encode_region_write_request(planned, false, &partition_versions).unwrap();
447            let Some(region_request::Body::BulkInsert(request)) = write.body else {
448                panic!("expected bulk")
449            };
450            assert_eq!(
451                request.partition_expr_version.unwrap().value,
452                7 + index as u64
453            );
454            let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
455                panic!("expected IPC")
456            };
457            let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
458                .unwrap()
459                .try_decode_record_batch(&ipc.data_header, &ipc.payload)
460                .unwrap();
461            // The partition-only Arrow projection must retain the complete sparse key,
462            // including the zone tag that a later partition rule can require.
463            assert!(batch.column_by_name("zone").is_none());
464            assert_eq!(batch, expected.slice(index, 1));
465            let tags = batch
466                .column_by_name("host")
467                .unwrap()
468                .as_any()
469                .downcast_ref::<StringArray>()
470                .unwrap();
471            assert_eq!(tags.len(), 1);
472            assert_eq!(tags.value(0), if index == 0 { "node-1" } else { "node-2" });
473        }
474    }
475
476    #[test]
477    fn test_encode_region_write_request_builds_bulk_insert_request() {
478        check_encode_region_write_request(false);
479        check_encode_region_write_request(true);
480    }
481
482    struct TwoRegionPartitionRule {
483        partition_columns: Vec<String>,
484    }
485
486    impl PartitionRule for TwoRegionPartitionRule {
487        fn as_any(&self) -> &dyn std::any::Any {
488            self
489        }
490
491        fn partition_columns(&self) -> &[String] {
492            &self.partition_columns
493        }
494
495        fn find_region(
496            &self,
497            _values: &[datatypes::prelude::Value],
498        ) -> partition::error::Result<store_api::storage::RegionNumber> {
499            unimplemented!()
500        }
501
502        fn split_record_batch(
503            &self,
504            _record_batch: &RecordBatch,
505        ) -> partition::error::Result<HashMap<store_api::storage::RegionNumber, RegionMask>>
506        {
507            Ok(HashMap::from([
508                (1, RegionMask::new(BooleanArray::from(vec![true, false]), 1)),
509                (2, RegionMask::new(BooleanArray::from(vec![false, true]), 1)),
510                (
511                    3,
512                    RegionMask::new(BooleanArray::from(vec![false, false]), 0),
513                ),
514            ]))
515        }
516    }
517
518    fn check_encode_region_write_request(skip_wal: bool) {
519        let planned_batch = PlannedRegionBatch {
520            region_id: RegionId::new(1024, 1),
521            batch: RecordBatch::try_new(
522                Arc::new(ArrowSchema::new(vec![
523                    Field::new("__primary_key", ArrowDataType::Binary, false),
524                    Field::new(
525                        "greptime_timestamp",
526                        ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
527                        false,
528                    ),
529                    Field::new("greptime_value", ArrowDataType::Float64, true),
530                ])),
531                vec![
532                    Arc::new(BinaryArray::from(vec![b"k1".as_slice()])),
533                    Arc::new(TimestampMillisecondArray::from(vec![1000_i64])),
534                    Arc::new(arrow::array::Float64Array::from(vec![1.0_f64])),
535                ],
536            )
537            .unwrap(),
538        };
539        let write =
540            encode_region_write_request(planned_batch, skip_wal, &HashMap::from([(1, Some(7))]))
541                .unwrap();
542
543        let Some(region_request::Body::BulkInsert(request)) = &write.body else {
544            panic!("expected bulk insert request");
545        };
546        assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id);
547        assert_eq!(skip_wal, request.skip_wal);
548        assert_eq!(
549            Some(PartitionExprVersion { value: 7 }),
550            request.partition_expr_version
551        );
552    }
553}
servers/batcher/logical_table/region_write.rs

servers/batcher/logical_table/
region_write.rs
