diff --git a/src/datatypes/src/timestamp.rs b/src/datatypes/src/timestamp.rs index bf6fdaa162..ac6ca1444f 100644 --- a/src/datatypes/src/timestamp.rs +++ b/src/datatypes/src/timestamp.rs @@ -13,7 +13,7 @@ // limitations under the License. use arrow_array::{ - ArrayRef, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, + Array, ArrayRef, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, }; use arrow_schema::DataType; @@ -179,8 +179,24 @@ pub fn timestamp_array_to_primitive( Some((ts_primitive, *unit)) } +/// Appends non-null timestamps in the source array's native time unit. +/// +/// Returns `None` for a non-timestamp array without changing `timestamps`. +pub fn append_timestamps(ts_array: &ArrayRef, timestamps: &mut Vec) -> Option<()> { + let (values, _) = timestamp_array_to_primitive(ts_array)?; + if values.null_count() == 0 { + timestamps.extend_from_slice(values.values()); + } else { + timestamps.extend(values.iter().flatten()); + } + Some(()) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + + use arrow_array::Int64Array; use common_time::timezone::set_default_timezone; use super::*; @@ -213,4 +229,46 @@ mod tests { assert_eq!(ts, ts.as_scalar_ref()); assert_eq!(ts, ts.to_owned_scalar()); } + + #[test] + fn test_append_timestamps() { + let cases = [ + vec![Some(i64::MIN), Some(-1), Some(0), Some(i64::MAX)], + vec![Some(-1), None, Some(2), None], + vec![None, None], + vec![], + ]; + for values in cases { + let arrays: [ArrayRef; 4] = [ + Arc::new(TimestampSecondArray::from(values.clone())), + Arc::new(TimestampMillisecondArray::from(values.clone())), + Arc::new(TimestampMicrosecondArray::from(values.clone())), + Arc::new(TimestampNanosecondArray::from(values.clone())), + ]; + let mut expected = vec![42]; + expected.extend(values.iter().flatten().copied()); + for array in arrays { + for _ in 0..2 { + let (primitive, _) = timestamp_array_to_primitive(&array).unwrap(); + let mut reference = vec![42]; + reference.extend(primitive.iter().flatten()); + + let mut timestamps = vec![42]; + assert_eq!(append_timestamps(&array, &mut timestamps), Some(())); + assert_eq!(timestamps, expected); + assert_eq!(timestamps, reference); + } + } + } + } + + #[test] + fn test_append_timestamps_invalid_array_preserves_prefix() { + for values in [vec![], vec![Some(1), None]] { + let array: ArrayRef = Arc::new(Int64Array::from(values)); + let mut timestamps = vec![42, -1]; + assert_eq!(append_timestamps(&array, &mut timestamps), None); + assert_eq!(timestamps, vec![42, -1]); + } + } } diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 22fae23619..0a3cf23eb6 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -20,7 +20,6 @@ use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, PartitionExprVersion}; -use arrow::array::Array; use arrow::record_batch::RecordBatch; use bytes::Bytes; use common_base::AffectedRows; @@ -28,12 +27,13 @@ use common_grpc::FlightData; use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; use crate::insert::Inserter; +use crate::req_convert::insert::extract_timestamps; use crate::{error, metrics}; impl Inserter { @@ -316,22 +316,3 @@ impl Inserter { }); } } - -/// Calculate the timestamp range of record batch. Return `None` if record batch is empty. -fn extract_timestamps(rb: &RecordBatch, timestamp_index_name: &str) -> error::Result> { - let ts_col = rb - .column_by_name(timestamp_index_name) - .context(error::ColumnNotFoundSnafu { - msg: timestamp_index_name, - })?; - if rb.num_rows() == 0 { - return Ok(vec![]); - } - let (primitive, _) = - datatypes::timestamp::timestamp_array_to_primitive(ts_col).with_context(|| { - error::InvalidTimeIndexTypeSnafu { - ty: ts_col.data_type().clone(), - } - })?; - Ok(primitive.iter().flatten().collect()) -} diff --git a/src/operator/src/req_convert/insert.rs b/src/operator/src/req_convert/insert.rs index f9bf83d926..72d6693c6b 100644 --- a/src/operator/src/req_convert/insert.rs +++ b/src/operator/src/req_convert/insert.rs @@ -17,6 +17,7 @@ mod fill_impure_default; mod row_to_region; mod stmt_to_region; mod table_to_region; +mod timestamps; use api::v1::SemanticType; pub use column_to_row::ColumnToRow; @@ -26,6 +27,7 @@ use snafu::{OptionExt, ResultExt}; pub use stmt_to_region::StatementToRegion; use table::metadata::TableInfo; pub use table_to_region::TableToRegion; +pub use timestamps::extract_timestamps; use crate::error::{ColumnNotFoundSnafu, MissingTimeIndexColumnSnafu, Result}; diff --git a/src/operator/src/req_convert/insert/timestamps.rs b/src/operator/src/req_convert/insert/timestamps.rs new file mode 100644 index 0000000000..f0315135f5 --- /dev/null +++ b/src/operator/src/req_convert/insert/timestamps.rs @@ -0,0 +1,38 @@ +// 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 arrow::record_batch::RecordBatch; +use datatypes::timestamp::append_timestamps; +use snafu::OptionExt; + +use crate::error::{self, Result}; + +/// Extracts non-null timestamps in the source column's native time unit. +pub fn extract_timestamps(rb: &RecordBatch, timestamp_index_name: &str) -> Result> { + let ts_col = rb + .column_by_name(timestamp_index_name) + .context(error::ColumnNotFoundSnafu { + msg: timestamp_index_name, + })?; + if rb.num_rows() == 0 { + return Ok(vec![]); + } + let mut timestamps = Vec::with_capacity(rb.num_rows()); + append_timestamps(ts_col, &mut timestamps).with_context(|| { + error::InvalidTimeIndexTypeSnafu { + ty: ts_col.data_type().clone(), + } + })?; + Ok(timestamps) +} diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 3c9a5d2a99..6f6991a542 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -23,7 +23,6 @@ use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, ColumnSchema, RowInsertRequests, Rows}; -use arrow::array::Array; use arrow::compute::{concat_batches, filter_record_batch}; use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; use arrow::record_batch::RecordBatch; @@ -44,6 +43,7 @@ use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_timestamp, greptim use common_runtime::spawn_global; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; +use datatypes::timestamp::append_timestamps; use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; use partition::manager::PartitionRuleManagerRef; use partition::partition::PartitionRuleRef; @@ -1389,21 +1389,13 @@ fn extract_timestamps(table_batch: &TableBatch) -> Vec { let mut timestamps = Vec::with_capacity(table_batch.row_count); for batch in &table_batch.batches { let timestamp_column = batch.batch.column(batch.timestamp_index); - let Some((timestamp_values, _)) = - datatypes::timestamp::timestamp_array_to_primitive(timestamp_column) - else { + let Some(()) = append_timestamps(timestamp_column, &mut timestamps) else { error!( "Failed to extract timestamps from record batch, table_id: {}, timestamp_index: {}", table_batch.table_id, batch.timestamp_index ); continue; }; - - if timestamp_values.null_count() == 0 { - timestamps.extend_from_slice(timestamp_values.values()); - } else { - timestamps.extend(timestamp_values.iter().flatten()); - } } timestamps }