1use std::cmp::Ordering;
5
6use icu_collator::options::CollatorOptions;
7use icu_collator::{Collator, CollatorBorrowed};
8
9use crate::Column;
10use crate::error::{PondError, Result};
11use crate::series::{Key, Labels, TimeSeries};
12use arrow_array::UInt32Array;
13
14#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
16pub enum JoinType {
17 Inner,
19 Left,
21 Right,
23 #[default]
25 Outer,
26}
27
28#[derive(Debug, Clone, Default, PartialEq, Eq)]
31pub enum OnConflict {
32 #[default]
34 Error,
35 Prefix(Vec<String>),
40}
41
42#[derive(Debug, Clone, Default, PartialEq, Eq)]
44pub struct Join {
45 pub how: JoinType,
47 pub on_conflict: OnConflict,
49}
50
51impl TimeSeries {
52 pub fn join(&self, other: &TimeSeries, options: Join) -> Result<TimeSeries> {
91 let prepared = prepare(&[self, other], &options.on_conflict)?;
92 let (left, right) = match &prepared {
93 Some(p) => (&p[0], &p[1]),
94 None => (self, other),
95 };
96 join_two(left, right, options.how)
97 }
98
99 pub fn join_many(series: &[&TimeSeries], options: Join) -> Result<TimeSeries> {
115 let Some((first, rest)) = series.split_first() else {
116 return Err(PondError::InvalidData {
117 call: "join_many",
118 message: "needs at least one series".into(),
119 });
120 };
121 let prepared = prepare(series, &options.on_conflict)?;
122 let all: Vec<&TimeSeries> = match &prepared {
123 Some(p) => p.iter().collect(),
124 None => std::iter::once(*first)
125 .chain(rest.iter().copied())
126 .collect(),
127 };
128 let mut joined = all[0].clone();
129 for next in &all[1..] {
130 joined = join_two(&joined, next, options.how)?;
131 }
132 Ok(joined)
133 }
134}
135
136fn prepare(series: &[&TimeSeries], on_conflict: &OnConflict) -> Result<Option<Vec<TimeSeries>>> {
140 const CALL: &str = "join";
141 let mut seen: Vec<&str> = Vec::new();
143 let mut repeated: Vec<&str> = Vec::new();
144 for s in series {
145 for c in s.columns() {
146 let name = c.name();
147 if seen.contains(&name) {
148 if !repeated.contains(&name) {
149 repeated.push(name);
150 }
151 } else {
152 seen.push(name);
153 }
154 }
155 }
156 if repeated.is_empty() {
157 return Ok(None);
158 }
159 repeated.sort_unstable();
160 let prefixes = match on_conflict {
161 OnConflict::Error => {
162 return Err(PondError::InvalidData {
163 call: CALL,
164 message: format!(
165 "cannot join series with duplicate column names: {}; rename one side or use OnConflict::Prefix",
166 repeated.join(", ")
167 ),
168 });
169 }
170 OnConflict::Prefix(p) => p,
171 };
172 if prefixes.len() != series.len() {
173 return Err(PondError::InvalidData {
174 call: CALL,
175 message: format!(
176 "prefix conflict handling requires exactly {} prefixes",
177 series.len()
178 ),
179 });
180 }
181 let renamed = series
182 .iter()
183 .zip(prefixes)
184 .map(|(s, prefix)| {
185 let names: Vec<(String, String)> = s
186 .columns()
187 .iter()
188 .map(|c| c.name())
189 .filter(|n| repeated.contains(n))
190 .map(|n| (n.to_string(), format!("{prefix}_{n}")))
191 .collect();
192 let renames: Vec<(&str, &str)> = names
193 .iter()
194 .map(|(a, b)| (a.as_str(), b.as_str()))
195 .collect();
196 s.rename(&renames)
197 })
198 .collect::<Result<Vec<_>>>()?;
199 let mut names: Vec<&str> = Vec::new();
201 let mut still: Vec<&str> = Vec::new();
202 for s in &renamed {
203 for c in s.columns() {
204 if names.contains(&c.name()) {
205 still.push(c.name());
206 } else {
207 names.push(c.name());
208 }
209 }
210 }
211 if !still.is_empty() {
212 still.sort_unstable();
213 still.dedup();
214 return Err(PondError::InvalidData {
215 call: CALL,
216 message: format!(
217 "prefix conflict handling still produced duplicate column names: {}",
218 still.join(", ")
219 ),
220 });
221 }
222 Ok(Some(renamed))
223}
224
225fn key_kind(key: &Key) -> &'static str {
226 match key {
227 Key::Time(_) => "time",
228 Key::TimeRange(_) => "timeRange",
229 Key::Interval(_) => "interval",
230 }
231}
232
233fn join_two(left: &TimeSeries, right: &TimeSeries, how: JoinType) -> Result<TimeSeries> {
235 const CALL: &str = "join";
236 let (lk, rk) = (left.key(), right.key());
237 if key_kind(lk) != key_kind(rk) {
238 return Err(PondError::InvalidData {
239 call: CALL,
240 message: format!(
241 "cannot join series with different key kinds ({} and {})",
242 key_kind(lk),
243 key_kind(rk)
244 ),
245 });
246 }
247 let key_name = left.key_name();
249 if right.columns().iter().any(|c| c.name() == key_name) {
250 return Err(PondError::InvalidData {
251 call: CALL,
252 message: format!(
253 "the right series has a value column named '{key_name}', the left key's name; rename it first"
254 ),
255 });
256 }
257
258 if lk.len().max(rk.len()) >= NONE as usize {
259 return Err(PondError::Unsupported {
260 call: CALL,
261 message: format!("a side has {} rows or more", NONE),
262 });
263 }
264 let (picks, left_only, right_only) = match_rows(lk, rk, how);
265 let (n, m, len) = (lk.len(), rk.len(), picks.len());
266 let left_identity = len == n && right_only == 0;
269 let right_identity = len == m && left_only == 0;
270
271 let key = if left_identity {
272 lk.clone()
273 } else {
274 output_key(lk, rk, &picks)?
275 };
276 let side = |s: &TimeSeries, identity: bool, pick: fn(Pair) -> u32| {
277 let mut indices: Option<UInt32Array> = None;
280 s.columns()
281 .iter()
282 .map(|c| {
283 let data = if identity {
284 c.data().clone()
285 } else {
286 let idx = indices.get_or_insert_with(|| {
287 picks
288 .iter()
289 .map(|&p| Some(pick(p)).filter(|&i| i != NONE))
290 .collect()
291 });
292 c.data().take_indices(idx)
293 };
294 Column::from_data(c.name().to_string(), data, false)
295 })
296 .collect::<Vec<_>>()
297 };
298 let mut columns = side(left, left_identity, |p| p.l);
299 columns.extend(side(right, right_identity, |p| p.r));
300 Ok(TimeSeries::from_parts(
301 left.name().to_string(),
302 key_name.to_string(),
303 key,
304 columns,
305 ))
306}
307
308#[derive(Clone, Copy)]
312struct Pair {
313 l: u32,
314 r: u32,
315}
316
317const NONE: u32 = u32::MAX;
318
319impl Pair {
320 fn both(i: usize, j: usize) -> Self {
321 Pair {
322 l: i as u32,
323 r: j as u32,
324 }
325 }
326 fn left_only(i: usize) -> Self {
327 Pair {
328 l: i as u32,
329 r: NONE,
330 }
331 }
332 fn right_only(j: usize) -> Self {
333 Pair {
334 l: NONE,
335 r: j as u32,
336 }
337 }
338 fn left(self) -> Option<usize> {
339 (self.l != NONE).then_some(self.l as usize)
340 }
341}
342
343fn match_rows(lk: &Key, rk: &Key, how: JoinType) -> (Vec<Pair>, usize, usize) {
346 let (n, m) = (lk.len(), rk.len());
347 match (lk, rk) {
348 (Key::Time(l), Key::Time(r)) => walk(n, m, how, |i, j| l[i].cmp(&r[j])),
350 _ => {
351 let mut order = KeyOrder::new(lk, rk);
352 walk(n, m, how, |i, j| order.cmp(i, j))
353 }
354 }
355}
356
357fn walk(
358 n: usize,
359 m: usize,
360 how: JoinType,
361 mut compare: impl FnMut(usize, usize) -> Ordering,
362) -> (Vec<Pair>, usize, usize) {
363 let keep_left = matches!(how, JoinType::Left | JoinType::Outer);
364 let keep_right = matches!(how, JoinType::Right | JoinType::Outer);
365 let mut picks = Vec::with_capacity(match how {
366 JoinType::Left => n,
367 JoinType::Right => m,
368 JoinType::Inner => n.min(m),
369 JoinType::Outer => n + m,
370 });
371 let (mut i, mut j) = (0, 0);
372 let (mut left_only, mut right_only) = (0, 0);
373 while i < n && j < m {
374 match compare(i, j) {
375 Ordering::Equal => {
376 picks.push(Pair::both(i, j));
377 i += 1;
378 j += 1;
379 }
380 Ordering::Less => {
381 if keep_left {
382 picks.push(Pair::left_only(i));
383 left_only += 1;
384 }
385 i += 1;
386 }
387 Ordering::Greater => {
388 if keep_right {
389 picks.push(Pair::right_only(j));
390 right_only += 1;
391 }
392 j += 1;
393 }
394 }
395 }
396 if keep_left {
397 left_only += n - i;
398 picks.extend((i..n).map(Pair::left_only));
399 }
400 if keep_right {
401 right_only += m - j;
402 picks.extend((j..m).map(Pair::right_only));
403 }
404 (picks, left_only, right_only)
405}
406
407struct KeyOrder<'a> {
410 lb: &'a [i64],
411 rb: &'a [i64],
412 ends: Option<(&'a [i64], &'a [i64])>,
413 labels: Option<(&'a Labels, &'a Labels)>,
414 collator: Option<CollatorBorrowed<'static>>,
415}
416
417impl<'a> KeyOrder<'a> {
418 fn new(lk: &'a Key, rk: &'a Key) -> Self {
419 let labels = match (lk, rk) {
420 (Key::Interval(l), Key::Interval(r)) => Some((&l.label, &r.label)),
421 _ => None,
422 };
423 Self {
424 lb: lk.begin(),
425 rb: rk.begin(),
426 ends: lk.end().zip(rk.end()),
427 labels,
428 collator: None,
429 }
430 }
431
432 fn cmp(&mut self, i: usize, j: usize) -> Ordering {
433 let c = self.lb[i].cmp(&self.rb[j]);
434 if c != Ordering::Equal {
435 return c;
436 }
437 let Some((le, re)) = self.ends else {
438 return Ordering::Equal;
439 };
440 let c = le[i].cmp(&re[j]);
441 if c != Ordering::Equal {
442 return c;
443 }
444 let Some((ll, rl)) = self.labels else {
445 return Ordering::Equal;
446 };
447 match (ll, rl) {
448 (Labels::Number(a), Labels::Number(b)) => {
450 a[i].partial_cmp(&b[j]).unwrap_or(Ordering::Equal)
451 }
452 (Labels::Number(_), Labels::String(_)) => Ordering::Less,
454 (Labels::String(_), Labels::Number(_)) => Ordering::Greater,
455 (Labels::String(a), Labels::String(b)) => {
456 let (a, b) = (a.get(i).unwrap_or(""), b.get(j).unwrap_or(""));
457 if a == b {
458 return Ordering::Equal;
459 }
460 self.collator
461 .get_or_insert_with(root_collator)
462 .compare(a, b)
463 }
464 }
465 }
466}
467
468pub(crate) fn root_collator() -> CollatorBorrowed<'static> {
472 Collator::try_new(Default::default(), CollatorOptions::default())
473 .expect("the root collation is compiled in")
474}
475
476fn output_key(lk: &Key, rk: &Key, picks: &[Pair]) -> Result<Key> {
479 let pick = |l: &[i64], r: &[i64]| -> Vec<i64> {
480 picks
481 .iter()
482 .map(|p| match p.left() {
483 Some(i) => l[i],
484 None => r[p.r as usize],
485 })
486 .collect()
487 };
488 match (lk, rk) {
489 (Key::Time(l), Key::Time(r)) => Ok(Key::Time(pick(l, r))),
490 (Key::TimeRange(l), Key::TimeRange(r)) => Ok(Key::TimeRange(crate::TimeRangeKeys {
491 begin: pick(&l.begin, &r.begin),
492 end: pick(&l.end, &r.end),
493 })),
494 (Key::Interval(l), Key::Interval(r)) => {
495 let from_left = picks.iter().filter(|p| p.l != NONE).count();
496 if from_left != 0
499 && from_left != picks.len()
500 && matches!(l.label, Labels::Number(_)) != matches!(r.label, Labels::Number(_))
501 {
502 let kind = |k: &Labels| match k {
503 Labels::Number(_) => "number",
504 Labels::String(_) => "string",
505 };
506 return Err(PondError::InvalidData {
507 call: "join",
508 message: format!(
509 "cannot combine interval keys with {} labels and interval keys with {} labels; an interval-keyed series must use one label type throughout",
510 kind(&l.label),
511 kind(&r.label)
512 ),
513 });
514 }
515 let picks: Vec<(u32, u32)> = picks
518 .iter()
519 .map(|p| if p.l != NONE { (0, p.l) } else { (1, p.r) })
520 .collect();
521 let sources: [&Key; 2] = [lk, rk];
522 if from_left == 0 && !picks.is_empty() {
525 let rows: Vec<(u32, u32)> = picks.iter().map(|&(_, r)| (0, r)).collect();
526 return Ok(Key::gather(&[rk], &rows));
527 }
528 Ok(Key::gather(&sources, &picks))
529 }
530 _ => unreachable!("key kinds were checked"),
531 }
532}
533
534#[cfg(test)]
535mod tests {
536 use super::*;
537 use crate::column::DecimalColumn;
538 use std::sync::Arc;
539
540 fn s(name: &str, times: Vec<i64>, col: &str, v: Vec<f64>) -> TimeSeries {
541 TimeSeries::new(name, times, vec![Column::required(col, v)]).unwrap()
542 }
543
544 #[test]
545 fn a_kept_side_shares_its_columns() {
546 let left = s("l", vec![0, 1, 2], "a", vec![1.0, 2.0, 3.0]);
547 let right = s("r", vec![1, 5], "b", vec![4.0, 5.0]);
548 let ptr =
549 |t: &TimeSeries, c: &str| t.column(c).unwrap().numbers().unwrap().values().as_ptr();
550 let out = left
551 .join(
552 &right,
553 Join {
554 how: JoinType::Left,
555 ..Join::default()
556 },
557 )
558 .unwrap();
559 assert_eq!(ptr(&out, "a"), ptr(&left, "a"));
560 let out = left
561 .join(
562 &right,
563 Join {
564 how: JoinType::Right,
565 ..Join::default()
566 },
567 )
568 .unwrap();
569 assert_eq!(ptr(&out, "b"), ptr(&right, "b"));
570 let same = s("m", vec![0, 1, 2], "c", vec![7.0, 8.0, 9.0]);
572 let out = left.join(&same, Join::default()).unwrap();
573 assert_eq!(ptr(&out, "a"), ptr(&left, "a"));
574 assert_eq!(ptr(&out, "c"), ptr(&same, "c"));
575 assert!(out.schema().iter().skip(1).all(|c| !c.required));
577 }
578
579 #[test]
580 fn decimal_columns_are_gathered_exactly() {
581 let left = s("l", vec![0, 2], "a", vec![1.0, 2.0]);
582 let d = DecimalColumn::from_options([Some(12_345), None, Some(-7)], 2).unwrap();
583 let right = TimeSeries::new("r", vec![1, 2, 3], vec![Column::optional("d", d)]).unwrap();
584 let out = left.join(&right, Join::default()).unwrap();
585 assert_eq!(out.key().begin(), &[0, 1, 2, 3]);
586 let got: Vec<Option<String>> = (0..4)
587 .map(|i| {
588 out.column("d")
589 .unwrap()
590 .decimals()
591 .unwrap()
592 .get(i)
593 .map(|x| x.to_string())
594 })
595 .collect();
596 assert_eq!(
597 got,
598 [None, Some("123.45".into()), None, Some("-0.07".into())]
599 );
600 }
601
602 #[test]
603 fn a_value_column_named_like_the_left_key_is_refused() {
604 use arrow_array::{Float64Array, RecordBatch, TimestampMillisecondArray};
606 use arrow_schema::{DataType, Field, Schema, TimeUnit};
607 let schema = Arc::new(Schema::new(vec![
608 Field::new(
609 "ts",
610 DataType::Timestamp(TimeUnit::Millisecond, None),
611 false,
612 ),
613 Field::new("v", DataType::Float64, false),
614 ]));
615 let batch = RecordBatch::try_new(
616 schema,
617 vec![
618 Arc::new(TimestampMillisecondArray::from(vec![0, 1])),
619 Arc::new(Float64Array::from(vec![1.0, 2.0])),
620 ],
621 )
622 .unwrap();
623 let left = TimeSeries::from_arrow("l", &[batch], "ts").unwrap();
624 let right = s("r", vec![0], "ts", vec![9.0]);
625 let e = left.join(&right, Join::default()).unwrap_err().to_string();
626 assert!(e.contains("named 'ts'"), "{e}");
627 assert_eq!(
629 right.join(&left, Join::default()).unwrap().key_name(),
630 "time"
631 );
632 }
633
634 #[test]
635 fn an_empty_interval_join_keeps_the_left_label_kind() {
636 let one = |label: Labels| {
637 let keys = crate::IntervalKeys {
638 label,
639 begin: vec![0],
640 end: vec![10],
641 };
642 TimeSeries::from_intervals("i", keys, vec![]).unwrap()
643 };
644 let text = one(vec!["x"].into());
645 let number = one(vec![5.0].into());
646 let inner = Join {
647 how: JoinType::Inner,
648 ..Join::default()
649 };
650 let is_text = |t: &TimeSeries| match t.key() {
651 Key::Interval(k) => matches!(k.label, Labels::String(_)),
652 _ => unreachable!(),
653 };
654 let out = text.join(&number, inner.clone()).unwrap();
655 assert_eq!(out.len(), 0);
656 assert!(is_text(&out));
657 let out = number.join(&text, inner).unwrap();
658 assert_eq!(out.len(), 0);
659 assert!(!is_text(&out));
660 }
661
662 #[test]
663 fn join_many_needs_a_series() {
664 assert!(TimeSeries::join_many(&[], Join::default()).is_err());
665 }
666
667 #[test]
668 fn labels_compare_as_javascript_locale_compare() {
669 let c = root_collator();
670 assert_eq!(c.compare("a", "B"), Ordering::Less);
671 assert_eq!(c.compare("B", "b"), Ordering::Greater);
672 assert_eq!(c.compare("\u{e9}", "e\u{301}"), Ordering::Equal);
673 assert_eq!(c.compare("_", "1"), Ordering::Less);
674 assert_eq!(c.compare("1", "a"), Ordering::Less);
675 }
676}