refactor: share timestamp extraction for batched writes

Signed-off-by: WenyXu <wenymedia@gmail.com>

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<i64>) -> 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]);
+        }
+    }
 }

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<Vec<i64>> {
-    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())
-}

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};
 

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<Vec<i64>> {
+    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)
+}

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<i64> {
     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
 }