use BTreeMap

Author: baojinri

proxy/src/opentsdb/types.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     default::Default,
     fmt::Debug,
 };
@@ -220,13 +220,13 @@ pub struct QueryResponse {
     /// A list of tags only returned when the results are for a single time
     /// series. If results are aggregated, this value may be null or an empty
     /// map
-    pub(crate) tags: HashMap<String, String>,
+    pub(crate) tags: BTreeMap<String, String>,
     /// If more than one timeseries were included in the result set, i.e. they
     /// were aggregated, this will display a list of tag names that were found
     /// in common across all time series.
     #[serde(rename = "aggregatedTags")]
     pub(crate) aggregated_tags: Vec<String>,
-    pub(crate) dps: HashMap<String, f64>,
+    pub(crate) dps: BTreeMap<String, f64>,
 }
 
 #[derive(Default)]
@@ -237,10 +237,10 @@ struct QueryConverter {
     // (column_name, index)
     tags_idx: Vec<(String, usize)>,
     aggregated_tags: Vec<String>,
-    // (tags_key, (tagk, tagv))
-    tags: HashMap<String, HashMap<String, String>>,
-    // (tags_key, (timestamp, value))
-    values: HashMap<String, HashMap<String, f64>>,
+    // (time_series, (tagk, tagv))
+    tags: BTreeMap<String, BTreeMap<String, String>>,
+    // (time_series, (timestamp, value))
+    values: BTreeMap<String, BTreeMap<String, f64>>,
 
     resp: Vec<QueryResponse>,
 }
@@ -249,18 +249,18 @@ impl QueryConverter {
     fn try_new(
         schema: &RecordSchema,
         metric: String,
-        timestamp_col_name: &str,
-        field_col_name: &str,
+        timestamp_col_name: String,
+        field_col_name: String,
         tags: Vec<String>,
         aggregated_tags: Vec<String>,
     ) -> Result<Self> {
         let timestamp_idx = schema
-            .index_of(timestamp_col_name)
+            .index_of(&timestamp_col_name)
             .context(InternalNoCause {
                 msg: "Timestamp column is missing in query response",
             })?;
 
-        let value_idx = schema.index_of(field_col_name).context(InternalNoCause {
+        let value_idx = schema.index_of(&field_col_name).context(InternalNoCause {
             msg: "Value column is missing in query response",
         })?;
 
@@ -316,7 +316,8 @@ impl QueryConverter {
     fn add_batch(&mut self, record_batch: RecordBatch) -> Result<()> {
         let row_num = record_batch.num_rows();
         for row_idx in 0..row_num {
-            let mut tags = HashMap::with_capacity(self.tags_idx.len());
+            let mut tags = BTreeMap::new();
+            // tags_key is used to identify a time series
             let mut tags_key = String::new();
             for (tag_key, idx) in &self.tags_idx {
                 let tag_value = record_batch
@@ -404,8 +405,8 @@ pub(crate) fn convert_output_to_response(
             QueryConverter::try_new(
                 record_schema,
                 metric,
-                &timestamp_col_name,
-                &field_col_name,
+                timestamp_col_name,
+                field_col_name,
                 tags,
                 aggregated_tags,
             )?
@@ -467,19 +468,30 @@ mod tests {
                     .unwrap(),
             )
             .unwrap()
+            .add_normal_column(
+                column_schema::Builder::new("tag2".to_string(), DatumKind::String)
+                    .is_tag(true)
+                    .build()
+                    .unwrap(),
+            )
+            .unwrap()
             .build()
             .unwrap()
     }
 
     fn build_record_batch(schema: &Schema) -> RecordBatchVec {
-        let tsid: ArrayRef = Arc::new(UInt64Array::from(vec![1]));
-        let timestamp: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![11111111]));
-        let values: ArrayRef = Arc::new(Float64Array::from(vec![100.0]));
-        let tag1: ArrayRef = Arc::new(StringArray::from(vec!["a"]));
+        let tsid: ArrayRef = Arc::new(UInt64Array::from(vec![1, 1, 2, 3, 3]));
+        let timestamp: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![
+            11111111, 11111112, 11111113, 11111111, 11111112,
+        ]));
+        let values: ArrayRef =
+            Arc::new(Float64Array::from(vec![100.0, 101.0, 200.0, 300.0, 301.0]));
+        let tag1: ArrayRef = Arc::new(StringArray::from(vec!["a", "a", "b", "c", "c"]));
+        let tag2: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "y", "z", "z"]));
 
         let batch = ArrowRecordBatch::try_new(
             schema.to_arrow_schema_ref(),
-            vec![tsid, timestamp, values, tag1],
+            vec![tsid, timestamp, values, tag1, tag2],
         )
         .unwrap();
 
@@ -489,7 +501,7 @@ mod tests {
     #[test]
     fn test_convert_output_to_response() {
         let metric = "metric".to_string();
-        let tags = vec!["tag1".to_string()];
+        let tags = vec!["tag1".to_string(), "tag2".to_string()];
         let schema = build_schema();
         let record_batch = build_record_batch(&schema);
         let result = convert_output_to_response(
@@ -498,19 +510,56 @@ mod tests {
             DEFAULT_FIELD.to_string(),
             TIMESTAMP_COLUMN.to_string(),
             tags,
-            vec!["tag1".to_string()],
+            vec![],
         )
         .unwrap();
 
         assert_eq!(
-            vec![QueryResponse {
-                metric: metric.clone(),
-                tags: vec![("tag1".to_string(), "a".to_string()),]
+            vec![
+                QueryResponse {
+                    metric: metric.clone(),
+                    tags: vec![
+                        ("tag1".to_string(), "a".to_string()),
+                        ("tag2".to_string(), "x".to_string()),
+                    ]
+                    .into_iter()
+                    .collect(),
+                    aggregated_tags: vec![],
+                    dps: vec![
+                        ("11111111".to_string(), 100.0),
+                        ("11111112".to_string(), 101.0),
+                    ]
+                    .into_iter()
+                    .collect(),
+                },
+                QueryResponse {
+                    metric: metric.clone(),
+                    tags: vec![
+                        ("tag1".to_string(), "b".to_string()),
+                        ("tag2".to_string(), "y".to_string()),
+                    ]
+                    .into_iter()
+                    .collect(),
+                    aggregated_tags: vec![],
+                    dps: vec![("11111113".to_string(), 200.0),].into_iter().collect(),
+                },
+                QueryResponse {
+                    metric: metric.clone(),
+                    tags: vec![
+                        ("tag1".to_string(), "c".to_string()),
+                        ("tag2".to_string(), "z".to_string()),
+                    ]
+                    .into_iter()
+                    .collect(),
+                    aggregated_tags: vec![],
+                    dps: vec![
+                        ("11111111".to_string(), 300.0),
+                        ("11111112".to_string(), 301.0),
+                    ]
                     .into_iter()
                     .collect(),
-                aggregated_tags: vec!["tag1".to_string()],
-                dps: vec![("11111111".to_string(), 100.0),].into_iter().collect(),
-            }],
+                },
+            ],
             result
         );
     }