wehub-resource-sync
142 lines
4.6 KiB
Python
142 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pytest
|
|
import rerun as rr
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterable
|
|
|
|
VALID_SEQUENCE_CASES = [
|
|
([0, 1, 2, 3], pa.array([0, 1, 2, 3], type=pa.int64())),
|
|
([-1, 0, 1], pa.array([-1, 0, 1], type=pa.int64())),
|
|
(np.array([10, 20, 30]), pa.array([10, 20, 30], type=pa.int64())),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("sequence,expected", VALID_SEQUENCE_CASES)
|
|
def test_sequence_column(sequence: Iterable[int], expected: pa.Array) -> None:
|
|
column = rr.TimeColumn("sequence", sequence=sequence)
|
|
assert column.as_arrow_array() == expected
|
|
|
|
|
|
VALID_DURATION_CASES = [
|
|
([0, 1, 2, 3], pa.array([0, 1_000_000_000, 2_000_000_000, 3_000_000_000], type=pa.duration("ns"))),
|
|
(
|
|
np.arange(10, 15, 1.0),
|
|
pa.array(
|
|
[10_000_000_000, 11_000_000_000, 12_000_000_000, 13_000_000_000, 14_000_000_000],
|
|
type=pa.duration("ns"),
|
|
),
|
|
),
|
|
([0.0, 1.5, 2.25, 3.0], pa.array([0, 1_500_000_000, 2_250_000_000, 3_000_000_000], type=pa.duration("ns"))),
|
|
(
|
|
[
|
|
timedelta(seconds=0),
|
|
timedelta(seconds=1),
|
|
timedelta(seconds=1, microseconds=500000),
|
|
timedelta(seconds=2, microseconds=250000),
|
|
],
|
|
pa.array([0, 1_000_000_000, 1_500_000_000, 2_250_000_000], type=pa.duration("ns")),
|
|
),
|
|
(
|
|
[np.timedelta64(0, "s"), np.timedelta64(1, "s"), np.timedelta64(1500, "ms"), np.timedelta64(2250, "ms")],
|
|
pa.array([0, 1_000_000_000, 1_500_000_000, 2_250_000_000], type=pa.duration("ns")),
|
|
),
|
|
([-1, 0, 1], pa.array([-1_000_000_000, 0, 1_000_000_000], type=pa.duration("ns"))),
|
|
(
|
|
np.array([np.timedelta64(0, "s"), np.timedelta64(1, "s"), np.timedelta64(1500, "ms")]),
|
|
pa.array([0, 1_000_000_000, 1_500_000_000], type=pa.duration("ns")),
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("duration,expected", VALID_DURATION_CASES)
|
|
def test_duration_column(
|
|
duration: Iterable[int] | Iterable[float] | Iterable[timedelta] | Iterable[np.timedelta64], expected: pa.Array
|
|
) -> None:
|
|
column = rr.TimeColumn("duration", duration=duration)
|
|
|
|
assert column.as_arrow_array() == expected
|
|
|
|
|
|
VALID_TIMESTAMP_CASES = [
|
|
(
|
|
[0, 1, 2, 3],
|
|
pa.array(
|
|
[
|
|
np.datetime64("1970-01-01T00:00:00", "ns"),
|
|
np.datetime64("1970-01-01T00:00:01", "ns"),
|
|
np.datetime64("1970-01-01T00:00:02", "ns"),
|
|
np.datetime64("1970-01-01T00:00:03", "ns"),
|
|
],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
),
|
|
(
|
|
[0.0, 1.5, 2.25, 3.0],
|
|
pa.array(
|
|
[
|
|
np.datetime64("1970-01-01T00:00:00", "ns"),
|
|
np.datetime64("1970-01-01T00:00:01.50", "ns"),
|
|
np.datetime64("1970-01-01T00:00:02.25", "ns"),
|
|
np.datetime64("1970-01-01T00:00:03", "ns"),
|
|
],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
),
|
|
(
|
|
np.array([0, 1, 2, 3]),
|
|
pa.array(
|
|
[
|
|
np.datetime64("1970-01-01T00:00:00", "ns"),
|
|
np.datetime64("1970-01-01T00:00:01", "ns"),
|
|
np.datetime64("1970-01-01T00:00:02", "ns"),
|
|
np.datetime64("1970-01-01T00:00:03", "ns"),
|
|
],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
),
|
|
(
|
|
np.array([0.0, 1.5, 2.25, 3.0]),
|
|
pa.array(
|
|
[
|
|
np.datetime64("1970-01-01T00:00:00", "ns"),
|
|
np.datetime64("1970-01-01T00:00:01.50", "ns"),
|
|
np.datetime64("1970-01-01T00:00:02.25", "ns"),
|
|
np.datetime64("1970-01-01T00:00:03", "ns"),
|
|
],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
),
|
|
(
|
|
[datetime(2020, 1, 1, 0, 0, 0), datetime(2020, 1, 1, 0, 0, 1)],
|
|
pa.array(
|
|
[
|
|
np.datetime64("2020-01-01T00:00:00", "ns"),
|
|
np.datetime64("2020-01-01T00:00:01", "ns"),
|
|
],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
),
|
|
(
|
|
np.array([np.datetime64("2020-01-01T00:00:00"), np.datetime64("2020-01-01T00:00:01")]),
|
|
pa.array(
|
|
[np.datetime64("2020-01-01T00:00:00", "ns"), np.datetime64("2020-01-01T00:00:01", "ns")],
|
|
type=pa.timestamp("ns"),
|
|
),
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("timestamp,expected", VALID_TIMESTAMP_CASES)
|
|
def test_timestamp_column(
|
|
timestamp: Iterable[int] | Iterable[float] | Iterable[datetime] | Iterable[np.datetime64], expected: pa.Array
|
|
) -> None:
|
|
column = rr.TimeColumn("timestamp", timestamp=timestamp)
|
|
assert column.as_arrow_array() == expected
|