1pub(crate) mod partitioner;
16
17use std::collections::HashMap;
18
19use api::helper::ColumnDataTypeWrapper;
20use api::v1::column_data_type_extension::TypeExt;
21use api::v1::column_def::options_from_column_schema;
22use api::v1::value::ValueData;
23use api::v1::{
24 Column, ColumnDataType, ColumnDataTypeExtension, ColumnSchema, JsonTypeExtension, Row,
25 RowDeleteRequest, RowInsertRequest, Rows, SemanticType, Value,
26};
27use common_base::BitVec;
28use datatypes::prelude::ConcreteDataType;
29use datatypes::vectors::VectorRef;
30use snafu::ResultExt;
31use snafu::prelude::*;
32use table::metadata::TableInfo;
33
34use crate::error::{
35 ColumnDataTypeSnafu, ColumnNotFoundSnafu, InvalidInsertRequestSnafu, InvalidJsonFormatSnafu,
36 MissingTimeIndexColumnSnafu, Result, UnexpectedSnafu,
37};
38
39fn encode_string_to_jsonb_binary(value_data: ValueData) -> Result<ValueData> {
41 if let ValueData::StringValue(json) = &value_data {
42 let binary = jsonb::parse_value_standard_mode(json.as_bytes())
43 .map_err(|_| InvalidJsonFormatSnafu { json }.build())
44 .map(|jsonb| jsonb.to_vec())?;
45 Ok(ValueData::BinaryValue(binary))
46 } else {
47 UnexpectedSnafu {
48 violated: "Expected to value data to be a string.",
49 }
50 .fail()
51 }
52}
53
54pub fn preprocess_row_insert_requests(requests: &mut Vec<RowInsertRequest>) -> Result<()> {
56 for request in requests {
57 validate_rows(&request.rows)?;
58 prepare_rows(&mut request.rows)?;
59 }
60
61 Ok(())
62}
63
64pub fn preprocess_row_delete_requests(requests: &mut Vec<RowDeleteRequest>) -> Result<()> {
66 for request in requests {
67 validate_rows(&request.rows)?;
68 prepare_rows(&mut request.rows)?;
69 }
70
71 Ok(())
72}
73
74fn prepare_rows(rows: &mut Option<Rows>) -> Result<()> {
75 if let Some(rows) = rows {
76 let indexes = rows
77 .schema
78 .iter()
79 .enumerate()
80 .filter_map(|(idx, schema)| {
81 if schema.datatype() == ColumnDataType::Json {
82 match schema
83 .datatype_extension
84 .as_ref()
85 .and_then(|x| x.type_ext.as_ref())
86 {
87 None | Some(TypeExt::JsonType(_)) => Some(Ok(idx)),
88 Some(TypeExt::JsonNativeType(_)) => None,
89 Some(_) => Some(
90 InvalidInsertRequestSnafu {
91 reason: format!(
92 "Invalid type extension for JSON column '{}'",
93 schema.column_name
94 ),
95 }
96 .fail(),
97 ),
98 }
99 } else {
100 None
101 }
102 })
103 .collect::<Result<Vec<_>>>()?;
104 for idx in &indexes {
105 let column = &mut rows.schema[*idx];
106 column.datatype_extension = Some(ColumnDataTypeExtension {
107 type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())),
108 });
109 column.datatype = ColumnDataType::Json.into();
110 }
111
112 for idx in &indexes {
113 for row in &mut rows.rows {
114 if let Some(value_data) = row.values[*idx].value_data.take() {
115 row.values[*idx].value_data = Some(encode_string_to_jsonb_binary(value_data)?);
116 }
117 }
118 }
119 }
120
121 Ok(())
122}
123
124fn validate_rows(rows: &Option<Rows>) -> Result<()> {
125 let Some(rows) = rows else {
126 return Ok(());
127 };
128
129 for (col_idx, schema) in rows.schema.iter().enumerate() {
130 let column_type =
131 ColumnDataTypeWrapper::try_new(schema.datatype, schema.datatype_extension.clone())
132 .context(ColumnDataTypeSnafu)?
133 .into();
134
135 let ConcreteDataType::Vector(d) = column_type else {
136 return Ok(());
137 };
138
139 for row in &rows.rows {
140 let value = &row.values[col_idx].value_data;
141 if let Some(data) = value {
142 validate_vector_col(data, d.dim)?;
143 }
144 }
145 }
146
147 Ok(())
148}
149
150fn validate_vector_col(data: &ValueData, dim: u32) -> Result<()> {
151 let data = match data {
152 ValueData::BinaryValue(data) => data,
153 _ => {
154 return InvalidInsertRequestSnafu {
155 reason: "Expecting binary data for vector column.".to_string(),
156 }
157 .fail();
158 }
159 };
160
161 let expected_len = dim as usize * std::mem::size_of::<f32>();
162 if data.len() != expected_len {
163 return InvalidInsertRequestSnafu {
164 reason: format!(
165 "Expecting {} bytes of data for vector column, but got {}.",
166 expected_len,
167 data.len()
168 ),
169 }
170 .fail();
171 }
172
173 Ok(())
174}
175
176pub fn columns_to_rows(columns: Vec<Column>, row_count: u32) -> Result<Rows> {
177 let row_count = row_count as usize;
178 let column_count = columns.len();
179 let mut schema = Vec::with_capacity(column_count);
180 let mut rows = vec![
181 Row {
182 values: Vec::with_capacity(column_count)
183 };
184 row_count
185 ];
186 for column in columns {
187 let column_schema = ColumnSchema {
188 column_name: column.column_name.clone(),
189 datatype: column.datatype,
190 semantic_type: column.semantic_type,
191 datatype_extension: column.datatype_extension.clone(),
192 options: column.options.clone(),
193 };
194 schema.push(column_schema);
195
196 push_column_to_rows(column, &mut rows)?;
197 }
198
199 Ok(Rows { schema, rows })
200}
201
202fn push_column_to_rows(column: Column, rows: &mut [Row]) -> Result<()> {
203 let null_mask = BitVec::from_vec(column.null_mask);
204 let column_type = ColumnDataTypeWrapper::try_new(column.datatype, column.datatype_extension)
205 .context(ColumnDataTypeSnafu)?
206 .datatype();
207 let column_values = column.values.unwrap_or_default();
208
209 macro_rules! push_column_values_match_types {
210 ($( ($arm:tt, $value_data_variant:tt, $field_name:tt), )*) => { match column_type { $(
211
212 ColumnDataType::$arm => {
213 let row_count = rows.len();
214 let actual_row_count = null_mask.count_ones() + column_values.$field_name.len();
215 ensure!(
216 actual_row_count == row_count,
217 InvalidInsertRequestSnafu {
218 reason: format!(
219 "Expecting {} rows of data for column '{}', but got {}.",
220 row_count, column.column_name, actual_row_count
221 ),
222 }
223 );
224
225 let mut null_mask_iter = null_mask.into_iter();
226 let mut values_iter = column_values.$field_name.into_iter();
227
228 for row in rows {
229 let value_is_null = null_mask_iter.next();
230 if value_is_null == Some(true) {
231 row.values.push(Value { value_data: None });
232 } else {
233 let value = values_iter.next().unwrap();
235 row.values.push(Value {
236 value_data: Some(ValueData::$value_data_variant(value)),
237 });
238 }
239 }
240 }
241
242 )* _ => {
243 return InvalidInsertRequestSnafu {
244 reason: format!(
245 "Column '{}' with type {:?} is not supported in row inserts.",
246 column.column_name, column_type
247 ),
248 }
249 .fail();
250 } }}
251 }
252
253 push_column_values_match_types!(
254 (Boolean, BoolValue, bool_values),
255 (Int8, I8Value, i8_values),
256 (Int16, I16Value, i16_values),
257 (Int32, I32Value, i32_values),
258 (Int64, I64Value, i64_values),
259 (Uint8, U8Value, u8_values),
260 (Uint16, U16Value, u16_values),
261 (Uint32, U32Value, u32_values),
262 (Uint64, U64Value, u64_values),
263 (Float32, F32Value, f32_values),
264 (Float64, F64Value, f64_values),
265 (Binary, BinaryValue, binary_values),
266 (String, StringValue, string_values),
267 (Json, StringValue, string_values),
268 (Date, DateValue, date_values),
269 (Datetime, DatetimeValue, datetime_values),
270 (
271 TimestampSecond,
272 TimestampSecondValue,
273 timestamp_second_values
274 ),
275 (
276 TimestampMillisecond,
277 TimestampMillisecondValue,
278 timestamp_millisecond_values
279 ),
280 (
281 TimestampMicrosecond,
282 TimestampMicrosecondValue,
283 timestamp_microsecond_values
284 ),
285 (
286 TimestampNanosecond,
287 TimestampNanosecondValue,
288 timestamp_nanosecond_values
289 ),
290 (TimeSecond, TimeSecondValue, time_second_values),
291 (
292 TimeMillisecond,
293 TimeMillisecondValue,
294 time_millisecond_values
295 ),
296 (
297 TimeMicrosecond,
298 TimeMicrosecondValue,
299 time_microsecond_values
300 ),
301 (TimeNanosecond, TimeNanosecondValue, time_nanosecond_values),
302 (
303 IntervalYearMonth,
304 IntervalYearMonthValue,
305 interval_year_month_values
306 ),
307 (
308 IntervalDayTime,
309 IntervalDayTimeValue,
310 interval_day_time_values
311 ),
312 (
313 IntervalMonthDayNano,
314 IntervalMonthDayNanoValue,
315 interval_month_day_nano_values
316 ),
317 (Decimal128, Decimal128Value, decimal128_values),
318 (Vector, BinaryValue, binary_values),
319 (List, ListValue, list_values),
320 (Struct, StructValue, struct_values),
321 );
322
323 Ok(())
324}
325
326pub fn row_count(columns: &HashMap<String, VectorRef>) -> Result<usize> {
327 let mut columns_iter = columns.values();
328
329 let len = columns_iter
330 .next()
331 .map(|column| column.len())
332 .unwrap_or_default();
333 ensure!(
334 columns_iter.all(|column| column.len() == len),
335 InvalidInsertRequestSnafu {
336 reason: "The row count of columns is not the same."
337 }
338 );
339
340 Ok(len)
341}
342
343pub fn column_schema(
344 table_info: &TableInfo,
345 columns: &HashMap<String, VectorRef>,
346) -> Result<Vec<ColumnSchema>> {
347 columns
348 .keys()
349 .map(|column_name| {
350 let column_schema = table_info
351 .meta
352 .schema
353 .column_schema_by_name(column_name)
354 .context(ColumnNotFoundSnafu {
355 msg: format!("unable to find column {column_name} in table schema"),
356 })?;
357
358 let (datatype, datatype_extension) =
359 ColumnDataTypeWrapper::try_from(column_schema.data_type.clone())
360 .context(ColumnDataTypeSnafu)?
361 .to_parts();
362
363 Ok(ColumnSchema {
364 column_name: column_name.clone(),
365 datatype: datatype as i32,
366 semantic_type: semantic_type(table_info, column_name)?.into(),
367 datatype_extension,
368 options: options_from_column_schema(column_schema),
369 })
370 })
371 .collect::<Result<Vec<_>>>()
372}
373
374fn semantic_type(table_info: &TableInfo, column: &str) -> Result<SemanticType> {
375 let table_meta = &table_info.meta;
376 let table_schema = &table_meta.schema;
377
378 let time_index_column = &table_schema
379 .timestamp_column()
380 .with_context(|| table::error::MissingTimeIndexColumnSnafu {
381 table_name: table_info.name.clone(),
382 })
383 .context(MissingTimeIndexColumnSnafu)?
384 .name;
385
386 let semantic_type = if column == time_index_column {
387 SemanticType::Timestamp
388 } else {
389 let column_index = table_schema.column_index_by_name(column);
390 let column_index = column_index.context(ColumnNotFoundSnafu {
391 msg: format!("unable to find column {column} in table schema"),
392 })?;
393
394 if table_meta.primary_key_indices.contains(&column_index) {
395 SemanticType::Tag
396 } else {
397 SemanticType::Field
398 }
399 };
400
401 Ok(semantic_type)
402}
403
404#[cfg(test)]
405mod tests {
406 use api::v1::column::Values;
407 use api::v1::{SemanticType, VectorTypeExtension};
408 use common_base::bit_vec::prelude::*;
409 use common_error::ext::WhateverResult;
410
411 use super::*;
412
413 #[test]
414 fn test_prepare_jsonb_rows() -> WhateverResult<()> {
415 let binary = Some(ColumnDataTypeExtension {
416 type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())),
417 });
418 for extension in [
419 None,
420 Some(ColumnDataTypeExtension::default()),
421 binary.clone(),
422 ] {
423 for value in [Some(ValueData::StringValue(r#"{"key":1}"#.into())), None] {
424 let expected = value.as_ref().map(|_| {
425 ValueData::BinaryValue(jsonb::parse_value(br#"{"key":1}"#).unwrap().to_vec())
426 });
427 let mut rows = Some(Rows {
428 schema: vec![ColumnSchema {
429 datatype: ColumnDataType::Json as i32,
430 datatype_extension: extension.clone(),
431 ..Default::default()
432 }],
433 rows: vec![Row {
434 values: vec![Value { value_data: value }],
435 }],
436 });
437 prepare_rows(&mut rows)?;
438 let rows = rows.unwrap();
439 assert_eq!(rows.schema[0].datatype_extension, binary);
440 assert_eq!(rows.rows[0].values[0].value_data, expected);
441 }
442 }
443 Ok(())
444 }
445
446 #[test]
447 fn test_preprocess_json_invalid_extension() {
448 let mut rows = Some(Rows {
449 schema: vec![ColumnSchema {
450 datatype: ColumnDataType::Json as i32,
451 datatype_extension: Some(ColumnDataTypeExtension {
452 type_ext: Some(TypeExt::VectorType(VectorTypeExtension { dim: 3 })),
453 }),
454 ..Default::default()
455 }],
456 rows: vec![Row {
457 values: vec![Value { value_data: None }],
458 }],
459 });
460 assert!(
461 prepare_rows(&mut rows)
462 .unwrap_err()
463 .to_string()
464 .contains("Invalid type extension")
465 );
466 }
467
468 #[test]
469 fn test_request_column_to_row() {
470 let columns = vec![
471 Column {
472 column_name: String::from("col1"),
473 datatype: ColumnDataType::Int32.into(),
474 semantic_type: SemanticType::Field.into(),
475 null_mask: bitvec![u8, Lsb0; 1, 0, 1].into_vec(),
476 values: Some(Values {
477 i32_values: vec![42],
478 ..Default::default()
479 }),
480 ..Default::default()
481 },
482 Column {
483 column_name: String::from("col2"),
484 datatype: ColumnDataType::String.into(),
485 semantic_type: SemanticType::Tag.into(),
486 null_mask: vec![],
487 values: Some(Values {
488 string_values: vec![
489 String::from("value1"),
490 String::from("value2"),
491 String::from("value3"),
492 ],
493 ..Default::default()
494 }),
495 ..Default::default()
496 },
497 Column {
498 column_name: String::from("col3"),
499 datatype: ColumnDataType::Vector.into(),
500 semantic_type: SemanticType::Field.into(),
501 null_mask: vec![],
502 values: Some(Values {
503 binary_values: vec![vec![0; 4], vec![1; 4], vec![2; 4]],
504 ..Default::default()
505 }),
506 datatype_extension: Some(ColumnDataTypeExtension {
507 type_ext: Some(TypeExt::VectorType(VectorTypeExtension { dim: 1 })),
508 }),
509 ..Default::default()
510 },
511 ];
512 let row_count = 3;
513
514 let result = columns_to_rows(columns, row_count);
515 let rows = result.unwrap();
516
517 assert_eq!(rows.schema.len(), 3);
518 assert_eq!(rows.schema[0].column_name, "col1");
519 assert_eq!(rows.schema[0].datatype, ColumnDataType::Int32 as i32);
520 assert_eq!(rows.schema[0].semantic_type, SemanticType::Field as i32);
521 assert_eq!(rows.schema[1].column_name, "col2");
522 assert_eq!(rows.schema[1].datatype, ColumnDataType::String as i32);
523 assert_eq!(rows.schema[1].semantic_type, SemanticType::Tag as i32);
524 assert_eq!(rows.schema[2].column_name, "col3");
525 assert_eq!(rows.schema[2].datatype, ColumnDataType::Vector as i32);
526 assert_eq!(rows.schema[2].semantic_type, SemanticType::Field as i32);
527 assert_eq!(
528 rows.schema[2].datatype_extension,
529 Some(ColumnDataTypeExtension {
530 type_ext: Some(TypeExt::VectorType(VectorTypeExtension { dim: 1 }))
531 })
532 );
533
534 assert_eq!(rows.rows.len(), 3);
535
536 assert_eq!(rows.rows[0].values.len(), 3);
537 assert_eq!(rows.rows[0].values[0].value_data, None);
538 assert_eq!(
539 rows.rows[0].values[1].value_data,
540 Some(ValueData::StringValue(String::from("value1")))
541 );
542 assert_eq!(
543 rows.rows[0].values[2].value_data,
544 Some(ValueData::BinaryValue(vec![0; 4]))
545 );
546
547 assert_eq!(rows.rows[1].values.len(), 3);
548 assert_eq!(
549 rows.rows[1].values[0].value_data,
550 Some(ValueData::I32Value(42))
551 );
552 assert_eq!(
553 rows.rows[1].values[1].value_data,
554 Some(ValueData::StringValue(String::from("value2")))
555 );
556 assert_eq!(
557 rows.rows[1].values[2].value_data,
558 Some(ValueData::BinaryValue(vec![1; 4]))
559 );
560
561 assert_eq!(rows.rows[2].values.len(), 3);
562 assert_eq!(rows.rows[2].values[0].value_data, None);
563 assert_eq!(
564 rows.rows[2].values[1].value_data,
565 Some(ValueData::StringValue(String::from("value3")))
566 );
567 assert_eq!(
568 rows.rows[2].values[2].value_data,
569 Some(ValueData::BinaryValue(vec![2; 4]))
570 );
571
572 let columns = vec![Column {
574 column_name: String::from("col1"),
575 datatype: ColumnDataType::Int32.into(),
576 semantic_type: SemanticType::Field.into(),
577 null_mask: bitvec![u8, Lsb0; 1, 0, 1].into_vec(),
578 values: Some(Values {
579 i8_values: vec![42],
580 ..Default::default()
581 }),
582 ..Default::default()
583 }];
584 let row_count = 3;
585 assert!(columns_to_rows(columns, row_count).is_err());
586
587 let columns = vec![Column {
589 column_name: String::from("col1"),
590 datatype: ColumnDataType::Int32.into(),
591 semantic_type: SemanticType::Field.into(),
592 null_mask: bitvec![u8, Lsb0; 0, 0, 1].into_vec(),
593 values: Some(Values {
594 i32_values: vec![42],
595 ..Default::default()
596 }),
597 ..Default::default()
598 }];
599 let row_count = 3;
600 assert!(columns_to_rows(columns, row_count).is_err());
601
602 let columns = vec![Column {
604 column_name: String::from("col1"),
605 datatype: ColumnDataType::Int32.into(),
606 semantic_type: SemanticType::Field.into(),
607 null_mask: vec![],
608 values: Some(Values {
609 i32_values: vec![42],
610 ..Default::default()
611 }),
612 ..Default::default()
613 }];
614 let row_count = 3;
615 assert!(columns_to_rows(columns, row_count).is_err());
616 }
617
618 #[test]
619 fn test_validate_vector_row_success() {
620 let data = ValueData::BinaryValue(vec![0; 4]);
621 let dim = 1;
622 assert!(validate_vector_col(&data, dim).is_ok());
623
624 let data = ValueData::BinaryValue(vec![0; 8]);
625 let dim = 2;
626 assert!(validate_vector_col(&data, dim).is_ok());
627
628 let data = ValueData::BinaryValue(vec![0; 12]);
629 let dim = 3;
630 assert!(validate_vector_col(&data, dim).is_ok());
631 }
632
633 #[test]
634 fn test_validate_vector_row_fail_wrong_type() {
635 let data = ValueData::I32Value(42);
636 let dim = 1;
637 assert!(validate_vector_col(&data, dim).is_err());
638 }
639
640 #[test]
641 fn test_validate_vector_row_fail_wrong_length() {
642 let data = ValueData::BinaryValue(vec![0; 8]);
643 let dim = 1;
644 assert!(validate_vector_col(&data, dim).is_err());
645
646 let data = ValueData::BinaryValue(vec![0; 4]);
647 let dim = 2;
648 assert!(validate_vector_col(&data, dim).is_err());
649 }
650}