Coverage for src/signalk_cli/_arrow.py: 92%

1"""A small Arrow table built with nanoarrow, handed to dataframe libraries via the PyCapsule interface.""" 

2 

3import json 

4from collections.abc import Mapping, Sequence 

5from datetime import datetime 

6from typing import Any 

7 

8import nanoarrow as na 

9 

10_TIMESTAMP = na.timestamp("us", timezone="UTC") 

11 

12FEATHER_EXTENSIONS = {".feather", ".arrow", ".fea"} 

13 

14 

15def _is_number(v: object) -> bool: 

16 return isinstance(v, (int, float)) and not isinstance(v, bool) 

17 

18 

19def as_text(v: object) -> str | None: 

20 if v is None or isinstance(v, str): 

21 return v 

22 if isinstance(v, (dict, list)): 

23 return json.dumps(v) 

24 return str(v) 

25 

26 

27def _epoch_us(ts: str | datetime | None) -> int | None: 

28 if ts is None: 

29 return None 

30 dt = datetime.fromisoformat(ts) if isinstance(ts, str) else ts 

31 return round(dt.timestamp() * 1_000_000) 

32 

33 

34def infer_column(values: Sequence[Any]) -> tuple[list[str | float | bool | None], Any]: 

35 """Pick one Arrow type for a column of JSON values: float64, bool, or string. 

36 

37 Numbers (and nulls) become float64, booleans stay bool, and anything else 

38 — strings, objects, arrays, or a mix — becomes a string column, with 

39 objects and arrays JSON-encoded. 

40 """ 

41 present = [v for v in values if v is not None] 

42 if all(_is_number(v) for v in present): 

43 return [None if v is None else float(v) for v in values], na.float64() 

44 if all(isinstance(v, bool) for v in present): 44 ↛ 45line 44 didn't jump to line 45 because the condition on line 44 was never true

45 return list(values), na.bool_() 

46 return [as_text(v) for v in values], na.string() 

47 

48 

49class ArrowTable: 

50 """A table of query results, ready to load into any Arrow-aware dataframe library. 

51 

52 It implements the [Arrow PyCapsule interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html), 

53 so it can be passed straight to polars, pandas, pyarrow, DuckDB and others 

54 without signalk-cli depending on any of them: 

55 

56 Examples: 

57 ```python 

58 import polars as pl 

59 df = pl.DataFrame(table) 

60 import pandas as pd 

61 df = pd.DataFrame.from_arrow(table) # pandas >= 3.0 

62 import pyarrow as pa 

63 t = pa.table(table) 

64 ``` 

65 """ 

66 

67 def __init__(self, columns: dict[str, tuple[list, Any]]) -> None: 

68 self._columns = {name: values for name, (values, _) in columns.items()} 

69 self._types = {name: t for name, (_, t) in columns.items()} 

70 lengths = {len(v) for v in self._columns.values()} 

71 if len(lengths) > 1: 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 raise ValueError(f"Columns have different lengths: {sorted(lengths)}") 

73 self._num_rows = lengths.pop() if lengths else 0 

74 

75 @classmethod 

76 def from_rows( 

77 cls, 

78 timestamps: Sequence[str | datetime | None], 

79 columns: Mapping[str, Sequence[Any]], 

80 *, 

81 text_columns: Sequence[str] = (), 

82 timestamp_name: str = "timestamp", 

83 ) -> "ArrowTable": 

84 """Build a table with a UTC timestamp column followed by the given columns. 

85 

86 Column types are inferred from the values (see `infer_column`), 

87 except `text_columns`, which are always strings so an empty result 

88 keeps the same schema. 

89 """ 

90 built: dict[str, tuple[list, Any]] = { 

91 timestamp_name: ([_epoch_us(t) for t in timestamps], _TIMESTAMP) 

92 } 

93 for name, values in columns.items(): 

94 if name in text_columns: 

95 built[name] = ([as_text(v) for v in values], na.string()) 

96 else: 

97 built[name] = infer_column(values) 

98 return cls(built) 

99 

100 @property 

101 def num_rows(self) -> int: 

102 return self._num_rows 

103 

104 @property 

105 def column_names(self) -> list[str]: 

106 return list(self._columns) 

107 

108 def __len__(self) -> int: 

109 return self._num_rows 

110 

111 def __repr__(self) -> str: 

112 cols = ", ".join(f"{n}: {self._types[n].type.name}" for n in self._columns) 

113 return f"ArrowTable({self._num_rows} rows; {cols})" 

114 

115 def to_pydict(self) -> dict[str, list]: 

116 """Return the columns as plain Python lists (timestamps as epoch microseconds).""" 

117 return {name: list(values) for name, values in self._columns.items()} 

118 

119 def _batch(self): 

120 children = [ 

121 na.c_array(values, self._types[name]) 

122 for name, values in self._columns.items() 

123 ] 

124 schema = na.struct({name: self._types[name] for name in self._columns}) 

125 return na.c_array_from_buffers( 

126 schema, self._num_rows, [None], children=children 

127 ) 

128 

129 def __arrow_c_schema__(self): 

130 return na.c_schema( 

131 na.struct({name: self._types[name] for name in self._columns}) 

132 ).__arrow_c_schema__() 

133 

134 def __arrow_c_stream__(self, requested_schema=None): 

135 return na.ArrayStream(self._batch()).__arrow_c_stream__(requested_schema) 

136 

137 

138def write_feather(table: ArrowTable, output: str) -> None: 

139 """Write a table to a Feather (Arrow IPC file) — needs pyarrow.""" 

140 try: 

141 import pyarrow as pa 

142 from pyarrow import feather 

143 except ImportError: 

144 raise ImportError( 

145 "pyarrow is required for Feather output: pip install 'signalk-cli[feather]'" 

146 ) from None 

147 feather.write_feather(pa.table(table), output)