common_function/aggrs/geo/
geo_path.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::sync::Arc;
16
17use datafusion::arrow::array::{Array, ArrayRef};
18use datafusion::common::cast::as_primitive_array;
19use datafusion::error::{DataFusionError, Result as DfResult};
20use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility};
21use datafusion::prelude::create_udaf;
22use datafusion_common::ScalarValue;
23use datafusion_common::cast::{as_list_array, as_struct_array};
24use datafusion_common::utils::SingleRowListArrayBuilder;
25use datatypes::arrow::array::{Float64Array, Int64Array, ListArray, StructArray};
26use datatypes::arrow::datatypes::{
27    DataType, Field, Float64Type, Int64Type, TimeUnit, TimestampNanosecondType,
28};
29use datatypes::compute::{self, sort_to_indices};
30
31pub const GEO_PATH_NAME: &str = "geo_path";
32
33const LATITUDE_FIELD: &str = "lat";
34const LONGITUDE_FIELD: &str = "lng";
35const TIMESTAMP_FIELD: &str = "timestamp";
36const DEFAULT_LIST_FIELD_NAME: &str = "item";
37
38#[derive(Debug, Default)]
39pub struct GeoPathAccumulator {
40    lat: Vec<Option<f64>>,
41    lng: Vec<Option<f64>>,
42    timestamp: Vec<Option<i64>>,
43}
44
45impl GeoPathAccumulator {
46    pub fn new() -> Self {
47        Self::default()
48    }
49
50    pub fn uadf_impl() -> AggregateUDF {
51        create_udaf(
52            GEO_PATH_NAME,
53            // Input types: lat, lng, timestamp
54            vec![
55                DataType::Float64,
56                DataType::Float64,
57                DataType::Timestamp(TimeUnit::Nanosecond, None),
58            ],
59            // Output type: list of points {[lat], [lng]}
60            Arc::new(DataType::Struct(
61                vec![
62                    Field::new(
63                        LATITUDE_FIELD,
64                        DataType::List(Arc::new(Field::new(
65                            DEFAULT_LIST_FIELD_NAME,
66                            DataType::Float64,
67                            true,
68                        ))),
69                        false,
70                    ),
71                    Field::new(
72                        LONGITUDE_FIELD,
73                        DataType::List(Arc::new(Field::new(
74                            DEFAULT_LIST_FIELD_NAME,
75                            DataType::Float64,
76                            true,
77                        ))),
78                        false,
79                    ),
80                ]
81                .into(),
82            )),
83            Volatility::Immutable,
84            // Create the accumulator
85            Arc::new(|_| Ok(Box::new(GeoPathAccumulator::new()))),
86            // Intermediate state types
87            Arc::new(vec![DataType::Struct(
88                vec![
89                    Field::new(
90                        LATITUDE_FIELD,
91                        DataType::List(Arc::new(Field::new(
92                            DEFAULT_LIST_FIELD_NAME,
93                            DataType::Float64,
94                            true,
95                        ))),
96                        false,
97                    ),
98                    Field::new(
99                        LONGITUDE_FIELD,
100                        DataType::List(Arc::new(Field::new(
101                            DEFAULT_LIST_FIELD_NAME,
102                            DataType::Float64,
103                            true,
104                        ))),
105                        false,
106                    ),
107                    Field::new(
108                        TIMESTAMP_FIELD,
109                        DataType::List(Arc::new(Field::new(
110                            DEFAULT_LIST_FIELD_NAME,
111                            DataType::Int64,
112                            true,
113                        ))),
114                        false,
115                    ),
116                ]
117                .into(),
118            )]),
119        )
120    }
121}
122
123impl DfAccumulator for GeoPathAccumulator {
124    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion::error::Result<()> {
125        if values.len() != 3 {
126            return Err(DataFusionError::Internal(format!(
127                "Expected 3 columns for geo_path, got {}",
128                values.len()
129            )));
130        }
131
132        let lat_array = as_primitive_array::<Float64Type>(&values[0])?;
133        let lng_array = as_primitive_array::<Float64Type>(&values[1])?;
134        let ts_array = as_primitive_array::<TimestampNanosecondType>(&values[2])?;
135
136        let size = lat_array.len();
137        self.lat.reserve(size);
138        self.lng.reserve(size);
139
140        for idx in 0..size {
141            self.lat.push(if lat_array.is_null(idx) {
142                None
143            } else {
144                Some(lat_array.value(idx))
145            });
146
147            self.lng.push(if lng_array.is_null(idx) {
148                None
149            } else {
150                Some(lng_array.value(idx))
151            });
152
153            self.timestamp.push(if ts_array.is_null(idx) {
154                None
155            } else {
156                Some(ts_array.value(idx))
157            });
158        }
159
160        Ok(())
161    }
162
163    fn evaluate(&mut self) -> DfResult<ScalarValue> {
164        let unordered_lng_array = Float64Array::from(self.lng.clone());
165        let unordered_lat_array = Float64Array::from(self.lat.clone());
166        let ts_array = Int64Array::from(self.timestamp.clone());
167
168        let ordered_indices = sort_to_indices(&ts_array, None, None)?;
169        let lat_array = compute::take(&unordered_lat_array, &ordered_indices, None)?;
170        let lng_array = compute::take(&unordered_lng_array, &ordered_indices, None)?;
171
172        let lat_list = Arc::new(SingleRowListArrayBuilder::new(lat_array).build_list_array());
173        let lng_list = Arc::new(SingleRowListArrayBuilder::new(lng_array).build_list_array());
174
175        let result = ScalarValue::Struct(Arc::new(StructArray::new(
176            vec![
177                Field::new(
178                    LATITUDE_FIELD,
179                    DataType::List(Arc::new(Field::new("item", DataType::Float64, true))),
180                    false,
181                ),
182                Field::new(
183                    LONGITUDE_FIELD,
184                    DataType::List(Arc::new(Field::new("item", DataType::Float64, true))),
185                    false,
186                ),
187            ]
188            .into(),
189            vec![lat_list, lng_list],
190            None,
191        )));
192
193        Ok(result)
194    }
195
196    fn size(&self) -> usize {
197        // Base size of GeoPathAccumulator struct fields
198        let mut total_size = std::mem::size_of::<Self>();
199
200        // Size of vectors (approximation)
201        total_size += self.lat.capacity() * std::mem::size_of::<Option<f64>>();
202        total_size += self.lng.capacity() * std::mem::size_of::<Option<f64>>();
203        total_size += self.timestamp.capacity() * std::mem::size_of::<Option<i64>>();
204
205        total_size
206    }
207
208    fn state(&mut self) -> datafusion::error::Result<Vec<ScalarValue>> {
209        let lat_array = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
210            Some(self.lat.clone()),
211        ]));
212        let lng_array = Arc::new(ListArray::from_iter_primitive::<Float64Type, _, _>(vec![
213            Some(self.lng.clone()),
214        ]));
215        let ts_array = Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
216            Some(self.timestamp.clone()),
217        ]));
218
219        let state_struct = StructArray::new(
220            vec![
221                Field::new(
222                    LATITUDE_FIELD,
223                    DataType::List(Arc::new(Field::new("item", DataType::Float64, true))),
224                    false,
225                ),
226                Field::new(
227                    LONGITUDE_FIELD,
228                    DataType::List(Arc::new(Field::new("item", DataType::Float64, true))),
229                    false,
230                ),
231                Field::new(
232                    TIMESTAMP_FIELD,
233                    DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
234                    false,
235                ),
236            ]
237            .into(),
238            vec![lat_array, lng_array, ts_array],
239            None,
240        );
241
242        Ok(vec![ScalarValue::Struct(Arc::new(state_struct))])
243    }
244
245    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion::error::Result<()> {
246        if states.len() != 1 {
247            return Err(DataFusionError::Internal(format!(
248                "Expected 1 states for geo_path, got {}",
249                states.len()
250            )));
251        }
252
253        let state = as_struct_array(&states[0])?;
254        let lat_lists = as_list_array(state.column(0))?;
255        let lng_lists = as_list_array(state.column(1))?;
256        let ts_lists = as_list_array(state.column(2))?;
257        for row in 0..state.len() {
258            if state.is_null(row) {
259                continue;
260            }
261            let lat_list = lat_lists.value(row);
262            let lng_list = lng_lists.value(row);
263            let ts_list = ts_lists.value(row);
264
265            self.lat
266                .extend(as_primitive_array::<Float64Type>(&lat_list)?);
267            self.lng
268                .extend(as_primitive_array::<Float64Type>(&lng_list)?);
269            self.timestamp
270                .extend(as_primitive_array::<Int64Type>(&ts_list)?);
271        }
272
273        Ok(())
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use datafusion::arrow::array::{Float64Array, TimestampNanosecondArray};
280    use datafusion::scalar::ScalarValue;
281
282    use super::*;
283
284    #[test]
285    fn test_geo_path_basic() {
286        let mut accumulator = GeoPathAccumulator::new();
287
288        // Create test data
289        let lat_array = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]));
290        let lng_array = Arc::new(Float64Array::from(vec![4.0, 5.0, 6.0]));
291        let ts_array = Arc::new(TimestampNanosecondArray::from(vec![100, 200, 300]));
292
293        // Update batch
294        accumulator
295            .update_batch(&[lat_array, lng_array, ts_array])
296            .unwrap();
297
298        // Evaluate
299        let result = accumulator.evaluate().unwrap();
300        if let ScalarValue::Struct(struct_array) = result {
301            // Verify structure
302            let fields = struct_array.fields().clone();
303            assert_eq!(fields.len(), 2);
304            assert_eq!(fields[0].name(), LATITUDE_FIELD);
305            assert_eq!(fields[1].name(), LONGITUDE_FIELD);
306
307            // Verify data
308            let columns = struct_array.columns();
309            assert_eq!(columns.len(), 2);
310
311            // Check latitude values
312            let lat_list = as_list_array(&columns[0]).unwrap().value(0);
313            let lat_array = as_primitive_array::<Float64Type>(&lat_list).unwrap();
314            assert_eq!(lat_array.len(), 3);
315            assert_eq!(lat_array.value(0), 1.0);
316            assert_eq!(lat_array.value(1), 2.0);
317            assert_eq!(lat_array.value(2), 3.0);
318
319            // Check longitude values
320            let lng_list = as_list_array(&columns[1]).unwrap().value(0);
321            let lng_array = as_primitive_array::<Float64Type>(&lng_list).unwrap();
322            assert_eq!(lng_array.len(), 3);
323            assert_eq!(lng_array.value(0), 4.0);
324            assert_eq!(lng_array.value(1), 5.0);
325            assert_eq!(lng_array.value(2), 6.0);
326        } else {
327            panic!("Expected Struct scalar value");
328        }
329    }
330
331    #[test]
332    fn test_geo_path_sort_by_timestamp() {
333        let mut accumulator = GeoPathAccumulator::new();
334
335        // Create test data with unordered timestamps
336        let lat_array = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]));
337        let lng_array = Arc::new(Float64Array::from(vec![4.0, 5.0, 6.0]));
338        let ts_array = Arc::new(TimestampNanosecondArray::from(vec![300, 100, 200]));
339
340        // Update batch
341        accumulator
342            .update_batch(&[lat_array, lng_array, ts_array])
343            .unwrap();
344
345        // Evaluate
346        let result = accumulator.evaluate().unwrap();
347        if let ScalarValue::Struct(struct_array) = result {
348            // Extract arrays
349            let columns = struct_array.columns();
350
351            // Check latitude values
352            let lat_list = as_list_array(&columns[0]).unwrap().value(0);
353            let lat_array = as_primitive_array::<Float64Type>(&lat_list).unwrap();
354            assert_eq!(lat_array.len(), 3);
355            assert_eq!(lat_array.value(0), 2.0); // timestamp 100
356            assert_eq!(lat_array.value(1), 3.0); // timestamp 200
357            assert_eq!(lat_array.value(2), 1.0); // timestamp 300
358
359            // Check longitude values (should be sorted by timestamp)
360            let lng_list = as_list_array(&columns[1]).unwrap().value(0);
361            let lng_array = as_primitive_array::<Float64Type>(&lng_list).unwrap();
362            assert_eq!(lng_array.len(), 3);
363            assert_eq!(lng_array.value(0), 5.0); // timestamp 100
364            assert_eq!(lng_array.value(1), 6.0); // timestamp 200
365            assert_eq!(lng_array.value(2), 4.0); // timestamp 300
366        } else {
367            panic!("Expected Struct scalar value");
368        }
369    }
370
371    #[test]
372    fn test_geo_path_merge() {
373        let mut accumulator1 = GeoPathAccumulator::new();
374        let mut accumulator2 = GeoPathAccumulator::new();
375
376        // Create test data for first accumulator
377        let lat_array1 = Arc::new(Float64Array::from(vec![1.0]));
378        let lng_array1 = Arc::new(Float64Array::from(vec![4.0]));
379        let ts_array1 = Arc::new(TimestampNanosecondArray::from(vec![100]));
380
381        // Create test data for second accumulator
382        let lat_array2 = Arc::new(Float64Array::from(vec![2.0]));
383        let lng_array2 = Arc::new(Float64Array::from(vec![5.0]));
384        let ts_array2 = Arc::new(TimestampNanosecondArray::from(vec![200]));
385
386        // Update batches
387        accumulator1
388            .update_batch(&[lat_array1, lng_array1, ts_array1])
389            .unwrap();
390        accumulator2
391            .update_batch(&[lat_array2, lng_array2, ts_array2])
392            .unwrap();
393
394        // Get states
395        let state1 = accumulator1.state().unwrap();
396        let state2 = accumulator2.state().unwrap();
397
398        // Create a merged accumulator
399        let mut merged = GeoPathAccumulator::new();
400
401        // Extract the struct arrays from the states
402        let state_array1 = match &state1[0] {
403            ScalarValue::Struct(array) => array.clone(),
404            _ => panic!("Expected Struct scalar value"),
405        };
406
407        let state_array2 = match &state2[0] {
408            ScalarValue::Struct(array) => array.clone(),
409            _ => panic!("Expected Struct scalar value"),
410        };
411
412        // States from different partitions arrive as rows of one batch
413        let states = compute::concat(&[state_array1.as_ref(), state_array2.as_ref()]).unwrap();
414        merged.merge_batch(&[states]).unwrap();
415
416        // Evaluate merged result
417        let result = merged.evaluate().unwrap();
418        if let ScalarValue::Struct(struct_array) = result {
419            // Extract arrays
420            let columns = struct_array.columns();
421
422            // Check latitude values
423            let lat_list = as_list_array(&columns[0]).unwrap().value(0);
424            let lat_array = as_primitive_array::<Float64Type>(&lat_list).unwrap();
425            assert_eq!(lat_array.len(), 2);
426            assert_eq!(lat_array.value(0), 1.0); // timestamp 100
427            assert_eq!(lat_array.value(1), 2.0); // timestamp 200
428
429            // Check longitude values (should be sorted by timestamp)
430            let lng_list = as_list_array(&columns[1]).unwrap().value(0);
431            let lng_array = as_primitive_array::<Float64Type>(&lng_list).unwrap();
432            assert_eq!(lng_array.len(), 2);
433            assert_eq!(lng_array.value(0), 4.0); // timestamp 100
434            assert_eq!(lng_array.value(1), 5.0); // timestamp 200
435        } else {
436            panic!("Expected Struct scalar value");
437        }
438    }
439}
common_function/aggrs/geo/geo_path.rs

common_function/aggrs/geo/
geo_path.rs
