diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index e22064dc0..ddb1b8308 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -52,11 +52,13 @@ mod write_routing; mod partition; mod snapshot; mod tag; +mod temporal; #[pymodule] fn pypaimon_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { error::register_module(py, m)?; context::register_module(py, m)?; + m.add_function(wrap_pyfunction!(temporal::temporal_match, m)?)?; variant::register_module(py, m)?; evolution::register_module(m)?; Ok(()) diff --git a/bindings/python/src/temporal.rs b/bindings/python/src/temporal.rs new file mode 100644 index 000000000..2f8d56757 --- /dev/null +++ b/bindings/python/src/temporal.rs @@ -0,0 +1,297 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Batch temporal matching over sorted, non-null int64 keys. + +use arrow::array::{Array, ArrayData, Int64Array}; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; +use arrow::record_batch::RecordBatch; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use std::sync::Arc; + +fn lower(times: &[i64], target: i128, mut start: usize, mut end: usize) -> usize { + while start < end { + let middle = start + (end - start) / 2; + if (times[middle] as i128) < target { + start = middle + 1; + } else { + end = middle; + } + } + start +} + +fn upper(times: &[i64], target: i128, mut start: usize, mut end: usize) -> usize { + while start < end { + let middle = start + (end - start) / 2; + if (times[middle] as i128) <= target { + start = middle + 1; + } else { + end = middle; + } + } + start +} + +#[derive(Clone, Copy)] +enum Mode { + Nearest, + Backward, + Forward, + Interpolate, + Window, +} + +struct Options { + mode: Mode, + exact: bool, + tolerance: Option, + preceding: i64, + following: i64, + left_closed: bool, + right_closed: bool, +} + +fn match_one( + times: &[i64], + target: i64, + start: usize, + end: usize, + options: &Options, +) -> (i64, i64) { + if start == end { + return (-1, -1); + } + let target = target as i128; + if matches!(options.mode, Mode::Window) { + let left = target - options.preceding as i128; + let right = target + options.following as i128; + let first = if options.left_closed { + lower(times, left, start, end) + } else { + upper(times, left, start, end) + }; + let last = if options.right_closed { + upper(times, right, start, end) + } else { + lower(times, right, start, end) + }; + return if last > first { + (first as i64, last as i64) + } else { + (0, 0) + }; + } + let left = lower(times, target, start, end); + let right = upper(times, target, left, end); + if matches!(options.mode, Mode::Interpolate) { + if left < end && times[left] as i128 == target { + return ((right - 1) as i64, (right - 1) as i64); + } + if left > start + && left < end + && options.tolerance.is_none_or(|tolerance| { + (target - times[left - 1] as i128).max(times[left] as i128 - target) + <= tolerance as i128 + }) + { + return ((left - 1) as i64, left as i64); + } + return (-1, -1); + } + let before = if options.exact { right } else { left }; + let after = if options.exact { left } else { right }; + let position = match options.mode { + Mode::Backward => (before > start).then(|| before - 1), + Mode::Forward => (after < end).then_some(after), + _ if options.exact && left < end && times[left] as i128 == target => Some(right - 1), + _ if left == start => (after < end).then_some(after), + _ if after == end || target - times[left - 1] as i128 <= times[after] as i128 - target => { + Some(left - 1) + } + _ => Some(after), + }; + match position { + Some(position) + if options.tolerance.is_none_or(|tolerance| { + (times[position] as i128 - target).abs() <= tolerance as i128 + }) => + { + (position as i64, -1) + } + _ => (-1, -1), + } +} + +fn int64_array(value: &Bound<'_, PyAny>) -> PyResult { + let mut data = ArrayData::from_pyarrow_bound(value)?; + if data.data_type() != &arrow::datatypes::DataType::Int64 || data.null_count() != 0 { + return Err(PyValueError::new_err( + "Expected a non-null int64 Arrow array", + )); + } + data.align_buffers(); + Ok(Int64Array::from(data)) +} + +/// Match each target within its [start, end) group of sorted times. +/// Returns indices, not row IDs; -1 denotes a missing match. Window end indices +/// are exclusive. Equal nearest/interpolation keys select the last duplicate; +/// nearest ties select the earlier key. Arithmetic uses i128 at int64 limits. +#[pyfunction] +#[pyo3(signature = (times, targets, starts, ends, mode, exact=true, tolerance=None, preceding=0, following=0, left_closed=true, right_closed=true))] +#[allow(clippy::too_many_arguments)] +pub fn temporal_match( + py: Python<'_>, + times: &Bound<'_, PyAny>, + targets: &Bound<'_, PyAny>, + starts: &Bound<'_, PyAny>, + ends: &Bound<'_, PyAny>, + mode: &str, + exact: bool, + tolerance: Option, + preceding: i64, + following: i64, + left_closed: bool, + right_closed: bool, +) -> PyResult> { + let times = int64_array(times)?; + let targets = int64_array(targets)?; + let starts = int64_array(starts)?; + let ends = int64_array(ends)?; + let mode = match mode { + "nearest" => Mode::Nearest, + "backward" => Mode::Backward, + "forward" => Mode::Forward, + "interpolate" => Mode::Interpolate, + "window" => Mode::Window, + _ => return Err(PyValueError::new_err("Unknown temporal match mode")), + }; + if targets.len() != starts.len() || targets.len() != ends.len() { + return Err(PyValueError::new_err( + "Target and group-bound arrays must have equal lengths", + )); + } + if tolerance.is_some_and(|value| value < 0) || preceding < 0 || following < 0 { + return Err(PyValueError::new_err( + "Temporal distances must be non-negative", + )); + } + for (&start, &end) in starts.values().iter().zip(ends.values()) { + if start < 0 || end < start || end as u64 > times.len() as u64 { + return Err(PyValueError::new_err("Invalid temporal group bounds")); + } + } + let options = Options { + mode, + exact, + tolerance, + preceding, + following, + left_closed, + right_closed, + }; + let (first, second): (Vec, Vec) = py.detach(|| { + targets + .values() + .iter() + .zip(starts.values()) + .zip(ends.values()) + .map(|((&target, &start), &end)| { + match_one( + times.values(), + target, + start as usize, + end as usize, + &options, + ) + }) + .unzip() + }); + let batch = RecordBatch::try_from_iter(vec![ + ("first", Arc::new(Int64Array::from(first)) as Arc), + ( + "second", + Arc::new(Int64Array::from(second)) as Arc, + ), + ]) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok(batch.to_pyarrow(py)?.unbind()) +} + +#[cfg(test)] +mod tests { + use super::*; + fn options(mode: Mode) -> Options { + Options { + mode, + exact: true, + tolerance: None, + preceding: 0, + following: 0, + left_closed: true, + right_closed: true, + } + } + #[test] + fn duplicate_and_tie_policy() { + let times = [0, 10, 10, 20]; + assert_eq!( + match_one(×, 10, 0, 4, &options(Mode::Nearest)), + (2, -1) + ); + assert_eq!( + match_one(×, 10, 0, 4, &options(Mode::Forward)), + (1, -1) + ); + assert_eq!( + match_one(×, 15, 0, 4, &options(Mode::Nearest)), + (2, -1) + ); + assert_eq!( + match_one(×, 10, 0, 4, &options(Mode::Interpolate)), + (2, 2) + ); + let mut option = options(Mode::Nearest); + option.exact = false; + assert_eq!(match_one(×, 10, 0, 4, &option), (0, -1)); + } + #[test] + fn limits_and_tolerance() { + let times = [i64::MIN, i64::MAX]; + assert_eq!(match_one(×, 0, 0, 2, &options(Mode::Nearest)), (1, -1)); + let mut option = options(Mode::Interpolate); + option.tolerance = Some(i64::MAX); + assert_eq!(match_one(×, 0, 0, 2, &option), (-1, -1)); + assert_eq!(match_one(×, 0, 1, 1, &option), (-1, -1)); + let mut option = options(Mode::Window); + option.preceding = i64::MAX; + option.following = i64::MAX; + assert_eq!(match_one(×, i64::MIN, 0, 2, &option), (0, 1)); + assert_eq!(match_one(×, i64::MAX, 0, 2, &option), (1, 2)); + } + #[test] + fn window_closure_and_groups() { + let times = [0, 10, 10, 20, 0, 5]; + let mut option = options(Mode::Window); + assert_eq!(match_one(×, 10, 0, 4, &option), (1, 3)); + option.left_closed = false; + assert_eq!(match_one(×, 10, 0, 4, &option), (0, 0)); + assert_eq!(match_one(×, 4, 4, 6, &options(Mode::Nearest)), (5, -1)); + } +} diff --git a/bindings/python/tests/test_temporal.py b/bindings/python/tests/test_temporal.py new file mode 100644 index 000000000..2549368eb --- /dev/null +++ b/bindings/python/tests/test_temporal.py @@ -0,0 +1,134 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. + +import bisect +import random +import struct + +import pyarrow as pa +import pytest +from pypaimon_rust import temporal_match + + +def call(times, targets, mode="nearest", starts=None, ends=None, **kwargs): + arrays = [times, targets, starts if starts is not None else [0] * len(targets), + ends if ends is not None else [len(times)] * len(targets)] + return temporal_match(*(pa.array(values, type=pa.int64()) for values in arrays), + mode, **kwargs).to_pydict() + + +@pytest.mark.parametrize("mode", ["nearest", "backward", "forward", "interpolate", "window"]) +def test_random_matching(mode): + rng = random.Random(42) + for _ in range(30): + times = sorted(rng.randrange(-20, 20) for _ in range(40)) + targets = [rng.randrange(-30, 30) for _ in range(50)] + for exact in (False, True): + for tolerance in (None, 0, 3): + actual = call(times, targets, mode, exact=exact, tolerance=tolerance, + preceding=3, following=5) + for i, target in enumerate(targets): + left, right = bisect.bisect_left(times, target), bisect.bisect_right(times, target) + if mode == "window": + expected = (bisect.bisect_left(times, target - 3), + bisect.bisect_right(times, target + 5)) + if expected[0] == expected[1]: + expected = (0, 0) + elif mode == "interpolate": + if left < right: + expected = (right - 1, right - 1) + elif left == 0 or left == len(times): + expected = (-1, -1) + elif tolerance is not None and max(target-times[left-1], times[left]-target) > tolerance: + expected = (-1, -1) + else: + expected = (left - 1, left) + else: + candidates = [(j, value) for j, value in enumerate(times) + if (exact or value != target) + and (mode != "backward" or value <= target) + and (mode != "forward" or value >= target) + and (tolerance is None or abs(value-target) <= tolerance)] + if not candidates: + expected = (-1, -1) + elif mode == "backward": + expected = (candidates[-1][0], -1) + elif mode == "forward": + expected = (candidates[0][0], -1) + else: + value = min((value for _, value in candidates), key=lambda v: (abs(v-target), v)) + indices = [j for j, v in candidates if v == value] + expected = (indices[-1] if value <= target else indices[0], -1) + assert (actual["first"][i], actual["second"][i]) == expected + + +@pytest.mark.parametrize("left", [False, True]) +@pytest.mark.parametrize("right", [False, True]) +def test_window_closure_and_limits(left, right): + low, high = -(1 << 63), (1 << 63) - 1 + times = [low, -1, 0, 0, 1, high] + for target in times: + result = call(times, [target], "window", preceding=high, following=high, + left_closed=left, right_closed=right) + positions = [i for i, value in enumerate(times) + if (value >= target-high if left else value > target-high) + and (value <= target+high if right else value < target+high)] + assert (result["first"][0], result["second"][0]) == ( + (positions[0], positions[-1]+1) if positions else (0, 0)) + + +def test_sliced_arrays_and_group_bounds(): + times = pa.array([999, 1, 3, 0, 2, 999], type=pa.int64()).slice(1, 4) + targets = pa.array([2, 1, 1], type=pa.int64()) + starts = pa.array([0, 2, 0], type=pa.int64()) + ends = pa.array([2, 4, 0], type=pa.int64()) + assert temporal_match(times, targets, starts, ends, "nearest").column(0).to_pylist() == [0, 2, -1] + assert call([], [], "nearest") == {"first": [], "second": []} + + +@pytest.mark.parametrize("unaligned_input", range(4)) +def test_unaligned_int64_buffers(unaligned_input): + values = ([1, 3, 5], [1], [0], [3]) + arrays = [pa.array(items, type=pa.int64()) for items in values] + items = values[unaligned_input] + raw = bytearray(8 + 8 * len(items)) + buffer = pa.py_buffer(raw) + offset = 1 if buffer.address % 8 == 0 else 0 + struct.pack_into("<{}q".format(len(items)), raw, offset, *items) + buffer = buffer.slice(offset, 8 * len(items)) + assert buffer.address % 8 != 0 + arrays[unaligned_input] = pa.Array.from_buffers( + pa.int64(), len(items), [None, buffer]) + arrays[unaligned_input].validate(full=True) + + result = temporal_match(*arrays, "nearest").to_pydict() + assert result == {"first": [0], "second": [-1]} + + +@pytest.mark.parametrize("kwargs", [dict(starts=[-1]), dict(ends=[2]), dict(starts=[1], ends=[0]), + dict(starts=[]), dict(tolerance=-1), dict(preceding=-1), + dict(mode="invalid")]) +def test_invalid_arguments(kwargs): + with pytest.raises(ValueError): + call([0], [0], **kwargs) + + +def test_invalid_arrow_keys(): + good = pa.array([0], type=pa.int64()) + for invalid in (pa.array([None], type=pa.int64()), pa.array([0.0])): + with pytest.raises(ValueError): + temporal_match(invalid, good, good, good, "nearest")