Google Git Google Git
Sign in
apache/paimon/refs/heads/master/./paimon-python/pypaimon/tests/multimodal_temporal_test.py
blob: 8c5199eefff85a7cf6bf423ac01070d7513476f8 [file]
# 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 json
import os
import shutil
import sys
import tempfile
import tracemalloc
import unittest
from datetime import datetime, timedelta
from unittest import mock
import numpy as np
import pyarrow as pa
import pytest
import pypaimon.multimodal as pmm
from pypaimon.multimodal import temporal
from pypaimon.catalog.table_query_auth import TableQueryAuthResult
from pypaimon.execution.python.read.reader.format_pyarrow_reader import (FormatPyArrowReader)
from pypaimon.read.table_scan import TableScan
class MultimodalTemporalTest(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_temporal_")
self.conn = pmm.connect(options={
"warehouse": os.path.join(self.temp_dir, "warehouse"),
})
def tearDown(self):
shutil.rmtree(self.temp_dir, ignore_errors=True)
def test_alignment_requires_an_explicit_group_boundary(self):
table = self._table("missing_group", {
"event_time": pa.int64(),
"value": pa.int32(),
})
with self.assertRaisesRegex(ValueError, "grouping column"):
pmm.join_asof(
table.scan(),
table.scan(),
on="event_time",
by=(),
direction="nearest",
tolerance=0,
)
def test_alignment_preserves_payload_names(self):
anchors = self._table("audit_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("audit_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"valid": pa.bool_(),
"matched_time": pa.int64(),
"time_delta": pa.int64(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
samples.add([{
"episode_id": 1,
"event_time": 100,
"valid": False,
"matched_time": 7,
"time_delta": 8,
}])
row = pmm.join_asof(
anchors.scan(),
samples.scan().select([
"valid", "matched_time", "time_delta"
]),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=0,
).to_list()[0]
self.assertFalse(row["valid"])
self.assertEqual(7, row["matched_time"])
self.assertEqual(8, row["time_delta"])
def test_alignment_handles_duplicate_right_timestamps(self):
anchors = self._table("duplicate_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("duplicate_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
samples.add([
{"episode_id": 1, "event_time": 100, "value": 1},
{"episode_id": 1, "event_time": 100, "value": 2},
])
for direction, expected in (
("backward", 2), ("forward", 1), ("nearest", 2)):
with self.subTest(direction=direction):
row = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction=direction, tolerance=0,
).to_list()[0]
self.assertEqual(expected, row["value"])
def test_alignment_treats_signed_zero_as_one_group(self):
schema = {"group": pa.float64(), "event_time": pa.int64()}
anchors = self._table("zero_anchors", schema)
samples = self._table("zero_samples", {**schema, "value": pa.int32()})
anchors.add([{"group": 0.0, "event_time": time} for time in (1, 2)])
samples.add(
[
{"group": -0.0, "event_time": 1, "value": 7},
{"group": 0.0, "event_time": 2, "value": 8},
]
)
rows = pmm.join_asof(
anchors.scan(),
samples.scan().select("value"),
on="event_time",
by="group",
tolerance=0,
).to_list()
self.assertEqual({1: 7, 2: 8}, {row["event_time"]: row["value"] for row in rows})
def test_alignment_preserves_composite_group_boundaries(self):
schema = {"host": pa.string(), "sensor": pa.int32(), "event_time": pa.int64()}
anchors = self._table("composite_anchors", schema)
samples = self._table("composite_samples", {**schema, "value": pa.int32()})
keys = [("a", 1), ("a", 2), ("b", 1)]
anchors.add([{"host": host, "sensor": sensor, "event_time": 10} for host, sensor in keys])
samples.add(
[
{"host": host, "sensor": sensor, "event_time": 9, "value": index}
for index, (host, sensor) in enumerate(keys)
]
)
rows = pmm.join_asof(
anchors.scan(),
samples.scan().select("value"),
on="event_time",
by=["host", "sensor"],
).to_list()
self.assertEqual(dict(zip(keys, range(3))), {(row["host"], row["sensor"]): row["value"] for row in rows})
def test_alignment_without_run_end_kernel(self):
schema = {"group": pa.string(), "event_time": pa.int64()}
anchors = self._table("fallback_anchors", schema)
samples = self._table("fallback_samples", {**schema, "value": pa.int32()})
anchors.add([{"group": "a", "event_time": 1}])
samples.add(
[
{"group": "a", "event_time": 1, "value": 7},
{"group": "b", "event_time": 1, "value": 8},
]
)
for kernel in (None, mock.Mock(side_effect=pa.ArrowNotImplementedError("No kernel"))):
with self.subTest(kernel=kernel), mock.patch.object(temporal.pc, "run_end_encode", kernel, create=True):
rows = pmm.join_asof(
anchors.scan(),
samples.scan().select("value"),
on="event_time",
by="group",
tolerance=0,
).to_list()
self.assertEqual([7], [row["value"] for row in rows])
def test_asof_excludes_exact_matches_within_group(self):
anchors = self._table("strict_anchors", {
"episode_id": pa.int32(), "event_time": pa.int64(),
})
samples = self._table("strict_samples", {
"episode_id": pa.int32(), "event_time": pa.int64(),
"value": pa.int32(),
})
groups = [
[9, 9, 10, 10, 11, 11],
[10, 10],
[10, 10, 11],
[9, 10, 10],
[],
[8, 11],
]
anchors.add([
{"episode_id": group, "event_time": 10}
for group in range(len(groups))
])
samples.add([
{"episode_id": group, "event_time": time, "value": index}
for group, times in enumerate(groups)
for index, time in enumerate(times)
])
for direction, expected in (
("backward", [1, None, None, 0, None, 0]),
("forward", [4, None, 2, None, None, 1]),
("nearest", [1, None, 2, 0, None, 1])):
with self.subTest(direction=direction):
rows = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id", direction=direction,
allow_exact_matches=False,
).to_list()
self.assertEqual(dict(enumerate(expected)), {
row["episode_id"]: row["value"] for row in rows
})
def test_asof_exact_match_options_are_independent_in_chain(self):
anchors = self._table("strict_chain_anchors", {
"episode_id": pa.int32(), "event_time": pa.int64(),
})
samples = self._table("strict_chain_samples", {
"episode_id": pa.int32(), "sample_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
samples.add([
{"episode_id": 1, "sample_time": time, "value": time}
for time in [9, 10, 10, 11]
])
for direction, expected in (
("backward", 9), ("forward", 11), ("nearest", 9)):
for tolerance in (0, 0.5, 1, None):
with self.subTest(direction=direction, tolerance=tolerance):
aligned = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
right_on="sample_time", direction=direction,
tolerance=tolerance, allow_exact_matches=False,
)
row = aligned.join_asof(
samples.scan().select("value"),
right_on="sample_time", allow_exact_matches=True,
).join_asof(
samples.scan().select("value"),
right_on="sample_time", allow_exact_matches=False,
suffix="_strict",
).to_list()[0]
self.assertEqual(
expected if tolerance in (1, None) else None,
row["value"],
)
self.assertEqual(10, row["value_right"])
self.assertEqual(9, row["value_strict"])
def test_asof_excludes_exact_nanosecond_timestamps(self):
anchors = self._table("strict_ns_anchors", {
"episode_id": pa.int32(), "event_time": pa.timestamp("ns"),
})
samples = self._table("strict_ns_samples", {
"episode_id": pa.int32(), "event_time": pa.timestamp("ns"),
"value": pa.int32(),
})
anchors.add(pa.table({
"episode_id": pa.array([1], type=pa.int32()),
"event_time": pa.array(
[1_000_000_001], type=pa.int64()).cast(pa.timestamp("ns")),
}))
samples.add(pa.table({
"episode_id": pa.array([1, 1, 1], type=pa.int32()),
"event_time": pa.array(
[1_000_000_000, 1_000_000_001, 1_000_000_002],
type=pa.int64()).cast(pa.timestamp("ns")),
"value": pa.array([9, 10, 11], type=pa.int32()),
}))
for direction, expected in (
("backward", 9), ("forward", 11), ("nearest", 9)):
for tolerance in (timedelta(0), timedelta(microseconds=1)):
with self.subTest(direction=direction, tolerance=tolerance):
row = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction=direction, tolerance=tolerance,
allow_exact_matches=False,
).to_list()[0]
self.assertEqual(
expected if tolerance else None, row["value"])
def test_asof_validates_allow_exact_matches(self):
table = self._table("strict_options", {
"episode_id": pa.int32(), "event_time": pa.int64(),
})
aligned = pmm.join_asof(
table.scan(), table.scan(), on="event_time", by="episode_id")
for value in (None, 0, 1, "false"):
with self.subTest(value=value):
with self.assertRaisesRegex(TypeError, "allow_exact_matches"):
pmm.join_asof(
table.scan(), table.scan(),
on="event_time", by="episode_id",
allow_exact_matches=value,
)
with self.assertRaisesRegex(TypeError, "allow_exact_matches"):
aligned.join_asof(table.scan(), allow_exact_matches=value)
def test_nearest_uses_candidate_side_for_duplicate_timestamps(self):
anchors = self._table("duplicate_nearest_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("duplicate_nearest_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([
{"episode_id": 1, "event_time": 90},
{"episode_id": 1, "event_time": 110},
])
samples.add([
{"episode_id": 1, "event_time": 100, "value": 1},
{"episode_id": 1, "event_time": 100, "value": 2},
])
rows = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=20,
).to_list()
self.assertEqual(
{90: 1, 110: 2},
{row["event_time"]: row["value"] for row in rows},
)
def test_linear_interpolation_stays_in_group_without_extrapolation(self):
anchors = self._table("linear_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([
{"episode_id": 1, "event_time": 5},
{"episode_id": 1, "event_time": 10},
{"episode_id": 1, "event_time": 20},
{"episode_id": 2, "event_time": 5},
])
states.add([
{"episode_id": 1, "event_time": 0, "value": 0},
{"episode_id": 1, "event_time": 10, "value": 20},
{"episode_id": 1, "event_time": 10, "value": 30},
{"episode_id": 2, "event_time": 0, "value": 100},
{"episode_id": 2, "event_time": 10, "value": 120},
])
result = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id", tolerance=5,
)
rows = sorted(
result.to_list(),
key=lambda row: (row["episode_id"], row["event_time"]),
)
self.assertIsInstance(result, pmm.TemporalAlignment)
self.assertEqual(pa.float64(), result.schema.field("value").type)
self.assertEqual([10.0, 30.0, None, 110.0], [
row["value"] for row in rows
])
def test_window_join_stays_in_group_and_skips_nulls(self):
anchors = self._table("window_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
"label": pa.string(),
})
anchors.add([
{"episode_id": 1, "event_time": 10},
{"episode_id": 2, "event_time": 10},
{"episode_id": 3, "event_time": 10},
])
samples.add([
{"episode_id": 1, "event_time": 5,
"value": 1, "label": "a"},
{"episode_id": 1, "event_time": 10,
"value": None, "label": None},
{"episode_id": 1, "event_time": 15,
"value": 5, "label": "c"},
{"episode_id": 2, "event_time": 10,
"value": 100, "label": "z"},
])
result = pmm.join_window(
anchors.scan(),
samples.scan().select(["value", "label"]),
on="event_time",
by="episode_id",
preceding=5,
following=5,
aggregations={
"average": ("value", "mean"),
"minimum": ("value", "min"),
"maximum": ("value", "max"),
"first_label": ("label", "first"),
"last_label": ("label", "last"),
"valid_count": ("label", "count"),
},
)
rows = sorted(result.to_list(), key=lambda row: row["episode_id"])
self.assertIsInstance(result, pmm.TemporalAlignment)
self.assertEqual(pa.float64(), result.schema.field("average").type)
self.assertEqual(pa.string(), result.schema.field("first_label").type)
self.assertEqual(pa.int64(), result.schema.field("valid_count").type)
self.assertEqual(
(3.0, 1, 5, "a", "c", 2),
tuple(rows[0][name] for name in (
"average", "minimum", "maximum", "first_label",
"last_label", "valid_count")),
)
self.assertEqual(100.0, rows[1]["average"])
self.assertIsNone(rows[2]["average"])
self.assertEqual(0, rows[2]["valid_count"])
def test_window_join_reuses_arrow_slices(self):
source = temporal._WindowJoinRight.__new__(temporal._WindowJoinRight)
source.by = ("group",)
source._index = {(1,): (0, 3), (2,): (3, 4)}
source._time_keys = np.array([5, 10, 15, 10], dtype=np.int64)
source._row_ids = pa.array([11, 12, 13, 14], type=pa.int64())
source.time_type = pa.int64()
source._preceding_key = 5
source._following_key = 0
source.closed = "both"
specifications = (
("average", "value", "mean"),
("first_label", "label", "first"),
("minimum", "value", "min"),
("last_label", "label", "last"),
("maximum", "value", "max"),
("valid_count", "label", "count"),
)
values = pa.table({
"value": pa.array([1, None, 5, None], type=pa.int32()),
"label": pa.array(["a", None, "c", None], type=pa.string()),
})
anchors = [
{"group": group, temporal._TIME_KEY: time}
for group, time in [(1, 10), (1, 15), (2, 10), (3, 10)]
]
expected = {
"mean": pa.array([1.0, 5.0, None, None], type=pa.float64()),
"min": pa.array([1, 5, None, None], type=pa.int32()),
"max": pa.array([1, 5, None, None], type=pa.int32()),
"first": pa.array(["a", "c", None, None], type=pa.string()),
"last": pa.array(["a", "c", None, None], type=pa.string()),
"count": pa.array([1, 1, 0, 0], type=pa.int64()),
}
for aggregations in (specifications[:1], specifications[::2],
specifications):
with self.subTest(aggregations=aggregations):
source.aggregations = aggregations
fetcher = mock.Mock(schema=values.schema)
fetcher.fetch.return_value = values
with mock.patch.object(
temporal.pc, "take", wraps=temporal.pc.take) as take:
arrays = source.build_arrays(anchors, fetcher)
fetcher.fetch.assert_called_once_with([11, 12, 13, 14])
self.assertEqual(len(aggregations), len(arrays))
for (_, _, operation), array in zip(aggregations, arrays):
self.assertEqual(expected[operation].type, array.type)
self.assertEqual(expected[operation], array)
take.assert_not_called()
def test_window_join_preserves_order_across_disjoint_ranges(self):
anchors = self._table("range_anchors", {
"episode_id": pa.int32(), "event_time": pa.int64(),
})
samples = self._table("range_samples", {
"episode_id": pa.int32(), "event_time": pa.int64(),
"label": pa.string(),
})
# Physical row IDs are neither contiguous nor ordered within a window.
samples.add([
{"episode_id": 1, "event_time": 12, "label": "late"},
{"episode_id": 2, "event_time": 2, "label": "other"},
{"episode_id": 1, "event_time": 2, "label": "first"},
{"episode_id": 1, "event_time": 6, "label": "excluded"},
{"episode_id": 1, "event_time": 2, "label": None},
{"episode_id": 1, "event_time": 2, "label": "last"},
{"episode_id": 1, "event_time": 1, "label": "early"},
])
anchors.add([
{"episode_id": group, "event_time": time}
for group, time in [(1, 12), (1, 3), (2, 2), (1, 2),
(1, 2), (3, 2), (1, 20)]
])
query = pmm.join_window(
anchors.scan(), samples.scan().select("label"),
on="event_time", by="episode_id", preceding=1,
aggregations={
"first": ("label", "first"),
"last": ("label", "last"),
"count": ("label", "count"),
})
rows = query.to_list()
self.assertEqual([
("late", "late", 1), ("first", "last", 2),
("other", "other", 1), ("early", "last", 3),
("early", "last", 3), (None, None, 0), (None, None, 0),
], [(row["first"], row["last"], row["count"]) for row in rows])
def test_window_join_bounds_overlapping_window_memory(self):
size = 1024
source = temporal._WindowJoinRight.__new__(temporal._WindowJoinRight)
source.by = ("group",)
source._index = {(1,): (0, size)}
source._time_keys = np.arange(size, dtype=np.int64)
source._row_ids = pa.array(range(10000, 10000 + size), type=pa.int64())
source.time_type = pa.int64()
source._preceding_key = size - 1
source._following_key = 0
source.closed = "both"
source.aggregations = (("count", "value", "count"),)
values = pa.table({"value": pa.array(range(size), type=pa.int64())})
fetcher = mock.Mock(schema=values.schema)
fetcher.fetch.return_value = values
anchors = [{"group": 1, temporal._TIME_KEY: size - 1}] * 256
tracemalloc.start()
try:
arrays = source.build_arrays(anchors, fetcher)
_, peak = tracemalloc.get_traced_memory()
finally:
tracemalloc.stop()
self.assertEqual([size] * len(anchors), arrays[0].to_pylist())
fetcher.fetch.assert_called_once_with(list(range(10000, 10000 + size)))
# Per-window row-ID and position lists alone exceed this budget.
self.assertLess(peak, 4 * 1024 * 1024)
def test_window_endpoints_preserve_values_across_chunks(self):
cases = [
(pa.string(), [None, "", "last", None], "", "last"),
(pa.binary(), [None, b"", b"last", None], b"", b"last"),
(pa.bool_(), [None, False, True, None], False, True),
(pa.int64(), [None, 0, -1, None], 0, -1),
(pa.float64(), [None, float("nan"), 1.0, None], float("nan"), 1.0),
(pa.list_(pa.int64()), [None, [], [1, None], None], [], [1, None]),
(pa.struct([pa.field("a", pa.int64())]),
[None, {"a": None}, {"a": 0}, None], {"a": None}, {"a": 0}),
(pa.string(), [None, None], None, None),
(pa.string(), [], None, None),
]
for data_type, values, first, last in cases:
array = pa.array(values, type=data_type)
empty = pa.array([], type=data_type)
padded = pa.concat_arrays([
pa.nulls(1, type=data_type), array,
pa.nulls(1, type=data_type)])
chunks = pa.chunked_array([
empty, array.slice(0, 1), empty, array.slice(1), empty])
sliced_chunks = pa.chunked_array([
padded.slice(0, 2), empty, padded.slice(2), empty,
]).slice(1, len(array))
for selected in (array, padded.slice(1, len(array)),
chunks, sliced_chunks):
for operation, expected in (("first", first), ("last", last)):
with self.subTest(data_type=data_type, operation=operation,
selected=selected):
result = temporal._aggregate_values(selected, operation)
if isinstance(expected, float) and np.isnan(expected):
self.assertTrue(np.isnan(result))
else:
self.assertEqual(expected, result)
def test_window_endpoints_skip_null_dictionary_values(self):
array = pa.DictionaryArray.from_arrays(
pa.array([0, 1, 0], type=pa.int8()),
pa.array([None, "value"], type=pa.string()))
for selected in (array, pa.chunked_array([array.slice(0, 1), array.slice(1)])):
for operation in ("first", "last"):
with self.subTest(operation=operation, selected=selected):
self.assertEqual(
"value", temporal._aggregate_values(selected, operation))
def test_window_endpoints_do_not_materialize_full_payload(self):
payload = b"x" * 1024
selected = pa.chunked_array([
pa.array([payload] * 8192, type=pa.binary())])
for operation in ("first", "last"):
with self.subTest(operation=operation):
tracemalloc.start()
try:
result = temporal._aggregate_values(selected, operation)
_, peak = tracemalloc.get_traced_memory()
finally:
tracemalloc.stop()
self.assertEqual(payload, result)
# Converting the entire window creates over 8 MiB of Python bytes.
self.assertLess(peak, 256 * 1024)
def test_window_join_supports_asymmetric_timestamp_bounds(self):
anchors = self._table("window_timestamp_anchors", {
"episode_id": pa.int32(),
"event_time": pa.timestamp("ms"),
})
samples = self._table("window_timestamp_samples", {
"episode_id": pa.int32(),
"captured_at": pa.timestamp("ms"),
"value": pa.float32(),
})
anchor = datetime(2026, 9, 1, 12, 0, 0)
anchors.add([{"episode_id": 1, "event_time": anchor}])
samples.add([
{"episode_id": 1,
"captured_at": anchor + timedelta(milliseconds=offset),
"value": value}
for offset, value in ((-11, 100.0), (-10, 1.0),
(0, 2.0), (5, 3.0), (6, 100.0))
])
row = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", right_on="captured_at", by="episode_id",
preceding=timedelta(milliseconds=10),
following=timedelta(milliseconds=5),
aggregations={"value": "mean"},
).to_list()[0]
self.assertEqual(2.0, row["value"])
right_closed = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", right_on="captured_at", by="episode_id",
preceding=timedelta(milliseconds=10),
following=timedelta(milliseconds=5),
aggregations={"value": "mean"},
closed="right",
).to_list()[0]
self.assertEqual(2.5, right_closed["value"])
def test_window_join_preserves_subunit_timestamp_bounds(self):
anchors = self._table("window_subunit_timestamp_anchors", {
"episode_id": pa.int32(),
"event_time": pa.timestamp("ms"),
})
samples = self._table("window_subunit_timestamp_samples", {
"episode_id": pa.int32(),
"event_time": pa.timestamp("ms"),
"value": pa.int32(),
})
anchor = datetime(2026, 9, 1, 12, 0, 0)
anchors.add([{"episode_id": 1, "event_time": anchor}])
samples.add([
{"episode_id": 1,
"event_time": anchor + timedelta(milliseconds=offset),
"value": offset}
for offset in (-1, 0, 1)
])
exact = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
preceding=timedelta(microseconds=500),
following=timedelta(0), closed="right",
aggregations={"matches": ("value", "count")},
).to_list()[0]
open_window = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
preceding=timedelta(microseconds=1_500),
following=timedelta(microseconds=1_500), closed="neither",
aggregations={"matches": ("value", "count")},
).to_list()[0]
self.assertEqual(1, exact["matches"])
self.assertEqual(3, open_window["matches"])
def test_window_join_normalizes_numpy_float_bounds(self):
anchors = self._table("window_numpy_float_anchors", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
})
samples = self._table("window_numpy_float_samples", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
"value": pa.int32(),
})
timestamp = 1_700_000_000.001
anchors.add([{"episode_id": 1, "event_time": timestamp}])
samples.add([{
"episode_id": 1, "event_time": timestamp, "value": 7,
}])
row = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
preceding=0, following=np.float32(0),
aggregations={"matches": ("value", "count")},
).to_list()[0]
self.assertEqual(1, row["matches"])
def test_window_join_keeps_numeric_bounds_exact_for_integer_time(self):
anchors = self._table("window_integer_bound_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_integer_bound_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
timestamp = 1_700_000_000_000_000_001
anchors.add(pa.Table.from_pydict({
"episode_id": [1], "event_time": [timestamp],
}))
samples.add(pa.Table.from_pydict({
"episode_id": [1, 1, 1],
"event_time": [timestamp - 1, timestamp, timestamp + 1],
"value": [1, 2, 3],
}))
for preceding, following in (
(0.0, 0.0), (0.125, 0.125),
(np.int64(0), 0), (0, np.int64(0)),
(np.int64(0), np.int64(0))):
with self.subTest(preceding=preceding, following=following):
row = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
preceding=preceding, following=following,
aggregations={
"matches": ("value", "count"),
"first_value": ("value", "first"),
},
).to_list()[0]
self.assertEqual(1, row["matches"])
self.assertEqual(2, row["first_value"])
def test_window_join_clips_bounds_to_integer_time_range(self):
anchors = self._table("window_range_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_range_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
minimum, maximum = -(1 << 63), (1 << 63) - 1
anchors.add(pa.Table.from_pydict({
"episode_id": [1, 2], "event_time": [minimum, maximum],
}))
times = {1: [minimum, minimum + 1], 2: [maximum - 1, maximum]}
samples.add(pa.Table.from_pydict({
"episode_id": [1, 1, 2, 2],
"event_time": times[1] + times[2],
"value": [1, 2, 3, 4],
}))
for preceding, following in ((0, 0), (0, 1), (1, 0), (1, 1)):
for closed in ("both", "left", "right", "neither"):
with self.subTest(
preceding=preceding, following=following,
closed=closed):
rows = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
preceding=preceding, following=following,
closed=closed,
aggregations={"matches": ("value", "count")},
).to_list()
for row in rows:
left = row["event_time"] - preceding
right = row["event_time"] + following
expected = sum(
(time >= left if closed in ("both", "left")
else time > left)
and (time <= right if closed in ("both", "right")
else time < right)
for time in times[row["episode_id"]]
)
self.assertEqual(expected, row["matches"])
@pytest.mark.python_read
def test_window_join_prunes_unaggregated_right_columns(self):
anchors = self._table("window_projection_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_projection_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"payload": pa.struct([
pa.field("unused_nested", pa.int32()),
pa.field("value", pa.int32()),
]),
"unused": pa.string(),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
samples.add([{
"episode_id": 1,
"event_time": 10,
"payload": {"value": 7, "unused_nested": 8},
"unused": "not read",
}])
payload_reads = []
original = FormatPyArrowReader._read_parquet_row_group_batches
def tracked(reader, row_group, columns):
if columns is not None and "payload" in columns:
payload_reads.append(tuple(columns))
yield from original(reader, row_group, columns)
with mock.patch.object(
FormatPyArrowReader,
"_read_parquet_row_group_batches", tracked):
row = pmm.join_window(
anchors.scan(),
samples.scan().select(["payload.value", "unused"]),
on="event_time", by="episode_id", preceding=0,
aggregations={"payload_value": "mean"},
).to_list()[0]
self.assertEqual(7.0, row["payload_value"])
self.assertTrue(payload_reads)
self.assertNotIn("unused", {
name for columns in payload_reads for name in columns
})
def test_window_join_preserves_nested_projection_aliases_when_pruned(self):
anchors = self._table("window_alias_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_alias_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"a_b": pa.int32(),
"a": pa.struct([pa.field("b", pa.int32())]),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
samples.add([{
"episode_id": 1,
"event_time": 10,
"a_b": 3,
"a": {"b": 7},
}])
row = pmm.join_window(
anchors.scan(), samples.scan().select(["a_b", "a.b"]),
on="event_time", by="episode_id", preceding=0,
aggregations={"nested_mean": ("a_b__0", "mean")},
).to_list()[0]
self.assertEqual(7.0, row["nested_mean"])
def test_window_join_preserves_masked_alias_types_when_pruned(self):
anchors = self._table("window_masked_alias_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_masked_alias_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"a_b": pa.int32(),
"a": pa.struct([pa.field("b", pa.int32())]),
})
anchors.add(pa.Table.from_pydict({
"episode_id": [1, 1], "event_time": [10, 20],
}))
samples.add(pa.Table.from_pydict({
"episode_id": [1], "event_time": [10],
"a_b": [3],
"a": pa.array([{"b": 7}],
type=pa.struct([pa.field("b", pa.int32())])),
}))
auth = TableQueryAuthResult(
filter=None,
column_masking={"a_b": json.dumps({
"name": "CAST",
"fieldRef": {"index": 2, "name": "a_b", "type": "INT"},
"type": "STRING",
})},
)
samples.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
for projection, source in (
(["a_b"], "a_b"), (["a.b", "a_b"], "a_b__0")):
with self.subTest(projection=projection):
result = pmm.join_window(
anchors.scan(), samples.scan().select(projection),
on="event_time", by="episode_id", preceding=0,
aggregations={
"first_value": (source, "first"),
"last_value": (source, "last"),
},
).to_arrow()
for name in ("first_value", "last_value"):
self.assertEqual(pa.string(), result[name].type)
self.assertEqual(["3", None], result[name].to_pylist())
for operation in ("mean", "min", "max"):
with self.subTest(projection=projection, operation=operation):
with self.assertRaisesRegex(
TypeError, "requires an integer or floating"):
pmm.join_window(
anchors.scan(), samples.scan().select(projection),
on="event_time", by="episode_id", preceding=0,
aggregations={"value": (source, operation)},
).to_arrow()
def test_window_join_matches_masks_by_original_nested_path(self):
anchors = self._table("window_nested_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_nested_mask_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"a_b": pa.int32(),
"a": pa.struct([pa.field("b", pa.int32())]),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
samples.add([{
"episode_id": 1, "event_time": 10,
"a_b": 3, "a": {"b": 7},
}])
auth = TableQueryAuthResult(
filter=None,
column_masking={"a_b": json.dumps({"name": "NULL"})},
)
samples.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
row = pmm.join_window(
anchors.scan(), samples.scan().select(["a_b", "a.b"]),
on="event_time", by="episode_id", preceding=0,
aggregations={"nested_mean": ("a_b__0", "mean")},
).to_list()[0]
self.assertEqual(7.0, row["nested_mean"])
def test_window_mean_uses_masked_numeric_alias_type_when_pruned(self):
anchors = self._table("window_numeric_alias_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_numeric_alias_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"a_b": pa.string(),
"a": pa.struct([pa.field("b", pa.int32())]),
})
anchors.add(pa.Table.from_pydict({
"episode_id": [1, 1], "event_time": [10, 20],
}))
samples.add(pa.Table.from_pydict({
"episode_id": [1, 1], "event_time": [9, 11],
"a_b": ["6", "10"],
"a": pa.array([{"b": 1}, {"b": 2}],
type=pa.struct([pa.field("b", pa.int32())])),
}))
auth = TableQueryAuthResult(
filter=None,
column_masking={"a_b": json.dumps({
"name": "CAST",
"fieldRef": {"index": 2, "name": "a_b", "type": "STRING"},
"type": "DOUBLE",
})},
)
samples.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
for projection, source in (
(["a_b"], "a_b"), (["a.b", "a_b"], "a_b__0")):
with self.subTest(projection=projection):
result = pmm.join_window(
anchors.scan(), samples.scan().select(projection),
on="event_time", by="episode_id", preceding=1, following=1,
aggregations={"value": (source, "mean")},
).to_arrow()
self.assertEqual(pa.float64(), result["value"].type)
self.assertEqual([8.0, None], result["value"].to_pylist())
def test_window_mean_avoids_numeric_overflow_and_integer_rounding(self):
anchors = self._table("window_numeric_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_numeric_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"integer_value": pa.int64(),
"float_value": pa.float64(),
})
anchors.add([{"episode_id": 1, "event_time": 1}])
samples.add([
{"episode_id": 1, "event_time": 0,
"integer_value": -(1 << 63),
"float_value": sys.float_info.max},
{"episode_id": 1, "event_time": 2,
"integer_value": (1 << 63) - 1,
"float_value": sys.float_info.max},
])
row = pmm.join_window(
anchors.scan(),
samples.scan().select(["integer_value", "float_value"]),
on="event_time", by="episode_id", preceding=1, following=1,
aggregations={
"integer_value": "mean",
"float_value": "mean",
},
).to_list()[0]
self.assertEqual(-0.5, row["integer_value"])
self.assertEqual(sys.float_info.max, row["float_value"])
def test_window_mean_preserves_finite_float_cancellation(self):
anchors = self._table("window_float_mean_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("window_float_mean_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.float64(),
})
anchors.add([{"episode_id": 1, "event_time": 1}])
samples.add([
{"episode_id": 1, "event_time": 0, "value": 1e16},
{"episode_id": 1, "event_time": 2, "value": -1e16 + 2},
])
row = pmm.join_window(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id", preceding=1, following=1,
aggregations={"value": "mean"},
).to_list()[0]
self.assertEqual(1.0, row["value"])
def test_window_join_can_follow_an_asof_join(self):
anchors = self._table("window_chain_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
images = self._table("window_chain_images", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"image": pa.string(),
})
imu = self._table("window_chain_imu", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"acceleration": pa.float32(),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
images.add([{"episode_id": 1, "event_time": 9, "image": "frame"}])
imu.add([
{"episode_id": 1, "event_time": 8, "acceleration": 1.0},
{"episode_id": 1, "event_time": 10, "acceleration": 3.0},
])
row = pmm.join_asof(
anchors.scan(), images.scan().select("image"),
on="event_time", by="episode_id", direction="nearest",
).join_window(
imu.scan().select("acceleration"),
preceding=2,
aggregations={"acceleration": "mean"},
).to_list()[0]
self.assertEqual("frame", row["image"])
self.assertEqual(2.0, row["acceleration"])
def test_window_join_validates_options(self):
table = self._table("window_validation", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
"text": pa.string(),
})
def scan():
return table.scan().select("value")
with self.assertRaisesRegex(ValueError, "non-negative"):
pmm.join_window(
scan(), scan(), on="event_time", by="episode_id",
preceding=-1, aggregations={"value": "mean"})
with self.assertRaisesRegex(ValueError, "Unsupported aggregation"):
pmm.join_window(
scan(), scan(), on="event_time", by="episode_id",
preceding=1, aggregations={"value": "median"})
with self.assertRaisesRegex(ValueError, "closed must be"):
pmm.join_window(
scan(), scan(), on="event_time", by="episode_id",
preceding=1, aggregations={"value": "mean"},
closed="middle")
with self.assertRaisesRegex(ValueError, "missing aggregation columns"):
pmm.join_window(
scan(), scan(), on="event_time", by="episode_id",
preceding=1, aggregations={"missing": "mean"})
with self.assertRaisesRegex(
TypeError, "requires an integer or floating"):
pmm.join_window(
scan(), table.scan().select("text"),
on="event_time", by="episode_id", preceding=1,
aggregations={"text": "mean"}).to_arrow()
def test_linear_interpolation_preserves_an_exact_infinite_float(self):
anchors = self._table("linear_exact_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_exact_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.float64(),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
states.add([{
"episode_id": 1, "event_time": 10, "value": float("inf")
}])
row = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id",
).to_list()[0]
self.assertEqual(float("inf"), row["value"])
def test_linear_interpolation_requires_both_neighbors_in_tolerance(self):
anchors = self._table("linear_tolerance_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_tolerance_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.float64(),
})
anchors.add([{"episode_id": 1, "event_time": 9}])
states.add([
{"episode_id": 1, "event_time": 0, "value": 0.0},
{"episode_id": 1, "event_time": 10, "value": 10.0},
])
row = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id", tolerance=8,
).to_list()[0]
self.assertIsNone(row["value"])
def test_linear_interpolation_supports_fixed_size_numeric_lists(self):
vector = pa.list_(pa.float32(), 2)
anchors = self._table("linear_vector_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_vector_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"state": vector,
})
anchors.add([{"episode_id": 1, "event_time": 5}])
states.add([
{"episode_id": 1, "event_time": 0, "state": [0.0, 10.0]},
{"episode_id": 1, "event_time": 10, "state": [10.0, 20.0]},
])
result = pmm.interpolate(
anchors.scan(), states.scan().select("state"),
on="event_time", by="episode_id",
)
self.assertEqual(vector, result.schema.field("state").type)
self.assertEqual([5.0, 15.0], result.to_list()[0]["state"])
def test_linear_interpolation_preserves_bigint_precision(self):
anchors = self._table("linear_bigint_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_bigint_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int64(),
"extreme": pa.int64(),
})
anchors.add([
{"episode_id": 1, "event_time": 1},
{"episode_id": 1, "event_time": 2},
])
states.add([
{"episode_id": 1, "event_time": 0,
"value": (1 << 53) + 1, "extreme": -(1 << 63)},
{"episode_id": 1, "event_time": 4,
"value": (1 << 53) + 3, "extreme": (1 << 63) - 1},
])
rows = pmm.interpolate(
anchors.scan(), states.scan().select(["value", "extreme"]),
on="event_time", by="episode_id",
).to_list()
by_time = {row["event_time"]: row for row in rows}
self.assertEqual(float((1 << 53) + 2), by_time[1]["value"])
self.assertEqual(-0.5, by_time[2]["extreme"])
def test_linear_interpolation_scales_extreme_float_time_axis(self):
anchors = self._table("linear_extreme_time_anchors", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
})
states = self._table("linear_extreme_time_states", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
"value": pa.float64(),
})
anchors.add([{"episode_id": 1, "event_time": 0.0}])
states.add([
{"episode_id": 1, "event_time": -sys.float_info.max,
"value": 0.0},
{"episode_id": 1, "event_time": sys.float_info.max,
"value": 10.0},
])
row = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id",
).to_list()[0]
self.assertEqual(5.0, row["value"])
def test_linear_interpolation_preserves_float_time_differences(self):
anchors = self._table("linear_nearby_time_anchors", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
})
states = self._table("linear_nearby_time_states", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
"value": pa.float64(),
})
base = 1_700_000_000.0
step = 2 ** -21
anchors.add([{"episode_id": 1, "event_time": base + step}])
states.add([
{"episode_id": 1, "event_time": base, "value": 0.0},
{"episode_id": 1, "event_time": base + 2 * step,
"value": 10.0},
])
row = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id",
).to_list()[0]
self.assertEqual(5.0, row["value"])
def test_float_time_preserves_bigint_interpolation_precision(self):
anchors = self._table("linear_float_bigint_anchors", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
})
states = self._table("linear_float_bigint_states", {
"episode_id": pa.int32(),
"event_time": pa.float64(),
"value": pa.int64(),
})
anchors.add([{"episode_id": 1, "event_time": 2.0}])
states.add([
{"episode_id": 1, "event_time": 0.0,
"value": -(1 << 63)},
{"episode_id": 1, "event_time": 3.0,
"value": 1 << 62},
])
row = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id",
).to_list()[0]
self.assertEqual(0.0, row["value"])
def test_linear_interpolation_handles_extreme_float_payloads(self):
anchors = self._table("linear_extreme_value_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_extreme_value_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.float64(),
"infinite": pa.float64(),
})
anchors.add([{"episode_id": 1, "event_time": 5}])
states.add([
{"episode_id": 1, "event_time": 0,
"value": -sys.float_info.max, "infinite": float("inf")},
{"episode_id": 1, "event_time": 10,
"value": sys.float_info.max, "infinite": float("inf")},
])
row = pmm.interpolate(
anchors.scan(), states.scan().select(["value", "infinite"]),
on="event_time", by="episode_id",
).to_list()[0]
self.assertEqual(0.0, row["value"])
self.assertEqual(float("inf"), row["infinite"])
def test_linear_interpolation_uses_effective_masked_payload_type(self):
anchors = self._table("linear_masked_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_masked_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.string(),
})
anchors.add([{"episode_id": 1, "event_time": 5}])
states.add([
{"episode_id": 1, "event_time": 0, "value": "0"},
{"episode_id": 1, "event_time": 10, "value": "10"},
])
auth = TableQueryAuthResult(
filter=None,
column_masking={"value": json.dumps({
"name": "CAST",
"fieldRef": {
"index": 2, "name": "value", "type": "STRING",
},
"type": "DOUBLE",
})},
)
states.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
row = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id",
).to_list()[0]
self.assertEqual(5.0, row["value"])
def test_linear_interpolation_rejects_masked_non_numeric_payload(self):
anchors = self._table("linear_invalid_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
states = self._table("linear_invalid_mask_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 5}])
states.add([
{"episode_id": 1, "event_time": 0, "value": 0},
{"episode_id": 1, "event_time": 10, "value": 10},
])
auth = TableQueryAuthResult(
filter=None,
column_masking={"value": json.dumps({
"name": "CAST",
"fieldRef": {
"index": 2, "name": "value", "type": "INT",
},
"type": "STRING",
})},
)
states.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
aligned = pmm.interpolate(
anchors.scan(), states.scan().select("value"),
on="event_time", by="episode_id",
)
with self.assertRaisesRegex(TypeError, "requires integer or floating"):
aligned.to_arrow()
def test_linear_interpolation_rejects_decimal_payloads(self):
anchors = self._table("linear_invalid_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
decimals = self._table("linear_invalid_decimals", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.decimal128(10, 2),
})
with self.assertRaisesRegex(TypeError, "requires integer or floating"):
pmm.interpolate(
anchors.scan(), decimals.scan().select("value"),
on="event_time", by="episode_id",
).to_arrow()
def test_linear_interpolation_can_follow_an_asof_join(self):
anchors = self._table("linear_chain_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
images = self._table("linear_chain_images", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"image": pa.string(),
})
states = self._table("linear_chain_states", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"state": pa.float32(),
})
anchors.add([{"episode_id": 1, "event_time": 5}])
images.add([{"episode_id": 1, "event_time": 4, "image": "frame"}])
states.add([
{"episode_id": 1, "event_time": 0, "state": 0.0},
{"episode_id": 1, "event_time": 10, "state": 10.0},
])
row = pmm.join_asof(
anchors.scan(), images.scan().select("image"),
on="event_time", by="episode_id",
direction="nearest", tolerance=2,
).interpolate(
states.scan().select("state"), tolerance=5,
).to_list()[0]
self.assertEqual("frame", row["image"])
self.assertEqual(5.0, row["state"])
def test_alignment_can_return_matched_timestamp(self):
anchors = self._table("timestamp_output_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("timestamp_output_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
samples.add([
{"episode_id": 1, "event_time": 95, "value": 7},
])
row = pmm.join_asof(
anchors.scan(),
samples.scan().select(["event_time", "value"]),
on="event_time", by="episode_id",
direction="backward", suffix="_matched",
).to_list()[0]
self.assertEqual(95, row["event_time_matched"])
self.assertEqual(5, row["event_time"] - row["event_time_matched"])
def test_alignment_rejects_nested_group_keys(self):
group_type = pa.struct([pa.field("part", pa.int32())])
anchors = self._table("nested_group_anchors", {
"group": group_type,
"event_time": pa.int64(),
})
samples = self._table("nested_group_samples", {
"group": group_type,
"event_time": pa.int64(),
})
with self.assertRaisesRegex(TypeError, "must have a scalar type"):
pmm.join_asof(
anchors.scan(),
samples.scan(),
on="event_time",
by="group",
direction="nearest",
tolerance=0,
)
def test_aligns_named_sources_in_episode_local_batches(self):
actions = self._table("actions", {
"episode_id": pa.string(),
"event_time": pa.int64(),
"action": pa.int32(),
})
images = self._table("images", {
"episode_id": pa.string(),
"event_time": pa.int64(),
"camera": pa.string(),
"image": pa.string(),
})
states = self._table("states", {
"episode_id": pa.string(),
"event_time": pa.int64(),
"state": pa.int32(),
})
commands = self._table("commands", {
"episode_id": pa.string(),
"event_time": pa.int64(),
"command": pa.string(),
})
actions.add([
{"episode_id": "ep-2", "event_time": 100, "action": 4},
{"episode_id": "ep-1", "event_time": 300, "action": 3},
{"episode_id": "ep-1", "event_time": 100, "action": 1},
{"episode_id": "ep-1", "event_time": 200, "action": 2},
])
images.add([
{"episode_id": "ep-1", "event_time": 90,
"camera": "left", "image": "early"},
{"episode_id": "ep-1", "event_time": 90,
"camera": "right", "image": "ignored"},
{"episode_id": "ep-1", "event_time": 110,
"camera": "left", "image": "late"},
{"episode_id": "ep-1", "event_time": 215,
"camera": "left", "image": "middle"},
{"episode_id": "ep-2", "event_time": 99,
"camera": "left", "image": "other"},
])
states.add([
{"episode_id": "ep-1", "event_time": 80, "state": 8},
{"episode_id": "ep-1", "event_time": 190, "state": 19},
{"episode_id": "ep-2", "event_time": 95, "state": 95},
])
commands.add([
{"episode_id": "ep-1", "event_time": 100, "command": "open"},
{"episode_id": "ep-1", "event_time": 220, "command": "close"},
{"episode_id": "ep-2", "event_time": 100, "command": "hold"},
])
aligned = pmm.join_asof(
actions.scan().select(["episode_id", "event_time", "action"]),
images.scan().where("camera = 'left'").select("image"),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=20,
).join_asof(
states.scan().select("state"),
direction="backward",
tolerance=25,
).join_asof(
commands.scan().select("command"),
direction="nearest",
tolerance=0,
).join_asof(
commands.scan().select("command"),
direction="forward",
tolerance=25,
suffix="_next",
)
reader = aligned.to_arrow_batch_reader(batch_size=2)
batches = list(reader)
rows = pa.Table.from_batches(batches).to_pylist()
self.assertEqual([2, 2], [batch.num_rows for batch in batches])
rows.sort(key=lambda row: (row["episode_id"], row["event_time"]))
self.assertEqual(
[("ep-1", 100), ("ep-1", 200), ("ep-1", 300), ("ep-2", 100)],
[(row["episode_id"], row["event_time"]) for row in rows],
)
# Equal-distance nearest ties select the earlier row.
self.assertEqual(
["early", "middle", None, "other"],
[row["image"] for row in rows],
)
self.assertEqual([8, 19, None, 95], [
row["state"] for row in rows
])
self.assertEqual(["open", None, None, "hold"], [
row["command"] for row in rows
])
self.assertEqual(["open", "close", None, "hold"], [
row["command_next"] for row in rows
])
def test_alignment_pins_each_scan_snapshot(self):
anchors = self._table("pinned_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.string(),
})
secondary = self._table("pinned_secondary", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.string(),
})
anchors.add([{"episode_id": 1, "event_time": 100, "value": "old"}])
secondary.add([
{"episode_id": 1, "event_time": 90, "value": "old-match"}
])
aligned = pmm.join_asof(
anchors.scan(),
secondary.scan(),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=20,
suffix="_secondary",
)
snapshots = aligned.resolved_snapshots
anchors.add([{"episode_id": 1, "event_time": 200, "value": "new"}])
secondary.add([
{"episode_id": 1, "event_time": 100, "value": "new-match"}
])
self.assertEqual([{
"episode_id": 1,
"event_time": 100,
"value": "old",
"value_secondary": "old-match",
}], aligned.to_list())
self.assertEqual(1, snapshots["left"]["snapshot_id"])
self.assertEqual(1, snapshots["right_1"]["snapshot_id"])
self.assertTrue(
snapshots["left"]["table"].endswith("pinned_anchors"))
self.assertTrue(
snapshots["right_1"]["table"].endswith("pinned_secondary"))
def test_alignment_reads_tag_after_snapshot_file_expires(self):
anchors = self._table("tagged_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.string(),
})
secondary = self._table("tagged_secondary", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.string(),
})
anchors.add([{"episode_id": 1, "event_time": 100, "value": "old"}])
secondary.add([
{"episode_id": 1, "event_time": 100, "value": "old-match"}
])
anchors.raw_table.create_tag("v1")
secondary.raw_table.create_tag("v1")
aligned = pmm.join_asof(
anchors.scan(tag_name="v1"),
secondary.scan(tag_name="v1"),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=0,
)
anchors.add([{"episode_id": 1, "event_time": 200, "value": "new"}])
secondary.add([
{"episode_id": 1, "event_time": 200, "value": "new-match"}
])
for table in (anchors.raw_table, secondary.raw_table):
manager = table.snapshot_manager()
table.file_io.delete(manager.get_snapshot_path(1))
self.assertEqual([100], [
row["event_time"] for row in aligned.to_list()
])
self.assertEqual(
"v1", aligned.resolved_snapshots["left"]["tag_name"])
def test_alignment_rejects_a_changed_tag(self):
table = self._table("changed_tag", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
table.add([{"episode_id": 1, "event_time": 100}])
table.raw_table.create_tag("v1")
aligned = pmm.join_asof(
table.scan(tag_name="v1"), table.scan(tag_name="v1"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
table.add([{"episode_id": 1, "event_time": 200}])
table.raw_table.replace_tag("v1")
with self.assertRaisesRegex(RuntimeError, "Tag 'v1' changed"):
aligned.to_list()
def test_alignment_rejects_tag_replacement_during_planning(self):
anchors = self._table("raced_tag_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("raced_tag_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{
"episode_id": 1, "event_time": 100, "value": 7,
}])
source.raw_table.create_tag("v1")
source.add([{
"episode_id": 1, "event_time": 200, "value": 99,
}])
aligned = pmm.join_asof(
anchors.scan(), source.scan(tag_name="v1").select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
original = temporal._validate_pinned_tag
replaced = []
def replace_after_validation(query):
original(query)
if query._temporal_tag_name == "v1" and not replaced:
source.raw_table.replace_tag("v1", snapshot_id=2)
replaced.append(True)
with mock.patch.object(
temporal, "_validate_pinned_tag",
side_effect=replace_after_validation):
with self.assertRaisesRegex(
RuntimeError, "changed from snapshot 1 to 2"):
aligned.to_list()
def test_alignment_normalizes_scan_mode_when_pinning(self):
anchors = self.conn.create_table(
"latest_full_anchors",
schema=pa.schema([
pa.field("episode_id", pa.int32()),
pa.field("event_time", pa.int64()),
]),
options={"scan.mode": "latest-full"},
)
secondary = self._table("latest_full_secondary", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
secondary.add([
{"episode_id": 1, "event_time": 100, "value": 7}
])
row = pmm.join_asof(
anchors.scan(),
secondary.scan(),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=0,
).to_list()[0]
self.assertEqual(7, row["value"])
def test_alignment_rejects_non_finite_temporal_values(self):
for value in (float("nan"), float("inf"), float("-inf")):
with self.subTest(value=value):
anchors = self._table("float_anchor_%s" % id(value), {
"episode_id": pa.int32(),
"event_time": pa.float64(),
})
secondary = self._table("float_source_%s" % id(value), {
"episode_id": pa.int32(),
"event_time": pa.float64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100.0}])
secondary.add([
{"episode_id": 1, "event_time": value, "value": 7}
])
aligned = pmm.join_asof(
anchors.scan(),
secondary.scan(),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=1.0,
)
with self.assertRaisesRegex(ValueError, "must be finite"):
aligned.to_list()
def test_alignment_validates_tolerance_type_and_value(self):
table = self._table("tolerance", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
for tolerance in (float("nan"), float("inf"), -1):
with self.subTest(tolerance=tolerance):
with self.assertRaises((TypeError, ValueError)):
pmm.join_asof(
table.scan(), table.scan(),
on="event_time", by="episode_id",
direction="nearest", tolerance=tolerance,
)
with self.assertRaisesRegex(TypeError, "Numeric alignment"):
pmm.join_asof(
table.scan(), table.scan(),
on="event_time", by="episode_id",
direction="nearest",
tolerance=timedelta(milliseconds=1),
)
with self.assertRaisesRegex(ValueError, "int64 maximum"):
pmm.join_asof(
table.scan(), table.scan(),
on="event_time", by="episode_id",
direction="nearest", tolerance=1 << 63,
)
with self.assertRaisesRegex(ValueError, "direction"):
pmm.join_asof(
table.scan(), table.scan(),
on="event_time", by="episode_id", direction="exact",
)
def test_alignment_keeps_internal_row_id_out_of_query_auth(self):
anchors = self._table("masked_row_id_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("masked_row_id_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{"episode_id": 1, "event_time": 100, "value": 7}])
selected = []
def query_auth(select):
selected.append(select)
self.assertTrue(select is None or "_ROW_ID" not in select)
return None
for table in (anchors.raw_table, source.raw_table):
table.catalog_environment.table_query_auth = (
lambda options, identifier: query_auth)
rows = pmm.join_asof(
anchors.scan(), source.scan(),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
self.assertEqual(7, rows[0]["value"])
self.assertTrue(selected)
def test_alignment_rejects_masked_internal_row_ids(self):
anchors = self._table("masked_internal_id_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("masked_internal_id_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{"episode_id": 1, "event_time": 100, "value": 7}])
auth = TableQueryAuthResult(
filter=None,
column_masking={"_ROW_ID": json.dumps({"name": "NULL"})},
)
anchors.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
aligned = pmm.join_asof(
anchors.scan(), source.scan(),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
with self.assertRaisesRegex(ValueError, "masks _ROW_ID"):
aligned.to_list()
def test_alignment_preserves_masked_output_schema(self):
anchors = self.conn.create_table(
"masked_schema_anchors",
schema=pa.schema([
pa.field("episode_id", pa.int32()),
pa.field("event_time", pa.int64()),
pa.field("secret", pa.int32(), nullable=False),
]),
)
source = self._table("masked_schema_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"sample": pa.int32(),
})
anchors.add([{
"episode_id": 1, "event_time": 100, "secret": 7,
}])
source.add([{
"episode_id": 1, "event_time": 100, "sample": 1,
}])
auth = TableQueryAuthResult(
filter=None,
column_masking={"secret": json.dumps({"name": "NULL"})},
)
anchors.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
reader = pmm.join_asof(
anchors.scan(), source.scan().select("sample"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_arrow_batch_reader()
result = reader.read_all()
self.assertTrue(result.schema.field("secret").nullable)
self.assertIsNone(result["secret"][0].as_py())
def test_alignment_preserves_type_changing_masked_output(self):
anchors = self._table("cast_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"anchor_value": pa.int32(),
})
source = self._table("cast_mask_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"source_value": pa.int32(),
})
anchors.add([{
"episode_id": 1, "event_time": 100, "anchor_value": 7,
}])
source.add([{
"episode_id": 1, "event_time": 100, "source_value": 8,
}])
for table, name in (
(anchors.raw_table, "anchor_value"),
(source.raw_table, "source_value")):
auth = TableQueryAuthResult(
filter=None,
column_masking={name: json.dumps({
"name": "CAST",
"fieldRef": {"index": 2, "name": name, "type": "INT"},
"type": "STRING",
})},
)
table.catalog_environment.table_query_auth = (
lambda options, identifier, result=auth:
lambda select: result)
result = pmm.join_asof(
anchors.scan(), source.scan().select("source_value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_arrow()
self.assertEqual(pa.string(), result.schema.field("anchor_value").type)
self.assertEqual(pa.string(), result.schema.field("source_value").type)
self.assertEqual("7", result["anchor_value"][0].as_py())
self.assertEqual("8", result["source_value"][0].as_py())
def test_alignment_preserves_masked_schema_for_empty_source(self):
anchors = self._table("empty_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("empty_mask_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
auth = TableQueryAuthResult(
filter=None,
column_masking={"value": json.dumps({
"name": "CAST",
"fieldRef": {
"index": 2, "name": "value", "type": "INT",
},
"type": "STRING",
})},
)
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
result = pmm.join_asof(
anchors.scan(), source.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_arrow()
self.assertEqual(pa.string(), result.schema.field("value").type)
self.assertIsNone(result["value"][0].as_py())
def test_alignment_rejects_type_changing_key_masks(self):
anchors = self._table("cast_key_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("cast_key_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{
"episode_id": 1, "event_time": 100, "value": 7,
}])
for name, index, field_type in (
("episode_id", 0, "INT"),
("event_time", 1, "BIGINT")):
with self.subTest(name=name):
auth = TableQueryAuthResult(
filter=None,
column_masking={name: json.dumps({
"name": "CAST",
"fieldRef": {
"index": index, "name": name, "type": field_type,
},
"type": "STRING",
})},
)
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier, result=auth:
lambda select: result)
with self.assertRaisesRegex(TypeError, name):
pmm.join_asof(
anchors.scan(), source.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
def test_alignment_supports_cross_column_key_masks(self):
anchors = self._table("cross_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("cross_mask_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"allowed_episode_id": pa.int32(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{
"episode_id": 99,
"event_time": 100,
"allowed_episode_id": 1,
"value": 7,
}])
auth = TableQueryAuthResult(
filter=None,
column_masking={"episode_id": json.dumps({
"name": "FIELD_REF",
"fieldRef": {
"index": 2,
"name": "allowed_episode_id",
"type": "INT",
},
})},
)
selected = []
def query_auth(select):
selected.append(select)
return auth
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: query_auth)
rows = pmm.join_asof(
anchors.scan(), source.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
self.assertEqual(7, rows[0]["value"])
self.assertTrue(selected)
self.assertTrue(all(
select is None or "allowed_episode_id" not in select
for select in selected
))
def test_alignment_rejects_a_mask_reading_a_masked_dependency(self):
anchors = self._table("dependency_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"allowed_episode_id": pa.int32(),
})
source = self._table("dependency_mask_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{
"episode_id": 99,
"event_time": 100,
"allowed_episode_id": 1,
}])
source.add([{
"episode_id": 1, "event_time": 100, "value": 7,
}])
auth = TableQueryAuthResult(
filter=None,
column_masking={
"episode_id": json.dumps({
"name": "FIELD_REF",
"fieldRef": {
"index": 2,
"name": "allowed_episode_id",
"type": "INT",
},
}),
"allowed_episode_id": json.dumps({
"name": "CAST",
"fieldRef": {
"index": 2,
"name": "allowed_episode_id",
"type": "INT",
},
"type": "STRING",
}),
},
)
anchors.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
with self.assertRaisesRegex(ValueError, "which is masked too"):
pmm.join_asof(
anchors.scan(), source.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_arrow()
def test_alignment_rejects_cross_column_and_json_null_masks(self):
anchors = self._table("mask_semantics_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("mask_semantics_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"first": pa.string(),
"second": pa.string(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{
"episode_id": 1,
"event_time": 100,
"first": "a",
"second": "b",
}])
auth = [TableQueryAuthResult(
filter=None,
column_masking={
"first": json.dumps({
"name": "FIELD_REF",
"fieldRef": {
"index": 3, "name": "second", "type": "STRING",
},
}),
"second": json.dumps({
"name": "FIELD_REF",
"fieldRef": {
"index": 2, "name": "first", "type": "STRING",
},
}),
},
)]
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth[0])
with self.assertRaisesRegex(ValueError, "which is masked too"):
pmm.join_asof(
anchors.scan(), source.scan().select(["first", "second"]),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
auth[0] = TableQueryAuthResult(
filter=None, column_masking={"first": "null"})
with self.assertRaisesRegex(ValueError, "JSON null"):
pmm.join_asof(
anchors.scan(), source.scan().select("first"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
def test_alignment_rejects_incremental_scans(self):
anchors = self.conn.create_table(
"incremental_anchors",
schema=pa.schema([
pa.field("episode_id", pa.int32()),
pa.field("event_time", pa.int64()),
]),
options={
"scan.mode": "incremental",
"incremental-between-timestamp": "0,9999999999999",
},
)
source = self._table("incremental_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{
"episode_id": 1, "event_time": 100, "value": 7,
}])
with self.assertRaisesRegex(
ValueError, "Temporal alignment.*incremental"):
pmm.join_asof(
anchors.scan(), source.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
def test_alignment_supports_zero_output_columns(self):
anchors = self._table("zero_output_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("zero_output_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{"episode_id": 1, "event_time": 100}])
result = pmm.join_asof(
anchors.scan().select("missing"),
source.scan().select("episode_id"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
self.assertEqual([], result.schema.names)
self.assertEqual([{}], result.to_list())
def test_alignment_supports_nested_projections(self):
anchors = self._table("nested_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"metadata": pa.struct([pa.field("value", pa.int32())]),
})
secondary = self._table("nested_secondary", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"payload": pa.struct([pa.field("value", pa.int32())]),
})
anchors.add([
{"episode_id": 1, "event_time": 100,
"metadata": {"value": 1}},
{"episode_id": 1, "event_time": 200, "metadata": None},
])
secondary.add([
{"episode_id": 1, "event_time": 100,
"payload": {"value": 7}},
])
rows = pmm.join_asof(
anchors.scan().select([
"episode_id", "event_time", "metadata.value"]),
secondary.scan().select("payload.value"),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=0,
).to_list()
self.assertEqual([1, None], [row["metadata_value"] for row in rows])
self.assertEqual(
[7, None], [row["payload_value"] for row in rows])
auth = TableQueryAuthResult(
filter=None,
column_masking={"payload": json.dumps({"name": "NULL"})},
)
selected = []
def query_auth(select):
selected.append(select)
return auth
secondary.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: query_auth)
masked = pmm.join_asof(
anchors.scan().select([
"episode_id", "event_time", "metadata.value"]),
secondary.scan().select("payload.value"),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=0,
).to_arrow()
self.assertEqual([None, None], masked["payload_value"].to_pylist())
self.assertTrue(masked.schema.field("payload_value").nullable)
self.assertIn(["payload"], selected)
self.assertTrue(all(
select is None or "payload_value" not in select
for select in selected
))
def test_nested_projection_cannot_shadow_internal_row_id(self):
anchors = self._table("nested_row_id_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"_ROW": pa.struct([pa.field("ID", pa.int32())]),
})
samples = self._table("nested_row_id_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([
{"episode_id": 1, "event_time": 10, "_ROW": {"ID": 1}},
{"episode_id": 1, "event_time": 20, "_ROW": {"ID": 0}},
])
samples.add([
{"episode_id": 1, "event_time": 10, "value": 7},
{"episode_id": 1, "event_time": 20, "value": 9},
])
rows = pmm.join_asof(
anchors.scan().select(["_ROW.ID", "event_time"]),
samples.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
self.assertEqual([1, 0], [row["_ROW_ID"] for row in rows])
self.assertEqual([7, 9], [row["value"] for row in rows])
def test_nested_payload_name_cannot_shadow_group_key(self):
anchors = self._table("nested_key_anchors", {
"payload_value": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("nested_key_samples", {
"payload_value": pa.int32(),
"event_time": pa.int64(),
"payload": pa.struct([pa.field("value", pa.int32())]),
})
anchors.add([{"payload_value": 1, "event_time": 10}])
samples.add([{
"payload_value": 1,
"event_time": 10,
"payload": {"value": 7},
}])
nested_only = pmm.join_asof(
anchors.scan(), samples.scan().select("payload.value"),
on="event_time", by="payload_value",
direction="nearest", tolerance=0,
).to_list()[0]
self.assertEqual(7, nested_only["payload_value_right"])
both = pmm.join_asof(
anchors.scan(),
samples.scan().select(["payload.value", "payload_value"]),
on="event_time", by="payload_value",
direction="nearest", tolerance=0,
).to_list()[0]
self.assertEqual(7, both["payload_value_right"])
self.assertNotIn("payload_value__0", both)
def test_alignment_rejects_unbound_nested_projection_masks(self):
anchors = self._table("nested_mask_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
source = self._table("nested_mask_source", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"payload": pa.struct([pa.field("value", pa.int32())]),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
source.add([{
"episode_id": 1,
"event_time": 100,
"payload": {"value": 7},
}])
auth = TableQueryAuthResult(
filter=None,
column_masking={"payload_value": json.dumps({"name": "NULL"})},
)
source.raw_table.catalog_environment.table_query_auth = (
lambda options, identifier: lambda select: auth)
with self.assertRaisesRegex(ValueError, "does not exist in table schema"):
pmm.join_asof(
anchors.scan(), source.scan().select("payload.value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_list()
def test_alignment_reuses_payload_scan_plans_across_batches(self):
anchors = self._table("planned_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
secondary = self._table("planned_secondary", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([
{"episode_id": 1, "event_time": value}
for value in range(8)
])
secondary.add([
{"episode_id": 1, "event_time": value, "value": value}
for value in range(8)
])
aligned = pmm.join_asof(
anchors.scan(), secondary.scan(),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
original_plan = TableScan.plan
with mock.patch.object(
TableScan, "plan", autospec=True,
side_effect=original_plan) as plan:
reader = aligned.to_arrow_batch_reader(batch_size=1)
self.assertEqual(8, sum(batch.num_rows for batch in reader))
self.assertEqual(4, plan.call_count)
def test_empty_source_stays_pinned_after_first_append(self):
anchors = self._table("pinned_empty_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("pinned_empty_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([{"episode_id": 1, "event_time": 10}])
aligned = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
samples.add([{
"episode_id": 1, "event_time": 10, "value": 7,
}])
self.assertIsNone(aligned.resolved_snapshots["right_1"]["snapshot_id"])
self.assertIsNone(aligned.to_list()[0]["value"])
@pytest.mark.python_read
def test_alignment_reuses_decoded_parquet_row_groups_across_batches(self):
anchors = self._table("cached_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("cached_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
row_count = 8192
rows = [
{"episode_id": 1, "event_time": value}
for value in range(row_count)
]
anchors.add(rows)
samples.add([
dict(row, value=row["event_time"])
for row in rows
])
original = FormatPyArrowReader._read_parquet_row_group_batches
for batch_size in (128, 1024, row_count):
with self.subTest(batch_size=batch_size):
decoded_rows = []
def tracked(reader, row_group, columns):
for batch in original(reader, row_group, columns):
if "value" in reader.existing_fields:
decoded_rows.append(batch.num_rows)
yield batch
with mock.patch.object(
FormatPyArrowReader,
"_read_parquet_row_group_batches", tracked):
aligned = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
reader = aligned.to_arrow_batch_reader(
batch_size=batch_size)
result = pa.Table.from_batches(list(reader))
self.assertEqual(
list(range(row_count)), result["value"].to_pylist())
self.assertEqual(row_count, sum(decoded_rows))
def test_alignment_streams_anchor_metadata(self):
anchors = self._table("streamed_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("streamed_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([
{"episode_id": 1, "event_time": value}
for value in range(4)
])
samples.add([
{"episode_id": 1, "event_time": value, "value": value}
for value in range(4)
])
aligned = pmm.join_asof(
anchors.scan(), samples.scan(),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
)
with mock.patch.object(
temporal, "_metadata_table",
wraps=temporal._metadata_table) as metadata_table:
self.assertEqual(4, len(aligned.to_list()))
# Only the indexed right side is collected into one metadata table.
self.assertEqual(1, metadata_table.call_count)
def test_alignment_splits_batches_before_arrow_offset_overflow(self):
child_count = 1 << 30
chunk = pa.ListArray.from_arrays(
pa.array([0, child_count], type=pa.int32()),
pa.nulls(child_count),
)
schema = pa.schema([pa.field("payload", chunk.type)])
class Fetcher:
def fetch(self, row_ids):
return pa.Table.from_arrays([
pa.chunked_array([chunk for _ in row_ids])
], schema=schema)
aligned = object.__new__(temporal.TemporalAlignment)
aligned._anchor_schema = schema
aligned._sources = ()
rows = [{temporal._ROW_ID: value} for value in range(2)]
batches = list(aligned._build_batches(
rows, Fetcher(), [], schema))
self.assertEqual([1, 1], [batch.num_rows for batch in batches])
self.assertTrue(all(batch.validate() is None for batch in batches))
self.assertEqual(schema, batches[0].schema)
def test_alignment_closes_anchor_stream_when_reader_closes(self):
anchors = self._table("closable_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
samples = self._table("closable_samples", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"value": pa.int32(),
})
anchors.add([
{"episode_id": 1, "event_time": value}
for value in range(2)
])
samples.add([
{"episode_id": 1, "event_time": value, "value": value}
for value in range(2)
])
closed = []
original = temporal._metadata_batches
def tracked_batches(*args):
try:
for batch in original(*args):
yield batch
finally:
closed.append(True)
with mock.patch.object(
temporal, "_metadata_batches", tracked_batches):
reader = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=0,
).to_arrow_batch_reader(batch_size=1)
next(reader)
reader.close()
self.assertEqual([True], closed)
def test_alignment_keeps_blob_payloads_as_descriptors(self):
anchors = self._table("blob_anchors", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
})
images = self._table("blob_images", {
"episode_id": pa.int32(),
"event_time": pa.int64(),
"image": pa.large_binary(),
})
anchors.add([{"episode_id": 1, "event_time": 100}])
images.add([{
"episode_id": 1,
"event_time": 100,
"image": b"encoded-image",
}])
row = pmm.join_asof(
anchors.scan(),
images.scan().select("image"),
on="event_time",
by="episode_id",
direction="nearest",
tolerance=0,
).to_list()[0]
descriptor = pmm.BlobDescriptor.deserialize(row["image"])
self.assertTrue(descriptor.uri.endswith(".blob"))
self.assertEqual(len(b"encoded-image"), descriptor.length)
def test_alignment_supports_timestamp_columns_with_different_names(self):
anchors = self._table("timestamp_anchors", {
"episode_id": pa.int32(),
"event_time": pa.timestamp("ms"),
})
samples = self._table("timestamp_samples", {
"episode_id": pa.int32(),
"captured_at": pa.timestamp("ms"),
"value": pa.int32(),
})
anchor_time = datetime(2026, 9, 1, 12, 0, 0, 100000)
sample_time = anchor_time - timedelta(milliseconds=5)
anchors.add([{"episode_id": 1, "event_time": anchor_time}])
samples.add([{
"episode_id": 1,
"captured_at": sample_time,
"value": 7,
}])
row = pmm.join_asof(
anchors.scan(),
samples.scan().select("value"),
on="event_time",
by="episode_id",
direction="nearest",
right_on="captured_at",
tolerance=timedelta(milliseconds=10),
).to_list()[0]
self.assertEqual(7, row["value"])
def test_alignment_preserves_nanosecond_timestamp_precision(self):
anchors = self._table("nanosecond_anchors", {
"episode_id": pa.int32(),
"event_time": pa.timestamp("ns"),
})
samples = self._table("nanosecond_samples", {
"episode_id": pa.int32(),
"event_time": pa.timestamp("ns"),
"value": pa.int32(),
})
anchors.add(pa.table({
"episode_id": pa.array([1], type=pa.int32()),
"event_time": pa.array(
[1_000_000_001], type=pa.int64()).cast(pa.timestamp("ns")),
}))
samples.add(pa.table({
"episode_id": pa.array([1, 1], type=pa.int32()),
"event_time": pa.array(
[1_000_000_000, 1_000_000_001],
type=pa.int64()).cast(pa.timestamp("ns")),
"value": pa.array([1, 2], type=pa.int32()),
}))
row = pmm.join_asof(
anchors.scan(), samples.scan().select("value"),
on="event_time", by="episode_id",
direction="nearest", tolerance=timedelta(0),
).to_list()[0]
self.assertEqual(2, row["value"])
def _table(self, name, fields):
return self.conn.create_table(name, schema=pa.schema([
pa.field(field_name, field_type)
for field_name, field_type in fields.items()
]))
if __name__ == "__main__":
unittest.main()