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

10 from collections.abc import Iterable

11

12 import pyarrow

13 from arrakis import Channel

14

15

16 def 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

40 def 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

52 def 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

72 def 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

84 def 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

103 def 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)])