operator/req_convert/common/
partitioner.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;
16use std::sync::Arc;
17
18use api::v1::region::{
19    BulkInsertRequest, DeleteRequest, DeleteRequests, InsertRequest, InsertRequests, RegionRequest,
20    bulk_insert_request, region_request::Body,
21};
22use api::v1::{ArrowIpc, PartitionExprVersion, Rows};
23use common_error::ext::BoxedError;
24use common_grpc::flight::{self, FlightDecoder};
25use datatypes::arrow::array::{Array, BinaryArray, StringBuilder};
26use datatypes::arrow::compute;
27use datatypes::arrow::datatypes::{DataType, Field, Schema};
28use datatypes::arrow::record_batch::RecordBatch;
29use mito_codec::row_converter::sparse::{SparseOffsetsCache, SparsePrimaryKeyView};
30use partition::PartitionRuleRef;
31use partition::manager::PartitionRuleManager;
32use partition::partition::RegionMask;
33use partition::splitter::RowSplitter;
34use snafu::{OptionExt, ResultExt, ensure};
35use store_api::storage::consts;
36use store_api::storage::{RegionId, RegionNumber, TableId};
37use table::metadata::TableInfo;
38
39use crate::error::{self, Result, SplitDeleteSnafu, SplitInsertSnafu};
40
41/// Loads table routing rules and converts rows to region write requests.
42pub struct Partitioner<'a> {
43    partition_manager: &'a PartitionRuleManager,
44}
45
46/// A single snapshot of a table's partition rule and expression versions.
47pub(crate) struct TablePartitioner {
48    table_id: TableId,
49    rule: PartitionRuleRef,
50    versions: HashMap<RegionNumber, Option<PartitionExprVersion>>,
51}
52
53impl<'a> Partitioner<'a> {
54    /// Creates a partitioner backed by the shared routing caches.
55    pub fn new(partition_manager: &'a PartitionRuleManager) -> Self {
56        Self { partition_manager }
57    }
58
59    /// Routes insert rows using the current table partitions.
60    pub async fn partition_insert_requests(
61        &self,
62        table_info: &TableInfo,
63        rows: Rows,
64        skip_wal: bool,
65    ) -> Result<Vec<InsertRequest>> {
66        TablePartitioner::try_new(self.partition_manager, table_info)
67            .await
68            .context(SplitInsertSnafu)?
69            .partition_insert_requests(rows, skip_wal)
70    }
71
72    /// Routes delete rows using the current table partitions.
73    pub async fn partition_delete_requests(
74        &self,
75        table_info: &TableInfo,
76        rows: Rows,
77    ) -> Result<Vec<DeleteRequest>> {
78        TablePartitioner::try_new(self.partition_manager, table_info)
79            .await
80            .context(SplitDeleteSnafu)?
81            .partition_delete_requests(rows)
82    }
83}
84
85impl TablePartitioner {
86    /// Reads a table's partition rule and expression versions into one routing snapshot.
87    pub async fn try_new(
88        partition_manager: &PartitionRuleManager,
89        table_info: &TableInfo,
90    ) -> partition::error::Result<Self> {
91        let (rule, versions) = partition_manager
92            .find_table_partition_rule(table_info)
93            .await?;
94        Ok(Self {
95            table_id: table_info.table_id(),
96            rule,
97            versions: versions
98                .into_iter()
99                .map(|(region, version)| {
100                    (region, version.map(|value| PartitionExprVersion { value }))
101                })
102                .collect(),
103        })
104    }
105
106    /// Rebuilds write requests using this partition snapshot after a route refresh.
107    ///
108    /// Takes ownership so row payloads can move into their new regions. Requests whose
109    /// region and partition version still match retain their original payloads.
110    /// The table metadata comes from the same refresh that created this partitioner.
111    pub(crate) fn rebuild_write_requests(
112        &self,
113        request: RegionRequest,
114        table_info: &TableInfo,
115    ) -> Result<Vec<RegionRequest>> {
116        let RegionRequest { header, body } = request;
117        let bodies = match body.context(error::UnexpectedSnafu {
118            violated: "Cannot rebuild a request without a body",
119        })? {
120            Body::Inserts(inserts) => self.rebuild_insert_requests(inserts)?,
121            Body::Deletes(deletes) => self.rebuild_delete_requests(deletes)?,
122            Body::BulkInsert(bulk) => self.rebuild_bulk_insert_request(bulk, table_info)?,
123            _ => {
124                return error::UnexpectedSnafu {
125                    violated: "Only writes can be rebuilt",
126                }
127                .fail();
128            }
129        };
130        Ok(bodies
131            .into_iter()
132            .map(|body| RegionRequest {
133                header: header.clone(),
134                body: Some(body),
135            })
136            .collect())
137    }
138
139    /// Moves insert rows into requests for this snapshot's destination regions.
140    pub fn partition_insert_requests(
141        &self,
142        rows: Rows,
143        skip_wal: bool,
144    ) -> Result<Vec<InsertRequest>> {
145        Ok(RowSplitter::new(self.rule.clone())
146            .split(rows)
147            .context(SplitInsertSnafu)?
148            .into_iter()
149            .map(|(region, rows)| InsertRequest {
150                region_id: RegionId::new(self.table_id, region).as_u64(),
151                rows: Some(rows),
152                partition_expr_version: self.version(region),
153                skip_wal,
154            })
155            .collect())
156    }
157
158    /// Moves delete rows into requests for this snapshot's destination regions.
159    pub fn partition_delete_requests(&self, rows: Rows) -> Result<Vec<DeleteRequest>> {
160        Ok(RowSplitter::new(self.rule.clone())
161            .split(rows)
162            .context(SplitDeleteSnafu)?
163            .into_iter()
164            .map(|(region, rows)| DeleteRequest {
165                region_id: RegionId::new(self.table_id, region).as_u64(),
166                rows: Some(rows),
167                partition_expr_version: self.version(region),
168            })
169            .collect())
170    }
171
172    /// Splits an Arrow batch, reusing the original IPC when one region selects all rows.
173    pub fn partition_bulk_insert_requests(
174        &self,
175        batch: RecordBatch,
176        original_ipc: Option<ArrowIpc>,
177        skip_wal: bool,
178    ) -> Result<Vec<BulkInsertRequest>> {
179        let masks = self
180            .rule
181            .split_record_batch(&batch)
182            .context(SplitInsertSnafu)?;
183        self.build_bulk_insert_requests(batch, masks, original_ipc, skip_wal)
184    }
185
186    fn build_bulk_insert_requests(
187        &self,
188        batch: RecordBatch,
189        masks: HashMap<RegionNumber, RegionMask>,
190        mut original_ipc: Option<ArrowIpc>,
191        skip_wal: bool,
192    ) -> Result<Vec<BulkInsertRequest>> {
193        let mut requests = Vec::with_capacity(masks.len());
194        for (region, mask) in masks {
195            if mask.select_none() {
196                continue;
197            }
198            let original = if mask.select_all() {
199                original_ipc.take()
200            } else {
201                None
202            };
203            let ipc = if let Some(ipc) = original {
204                ipc
205            } else {
206                let selected = if mask.select_all() {
207                    batch.clone()
208                } else {
209                    compute::filter_record_batch(&batch, mask.array())
210                        .context(error::ComputeArrowSnafu)?
211                };
212                let (schema, data_header, payload) = flight::record_batch_to_ipc(selected)
213                    .map_err(BoxedError::new)
214                    .context(error::ExternalSnafu)?;
215                ArrowIpc {
216                    schema,
217                    data_header,
218                    payload,
219                }
220            };
221            requests.push(BulkInsertRequest {
222                region_id: RegionId::new(self.table_id, region).as_u64(),
223                partition_expr_version: self.version(region),
224                // The destination datanode validates against its current schema.
225                aligned_schema_version: None,
226                skip_wal,
227                body: Some(bulk_insert_request::Body::ArrowIpc(ipc)),
228            });
229        }
230        Ok(requests)
231    }
232
233    fn rebuild_insert_requests(&self, inserts: InsertRequests) -> Result<Vec<Body>> {
234        let mut bodies = Vec::new();
235        for insert in inserts.requests {
236            let region_id = RegionId::from_u64(insert.region_id);
237            if self.is_partition_expr_unchanged(region_id, insert.partition_expr_version) {
238                bodies.push(Body::Inserts(InsertRequests {
239                    requests: vec![insert],
240                }));
241                continue;
242            }
243            let rows = insert.rows.context(error::UnexpectedSnafu {
244                violated: "Cannot rebuild an insert without rows",
245            })?;
246            for request in self.partition_insert_requests(rows, insert.skip_wal)? {
247                bodies.push(Body::Inserts(InsertRequests {
248                    requests: vec![request],
249                }));
250            }
251        }
252        Ok(bodies)
253    }
254
255    fn rebuild_delete_requests(&self, deletes: DeleteRequests) -> Result<Vec<Body>> {
256        let mut bodies = Vec::new();
257        for delete in deletes.requests {
258            let region_id = RegionId::from_u64(delete.region_id);
259            if self.is_partition_expr_unchanged(region_id, delete.partition_expr_version) {
260                bodies.push(Body::Deletes(DeleteRequests {
261                    requests: vec![delete],
262                }));
263                continue;
264            }
265            let rows = delete.rows.context(error::UnexpectedSnafu {
266                violated: "Cannot rebuild a delete without rows",
267            })?;
268            for request in self.partition_delete_requests(rows)? {
269                bodies.push(Body::Deletes(DeleteRequests {
270                    requests: vec![request],
271                }));
272            }
273        }
274        Ok(bodies)
275    }
276
277    fn rebuild_bulk_insert_request(
278        &self,
279        bulk: BulkInsertRequest,
280        table_info: &TableInfo,
281    ) -> Result<Vec<Body>> {
282        let region_id = RegionId::from_u64(bulk.region_id);
283        if self.is_partition_expr_unchanged(region_id, bulk.partition_expr_version) {
284            return Ok(vec![Body::BulkInsert(bulk)]);
285        }
286        let bulk_insert_request::Body::ArrowIpc(ipc) =
287            bulk.body.context(error::UnexpectedSnafu {
288                violated: "Cannot rebuild a bulk insert without Arrow IPC data",
289            })?;
290        let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
291            .and_then(|mut decoder| decoder.try_decode_record_batch(&ipc.data_header, &ipc.payload))
292            .map_err(BoxedError::new)
293            .context(error::ExternalSnafu)?;
294        let routing_batch = self.restore_partition_tags(&batch, table_info)?;
295        let masks = self
296            .rule
297            .split_record_batch(routing_batch.as_ref().unwrap_or(&batch))
298            .context(SplitInsertSnafu)?;
299        Ok(self
300            .build_bulk_insert_requests(batch, masks, Some(ipc), bulk.skip_wal)?
301            .into_iter()
302            .map(Body::BulkInsert)
303            .collect())
304    }
305
306    /// Restores omitted physical-table tags for routing without changing the write payload.
307    fn restore_partition_tags(
308        &self,
309        batch: &RecordBatch,
310        table_info: &TableInfo,
311    ) -> Result<Option<RecordBatch>> {
312        let meta = &table_info.meta;
313        if !table_info.is_physical_table() || meta.column_ids.len() != meta.schema.num_columns() {
314            return Ok(None);
315        }
316        let missing = self
317            .rule
318            .partition_columns()
319            .iter()
320            .filter(|name| batch.column_by_name(name).is_none())
321            .filter_map(|name| {
322                let index = meta.schema.column_index_by_name(name)?;
323                (meta.primary_key_indices.contains(&index)
324                    && meta.schema.column_schemas()[index].data_type.is_string())
325                .then_some((name, meta.column_ids[index]))
326            })
327            .collect::<Vec<_>>();
328        if missing.is_empty() {
329            return Ok(None);
330        }
331        let Some(primary_key) = batch.column_by_name(consts::PRIMARY_KEY_COLUMN_NAME) else {
332            return Ok(None);
333        };
334        let primary_key =
335            primary_key
336                .as_any()
337                .downcast_ref::<BinaryArray>()
338                .context(error::UnexpectedSnafu {
339                    violated: "Sparse primary key column must be Binary",
340                })?;
341        let mut builders = missing
342            .iter()
343            .map(|_| StringBuilder::new())
344            .collect::<Vec<_>>();
345        let mut cache = SparseOffsetsCache::new();
346        let mut buffer = Vec::new();
347        for row in 0..batch.num_rows() {
348            ensure!(
349                !primary_key.is_null(row),
350                error::UnexpectedSnafu {
351                    violated: "Sparse primary key must not be null",
352                }
353            );
354            let mut key = SparsePrimaryKeyView::new(primary_key.value(row), &mut cache)
355                .map_err(BoxedError::new)
356                .context(error::ExternalSnafu)?;
357            for ((_, column_id), builder) in missing.iter().zip(&mut builders) {
358                let label = key
359                    .label(*column_id, &mut buffer)
360                    .map_err(BoxedError::new)
361                    .context(error::ExternalSnafu)?;
362                builder.append_option(label);
363            }
364        }
365        let schema = batch.schema();
366        let mut fields = schema.fields().to_vec();
367        let mut columns = batch.columns().to_vec();
368        for ((name, _), mut builder) in missing.into_iter().zip(builders) {
369            fields.push(Arc::new(Field::new(name, DataType::Utf8, true)));
370            columns.push(Arc::new(builder.finish()));
371        }
372        RecordBatch::try_new(
373            Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())),
374            columns,
375        )
376        .map(Some)
377        .context(error::ComputeArrowSnafu)
378    }
379
380    /// Whether the request's region still exists with the same partition expression version.
381    ///
382    /// A leader change leaves the partition unchanged, so the original payload can be
383    /// sent to the new peer. A removed region or changed partition expression version
384    /// requires rebuilding the request using this snapshot's partition rule.
385    fn is_partition_expr_unchanged(
386        &self,
387        region_id: RegionId,
388        version: Option<PartitionExprVersion>,
389    ) -> bool {
390        region_id.table_id() == self.table_id
391            && self.versions.get(&region_id.region_number()) == Some(&version)
392    }
393
394    fn version(&self, region: RegionNumber) -> Option<PartitionExprVersion> {
395        self.versions.get(&region).copied().flatten()
396    }
397}
398
399#[cfg(test)]
400mod tests {
401    use std::sync::Arc;
402
403    use api::v1::region::{
404        AlignedSchemaVersion, BulkInsertRequest, DeleteRequest, InsertRequest, RegionRequestHeader,
405    };
406    use api::v1::value::ValueData;
407    use api::v1::{ArrowIpc, PartitionExprVersion};
408    use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, Value};
409    use common_error::ext::{ErrorExt, RetryHint};
410    use common_error::status_code::StatusCode;
411    use common_grpc::flight;
412    use common_meta::cache::{self, TableInfoCache};
413    use common_meta::key::TableMetadataManager;
414    use common_meta::key::table_route::TableRouteValue;
415    use common_meta::kv_backend::memory::MemoryKvBackend;
416    use common_meta::kv_backend::test_util::MockKvBackendBuilder;
417    use common_meta::rpc::router::{Region, RegionRoute};
418    use common_meta::rpc::store::RangeResponse;
419    use datatypes::arrow::array::{Int32Array, StringArray};
420    use datatypes::arrow::datatypes::{DataType, Field, Schema};
421    use datatypes::arrow::record_batch::RecordBatch;
422    use datatypes::data_type::ConcreteDataType;
423    use datatypes::schema;
424    use datatypes::value::{Value as DataValue, ValueRef};
425    use mito_codec::row_converter::sparse;
426    use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec};
427    use moka::future::Cache;
428    use partition::expr;
429    use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
430    use partition::multi_dim::MultiDimPartitionRule;
431    use store_api::metric_engine_consts;
432
433    use super::*;
434    use crate::error::Error;
435    use crate::test_util;
436    use crate::test_util::{test_new_partition_info_cache, test_new_table_route_cache};
437
438    fn build_sparse_partitioner() -> TablePartitioner {
439        TablePartitioner {
440            table_id: 1,
441            rule: Arc::new(
442                MultiDimPartitionRule::try_new(
443                    vec!["zone".to_string()],
444                    vec![1, 2],
445                    vec![
446                        expr::col("zone").lt(DataValue::String("m".into())),
447                        expr::col("zone").gt_eq(DataValue::String("m".into())),
448                    ],
449                    true,
450                )
451                .unwrap(),
452            ),
453            versions: [
454                (1, Some(PartitionExprVersion { value: 2 })),
455                (2, Some(PartitionExprVersion { value: 2 })),
456            ]
457            .into_iter()
458            .collect(),
459        }
460    }
461
462    fn build_sparse_table_info() -> TableInfo {
463        let mut table_info = test_util::new_test_table_info(1, "physical_tags", [1, 2].into_iter());
464        table_info.meta.schema = Arc::new(schema::Schema::new(vec![
465            schema::ColumnSchema::new("host", ConcreteDataType::string_datatype(), true),
466            schema::ColumnSchema::new("zone", ConcreteDataType::string_datatype(), true),
467        ]));
468        table_info.meta.primary_key_indices = vec![0, 1];
469        table_info.meta.partition_key_indices = vec![1];
470        table_info.meta.column_ids = vec![10, 20];
471        table_info.meta.options.extra_options.insert(
472            metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY.to_string(),
473            String::new(),
474        );
475        table_info
476    }
477
478    fn build_sparse_batch() -> RecordBatch {
479        let codec = SparsePrimaryKeyCodec::schemaless();
480        let keys = [Some("east"), Some("west"), Some(""), None, None]
481            .into_iter()
482            .enumerate()
483            .map(|(index, zone)| {
484                let mut key = Vec::new();
485                let values = [
486                    (
487                        sparse::RESERVED_COLUMN_ID_TABLE_ID,
488                        ValueRef::UInt32(42 + index as u32),
489                    ),
490                    (
491                        sparse::RESERVED_COLUMN_ID_TSID,
492                        ValueRef::UInt64(index as u64),
493                    ),
494                    (10, ValueRef::String("node")),
495                ];
496                codec.encode_value_refs(&values, &mut key).unwrap();
497                if let Some(zone) = zone {
498                    codec
499                        .encode_raw_tag_value([(20, zone.as_bytes())].into_iter(), &mut key)
500                        .unwrap();
501                }
502                if index == 4 {
503                    key.extend_from_slice(&20_u32.to_be_bytes());
504                    key.push(0);
505                }
506                key
507            })
508            .collect::<Vec<_>>();
509        RecordBatch::try_new(
510            Arc::new(Schema::new(vec![
511                Field::new("host", DataType::Utf8, false),
512                Field::new(consts::PRIMARY_KEY_COLUMN_NAME, DataType::Binary, false),
513            ])),
514            vec![
515                Arc::new(StringArray::from(vec!["node"; keys.len()])),
516                Arc::new(BinaryArray::from_iter_values(keys.iter())),
517            ],
518        )
519        .unwrap()
520    }
521
522    #[tokio::test]
523    async fn test_partition_tags_use_physical_table_metadata() {
524        let backend = Arc::new(MemoryKvBackend::new());
525        let table_route_cache = test_new_table_route_cache(backend.clone());
526        let partition_info_cache = test_new_partition_info_cache(table_route_cache.clone());
527        let manager =
528            PartitionRuleManager::new(backend.clone(), table_route_cache, partition_info_cache);
529        let table_info = build_sparse_table_info();
530        let routes = [
531            expr::col("zone").lt(DataValue::String("m".into())),
532            expr::col("zone").gt_eq(DataValue::String("m".into())),
533        ]
534        .into_iter()
535        .enumerate()
536        .map(|(index, expr)| RegionRoute {
537            region: Region {
538                id: RegionId::new(1, index as u32 + 1),
539                partition_expr: expr.as_json_str().unwrap(),
540                ..Default::default()
541            },
542            ..Default::default()
543        })
544        .collect();
545        TableMetadataManager::new(backend)
546            .create_table_metadata(
547                table_info.clone(),
548                TableRouteValue::physical(routes),
549                Default::default(),
550            )
551            .await
552            .unwrap();
553        let batch = build_sparse_batch();
554        let partitioner = TablePartitioner::try_new(&manager, &table_info)
555            .await
556            .unwrap();
557        let restored = partitioner
558            .restore_partition_tags(&batch, &table_info)
559            .unwrap()
560            .unwrap();
561        let zones = restored
562            .column_by_name("zone")
563            .unwrap()
564            .as_any()
565            .downcast_ref::<StringArray>()
566            .unwrap();
567        assert_eq!(
568            zones.iter().collect::<Vec<_>>(),
569            vec![Some("east"), Some("west"), Some(""), None, None]
570        );
571
572        for (physical, tag, valid_ids) in [
573            (false, true, true),
574            (true, false, true),
575            (true, true, false),
576        ] {
577            let mut info = table_info.clone();
578            if !physical {
579                info.meta
580                    .options
581                    .extra_options
582                    .remove(metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY);
583            }
584            if !tag {
585                info.meta.primary_key_indices = vec![0];
586            }
587            if !valid_ids {
588                info.meta.column_ids.clear();
589            }
590            let partitioner = TablePartitioner::try_new(&manager, &info).await.unwrap();
591            assert!(
592                partitioner
593                    .restore_partition_tags(&batch, &info)
594                    .unwrap()
595                    .is_none()
596            );
597            assert!(
598                partitioner
599                    .partition_bulk_insert_requests(batch.clone(), None, false)
600                    .is_err()
601            );
602        }
603    }
604
605    #[test]
606    fn test_rebuild_sparse_bulk_with_changed_partition_column() {
607        let partitioner = build_sparse_partitioner();
608        let table_info = build_sparse_table_info();
609        let batch = build_sparse_batch();
610        let restored = partitioner
611            .restore_partition_tags(&batch, &table_info)
612            .unwrap()
613            .unwrap();
614        let zones = restored
615            .column_by_name("zone")
616            .unwrap()
617            .as_any()
618            .downcast_ref::<StringArray>()
619            .unwrap();
620        assert_eq!(
621            zones.iter().collect::<Vec<_>>(),
622            vec![Some("east"), Some("west"), Some(""), None, None]
623        );
624        let (schema, data_header, payload) = flight::record_batch_to_ipc(batch.clone()).unwrap();
625        let rebuilt = partitioner
626            .rebuild_write_requests(
627                build_request(Body::BulkInsert(BulkInsertRequest {
628                    region_id: RegionId::new(1, 1).as_u64(),
629                    partition_expr_version: Some(PartitionExprVersion { value: 1 }),
630                    skip_wal: true,
631                    body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
632                        schema,
633                        data_header,
634                        payload,
635                    })),
636                    ..Default::default()
637                })),
638                &table_info,
639            )
640            .unwrap();
641        assert_eq!(rebuilt.len(), 2);
642        for request in rebuilt {
643            let Some(Body::BulkInsert(bulk)) = request.body else {
644                unreachable!()
645            };
646            assert!(bulk.skip_wal);
647            assert_eq!(
648                bulk.partition_expr_version,
649                Some(PartitionExprVersion { value: 2 })
650            );
651            let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else {
652                unreachable!()
653            };
654            let actual = FlightDecoder::try_from_schema_bytes(&ipc.schema)
655                .unwrap()
656                .try_decode_record_batch(&ipc.data_header, &ipc.payload)
657                .unwrap();
658            let mask = if RegionId::from_u64(bulk.region_id).region_number() == 1 {
659                vec![true, false, true, true, true]
660            } else {
661                vec![false, true, false, false, false]
662            };
663            let expected = compute::filter_record_batch(&batch, &mask.into()).unwrap();
664            assert_eq!(actual, expected);
665        }
666    }
667
668    #[test]
669    fn test_restore_partition_tags_checks_only_missing_tags() {
670        let partitioner = build_sparse_partitioner();
671        let table_info = build_sparse_table_info();
672        let batch = build_sparse_batch();
673        let invalid = RecordBatch::try_new(
674            batch.schema(),
675            vec![
676                batch.column(0).clone(),
677                Arc::new(BinaryArray::from_iter_values([b"invalid"; 5])),
678            ],
679        )
680        .unwrap();
681        assert!(
682            partitioner
683                .restore_partition_tags(&invalid, &table_info)
684                .is_err()
685        );
686        let mut fields = invalid.schema().fields().to_vec();
687        fields.push(Arc::new(Field::new("zone", DataType::Utf8, true)));
688        let mut columns = invalid.columns().to_vec();
689        columns.push(Arc::new(StringArray::from(vec![Some("east"); 5])));
690        let complete = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap();
691        assert!(
692            partitioner
693                .restore_partition_tags(&complete, &table_info)
694                .unwrap()
695                .is_none()
696        );
697        let (schema, data_header, payload) = flight::record_batch_to_ipc(invalid).unwrap();
698        let request = build_request(Body::BulkInsert(BulkInsertRequest {
699            region_id: RegionId::new(1, 1).as_u64(),
700            partition_expr_version: Some(PartitionExprVersion { value: 2 }),
701            body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
702                schema,
703                data_header,
704                payload,
705            })),
706            ..Default::default()
707        }));
708        let expected = request.clone();
709        assert_eq!(
710            partitioner
711                .rebuild_write_requests(request, &table_info)
712                .unwrap(),
713            vec![expected]
714        );
715    }
716
717    #[tokio::test]
718    async fn test_partition_requests_preserve_routing_errors() {
719        let table_info = test_util::new_test_table_info(1, "table_1", [1].into_iter());
720        for (backend_unavailable, expected_status, expected_hint) in [
721            (false, StatusCode::TableNotFound, RetryHint::NonRetryable),
722            (true, StatusCode::Internal, RetryHint::Retryable),
723        ] {
724            let backend = Arc::new(
725                MockKvBackendBuilder::default()
726                    .range_fn(Arc::new(move |_| {
727                        if backend_unavailable {
728                            common_meta::error::NoLeaderSnafu { table_id: 1u32 }.fail()
729                        } else {
730                            Ok(RangeResponse {
731                                kvs: Vec::new(),
732                                more: false,
733                            })
734                        }
735                    }))
736                    .build()
737                    .unwrap(),
738            );
739            let routes = Arc::new(cache::new_table_route_cache(
740                "partitioner_test_routes".to_string(),
741                Cache::new(16),
742                backend.clone(),
743            ));
744            let partitions = Arc::new(partition::cache::new_partition_info_cache(
745                "partitioner_test_partitions".to_string(),
746                Cache::new(16),
747                routes.clone(),
748            ));
749            let manager = PartitionRuleManager::new(backend, routes, partitions);
750            let partitioner = Partitioner::new(&manager);
751            let insert_error = partitioner
752                .partition_insert_requests(&table_info, Rows::default(), false)
753                .await
754                .unwrap_err();
755            let delete_error = partitioner
756                .partition_delete_requests(&table_info, Rows::default())
757                .await
758                .unwrap_err();
759
760            assert!(matches!(insert_error, Error::SplitInsert { .. }));
761            assert!(matches!(delete_error, Error::SplitDelete { .. }));
762            for error in [insert_error, delete_error] {
763                assert_eq!(error.status_code(), expected_status);
764                assert_eq!(error.retry_hint(), expected_hint);
765            }
766        }
767    }
768
769    async fn create_fixture() -> (PartitionRuleManagerRef, TableInfoCache) {
770        let backend = Arc::new(MemoryKvBackend::new());
771        let manager = test_util::create_partition_rule_manager(backend.clone()).await;
772        TableMetadataManager::new(backend.clone())
773            .create_table_metadata(
774                test_util::new_test_table_info(42, "logical_reroute", [1, 2, 3].into_iter()),
775                TableRouteValue::logical(1),
776                Default::default(),
777            )
778            .await
779            .unwrap();
780        let cache = cache::new_table_info_cache(
781            "reroute_test_table_info".to_string(),
782            Cache::new(16),
783            backend,
784        );
785        (manager, cache)
786    }
787
788    fn build_request(body: Body) -> RegionRequest {
789        RegionRequest {
790            header: Some(RegionRequestHeader {
791                dbname: "greptime.public".to_string(),
792                ..Default::default()
793            }),
794            body: Some(body),
795        }
796    }
797
798    fn build_rows(values: &[i32]) -> Rows {
799        Rows {
800            schema: vec![ColumnSchema {
801                column_name: "a".to_string(),
802                datatype: ColumnDataType::Int32 as i32,
803                ..Default::default()
804            }],
805            rows: values
806                .iter()
807                .map(|value| Row {
808                    values: vec![Value {
809                        value_data: Some(ValueData::I32Value(*value)),
810                    }],
811                })
812                .collect(),
813        }
814    }
815
816    fn build_row_request(
817        rows: Rows,
818        insert: bool,
819        region: u32,
820        version: Option<u64>,
821    ) -> RegionRequest {
822        let region_id = RegionId::new(1, region).as_u64();
823        let partition_expr_version = version.map(|value| PartitionExprVersion { value });
824        let body = if insert {
825            Body::Inserts(InsertRequests {
826                requests: vec![InsertRequest {
827                    region_id,
828                    partition_expr_version,
829                    rows: Some(rows),
830                    skip_wal: true,
831                }],
832            })
833        } else {
834            Body::Deletes(DeleteRequests {
835                requests: vec![DeleteRequest {
836                    region_id,
837                    partition_expr_version,
838                    rows: Some(rows),
839                }],
840            })
841        };
842        build_request(body)
843    }
844
845    fn row_request_parts(request: &RegionRequest) -> (u64, Option<u64>, &Rows) {
846        match request.body.as_ref().unwrap() {
847            Body::Inserts(inserts) => {
848                assert_eq!(inserts.requests.len(), 1);
849                let insert = &inserts.requests[0];
850                assert!(insert.skip_wal);
851                (
852                    insert.region_id,
853                    insert.partition_expr_version.map(|v| v.value),
854                    insert.rows.as_ref().unwrap(),
855                )
856            }
857            Body::Deletes(deletes) => {
858                assert_eq!(deletes.requests.len(), 1);
859                let delete = &deletes.requests[0];
860                (
861                    delete.region_id,
862                    delete.partition_expr_version.map(|v| v.value),
863                    delete.rows.as_ref().unwrap(),
864                )
865            }
866            _ => panic!("expected row write"),
867        }
868    }
869
870    #[tokio::test]
871    async fn test_rebuild_rows_updates_partitions_without_losing_rows() {
872        let (manager, cache) = create_fixture().await;
873        let table_info = cache.get(1).await.unwrap().unwrap();
874        let (_, versions) = manager
875            .find_table_partition_rule(&table_info)
876            .await
877            .unwrap();
878        for insert in [true, false] {
879            let request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0));
880            let header = request.header.clone();
881            let rerouted = rebuild_with_cached_table(&manager, &cache, request)
882                .await
883                .unwrap();
884            assert_eq!(rerouted.len(), 3);
885            let mut actual = Vec::new();
886            for request in rerouted {
887                assert_eq!(request.header, header);
888                let (id, version, rows) = row_request_parts(&request);
889                let region_id = RegionId::from_u64(id);
890                assert_eq!(region_id.table_id(), 1);
891                assert_eq!(version, versions[&region_id.region_number()]);
892                assert_eq!(rows.schema, build_rows(&[]).schema);
893                assert_eq!(rows.rows.len(), 1);
894                actual.push((region_id.region_number(), rows.rows[0].clone()));
895            }
896            actual.sort_by_key(|(region, _)| *region);
897            assert_eq!(
898                actual,
899                vec![
900                    (1, build_rows(&[70]).rows.remove(0)),
901                    (2, build_rows(&[20]).rows.remove(0)),
902                    (3, build_rows(&[1]).rows.remove(0))
903                ]
904            );
905        }
906    }
907
908    #[tokio::test]
909    async fn test_unchanged_row_routes_preserve_allocations() {
910        let (manager, cache) = create_fixture().await;
911        let table_info = cache.get(1).await.unwrap().unwrap();
912        let (_, versions) = manager
913            .find_table_partition_rule(&table_info)
914            .await
915            .unwrap();
916        for insert in [true, false] {
917            let request = build_row_request(build_rows(&[1]), insert, 3, versions[&3]);
918            let original_rows = row_request_parts(&request).2;
919            let row_ptr = original_rows.rows.as_ptr();
920            let value_ptr = original_rows.rows[0].values.as_ptr();
921            let rerouted = rebuild_with_cached_table(&manager, &cache, request)
922                .await
923                .unwrap();
924            assert_eq!(rerouted.len(), 1);
925            let rows = row_request_parts(&rerouted[0]).2;
926            assert_eq!(rows.rows.as_ptr(), row_ptr);
927            assert_eq!(rows.rows[0].values.as_ptr(), value_ptr);
928            assert_eq!(rows, &build_rows(&[1]));
929        }
930    }
931
932    fn build_bulk_request(region: u32, version: Option<u64>, values: Vec<i32>) -> RegionRequest {
933        let batch = RecordBatch::try_new(
934            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
935            vec![Arc::new(Int32Array::from(values))],
936        )
937        .unwrap();
938        let (schema, data_header, payload) = flight::record_batch_to_ipc(batch).unwrap();
939        build_request(Body::BulkInsert(BulkInsertRequest {
940            region_id: RegionId::new(1, region).as_u64(),
941            partition_expr_version: version.map(|value| PartitionExprVersion { value }),
942            aligned_schema_version: Some(AlignedSchemaVersion {
943                schema_version: 123,
944            }),
945            skip_wal: true,
946            body: Some(bulk_insert_request::Body::ArrowIpc(ArrowIpc {
947                schema,
948                data_header,
949                payload,
950            })),
951        }))
952    }
953
954    #[tokio::test]
955    async fn test_rebuild_bulk_reencodes_only_changed_partitions() {
956        let (manager, cache) = create_fixture().await;
957        let table_info = cache.get(1).await.unwrap().unwrap();
958        let (_, versions) = manager
959            .find_table_partition_rule(&table_info)
960            .await
961            .unwrap();
962        let request = build_bulk_request(99, Some(0), vec![1, 20, 70]);
963        let header = request.header.clone();
964        let rerouted = rebuild_with_cached_table(&manager, &cache, request)
965            .await
966            .unwrap();
967        assert_eq!(rerouted.len(), 3);
968        let mut actual = Vec::new();
969        for request in rerouted {
970            assert_eq!(request.header, header);
971            let Some(Body::BulkInsert(bulk)) = request.body else {
972                panic!("expected bulk insert")
973            };
974            let region = RegionId::from_u64(bulk.region_id);
975            assert_eq!(region.table_id(), 1);
976            assert!(bulk.skip_wal);
977            assert_eq!(bulk.aligned_schema_version, None);
978            assert_eq!(
979                bulk.partition_expr_version.map(|v| v.value),
980                versions[&region.region_number()]
981            );
982            let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = bulk.body else {
983                panic!("expected Arrow IPC")
984            };
985            let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
986                .unwrap()
987                .try_decode_record_batch(&ipc.data_header, &ipc.payload)
988                .unwrap();
989            assert_eq!(batch.num_rows(), 1);
990            let values = batch
991                .column(0)
992                .as_any()
993                .downcast_ref::<Int32Array>()
994                .unwrap();
995            actual.push((region.region_number(), values.value(0)));
996        }
997        actual.sort_unstable();
998        assert_eq!(actual, [(1, 70), (2, 20), (3, 1)]);
999
1000        // A leader-only change retains the original IPC buffers and schema hint.
1001        let request = build_bulk_request(3, versions[&3], vec![1]);
1002        let Some(Body::BulkInsert(bulk)) = request.body.as_ref() else {
1003            unreachable!()
1004        };
1005        let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else {
1006            unreachable!()
1007        };
1008        let pointers = (
1009            ipc.schema.as_ptr(),
1010            ipc.data_header.as_ptr(),
1011            ipc.payload.as_ptr(),
1012        );
1013        let rerouted = rebuild_with_cached_table(&manager, &cache, request)
1014            .await
1015            .unwrap();
1016        assert_eq!(rerouted.len(), 1);
1017        let Some(Body::BulkInsert(bulk)) = rerouted[0].body.as_ref() else {
1018            unreachable!()
1019        };
1020        assert_eq!(
1021            bulk.aligned_schema_version,
1022            Some(AlignedSchemaVersion {
1023                schema_version: 123
1024            })
1025        );
1026        let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = &bulk.body else {
1027            unreachable!()
1028        };
1029        assert_eq!(
1030            (
1031                ipc.schema.as_ptr(),
1032                ipc.data_header.as_ptr(),
1033                ipc.payload.as_ptr()
1034            ),
1035            pointers
1036        );
1037    }
1038
1039    #[tokio::test]
1040    async fn test_rebuild_logical_table_preserves_logical_region_ids() {
1041        let (manager, cache) = create_fixture().await;
1042        for insert in [true, false] {
1043            let mut request = build_row_request(build_rows(&[1, 20, 70]), insert, 99, Some(0));
1044            match request.body.as_mut().unwrap() {
1045                Body::Inserts(inserts) => {
1046                    inserts.requests[0].region_id = RegionId::new(42, 99).as_u64()
1047                }
1048                Body::Deletes(deletes) => {
1049                    deletes.requests[0].region_id = RegionId::new(42, 99).as_u64()
1050                }
1051                _ => unreachable!(),
1052            }
1053            let rerouted = rebuild_with_cached_table(&manager, &cache, request)
1054                .await
1055                .unwrap();
1056            assert_eq!(rerouted.len(), 3);
1057            let mut regions = rerouted
1058                .iter()
1059                .map(|request| {
1060                    let region = RegionId::from_u64(row_request_parts(request).0);
1061                    assert_eq!(region.table_id(), 42);
1062                    region.region_number()
1063                })
1064                .collect::<Vec<_>>();
1065            regions.sort_unstable();
1066            assert_eq!(regions, [1, 2, 3]);
1067        }
1068    }
1069
1070    async fn rebuild_with_cached_table(
1071        manager: &PartitionRuleManager,
1072        cache: &TableInfoCache,
1073        request: RegionRequest,
1074    ) -> Result<Vec<RegionRequest>> {
1075        let region_id = match request.body.as_ref().unwrap() {
1076            Body::Inserts(inserts) => inserts.requests[0].region_id,
1077            Body::Deletes(deletes) => deletes.requests[0].region_id,
1078            Body::BulkInsert(bulk) => bulk.region_id,
1079            _ => unreachable!(),
1080        };
1081        let table_id = RegionId::from_u64(region_id).table_id();
1082        let table_info = cache.get(table_id).await.unwrap().unwrap();
1083        let partitioner = TablePartitioner::try_new(manager, &table_info)
1084            .await
1085            .context(error::InvalidPartitionSnafu)?;
1086        partitioner.rebuild_write_requests(request, &table_info)
1087    }
1088}
operator/req_convert/common/partitioner.rs

operator/req_convert/common/
partitioner.rs
