Coverage for arrakis_server/schemas.py: 85.7%

1# Copyright (c) 2022, California Institute of Technology and contributors 

2# 

3# You should have received a copy of the licensing terms for this 

4# software included in the file "LICENSE" located in the top-level 

5# directory of this package. If you did not, you can view a copy at 

6# https://git.ligo.org/ngdd/arrakis-server/-/raw/main/LICENSE 

7 

8"""Arrow Flight schema definitions.""" 

9 

10from collections.abc import Iterable 

11 

12import pyarrow 

13from arrakis import Channel 

14 

15 

16def stream(channels: Iterable[Channel]) -> pyarrow.Schema: 

17 """Create an Arrow Flight schema for `stream`. 

18 

19 Parameters 

20 ---------- 

21 channels : Iterable[Channel] 

22 The list of channels for the stream request. 

23 

24 Returns 

25 ------- 

26 pyarrow.Schema 

27 The stream schema. 

28 

29 """ 

30 columns = [pyarrow.field("time", pyarrow.int64(), nullable=False)] 

31 for channel in channels: 

32 dtype = pyarrow.from_numpy_dtype(channel.data_type) 

33 field = pyarrow.field(channel.name, pyarrow.list_(dtype)).with_metadata( 

34 {"rate": str(channel.sample_rate)} 

35 ) 

36 columns.append(field) 

37 return pyarrow.schema(columns) 

38 

39 

40def describe() -> pyarrow.Schema: 

41 """Create an Arrow Flight schema for `describe`. 

42 

43 Returns 

44 ------- 

45 pyarrow.Schema 

46 The describe schema. 

47 

48 """ 

49 return find() 

50 

51 

52def find() -> pyarrow.Schema: 

53 """Create an Arrow Flight schema for `find`. 

54 

55 Returns 

56 ------- 

57 pyarrow.Schema 

58 The find schema. 

59 

60 """ 

61 return pyarrow.schema( 

62 [ 

63 pyarrow.field("channel", pyarrow.string(), nullable=False), 

64 pyarrow.field("data_type", pyarrow.string(), nullable=False), 

65 pyarrow.field("sample_rate", pyarrow.int32(), nullable=False), 

66 pyarrow.field("partition_id", pyarrow.string()), 

67 pyarrow.field("publisher", pyarrow.string()), 

68 ] 

69 ) 

70 

71 

72def count() -> pyarrow.Schema: 

73 """Create an Arrow Flight schema for `count`. 

74 

75 Returns 

76 ------- 

77 pyarrow.Schema 

78 The count schema. 

79 

80 """ 

81 return pyarrow.schema([pyarrow.field("count", pyarrow.int64())]) 

82 

83 

84def partition() -> pyarrow.Schema: 

85 """Create an Arrow Flight schema for `partition`. 

86 

87 Returns 

88 ------- 

89 pyarrow.Schema 

90 The partition schema. 

91 

92 """ 

93 return pyarrow.schema( 

94 [ 

95 pyarrow.field("channel", pyarrow.string(), nullable=False), 

96 pyarrow.field("data_type", pyarrow.string(), nullable=False), 

97 pyarrow.field("sample_rate", pyarrow.int32(), nullable=False), 

98 pyarrow.field("partition_id", pyarrow.string()), 

99 ] 

100 ) 

101 

102 

103def publish() -> pyarrow.Schema: 

104 """Create an Arrow Flight schema for `publish`. 

105 

106 Returns 

107 ------- 

108 pyarrow.Schema 

109 The publish schema. 

110 

111 """ 

112 dtype = pyarrow.map_(pyarrow.string(), pyarrow.string()) 

113 return pyarrow.schema([pyarrow.field("properties", dtype, nullable=False)])