Coverage for tests/unit/no_torch/test_zanj_basic.py: 100%

1 from __future__ import annotations

2

3 import json

4 import typing

5 from pathlib import Path

6

7 import numpy as np

8 import pandas as pd # type: ignore

9

10 from zanj import ZANJ

11

12 np.random.seed(0)

13

14

15 TEST_DATA_PATH: Path = Path("tests/junk_data")

16

17

18 def array_meta(x: typing.Any) -> dict:

19 if isinstance(x, np.ndarray):

20 return dict(

21 shape=list(x.shape),

22 dtype=str(x.dtype),

23 contents=str(x),

24 )

25 else:

26 return dict(

27 type=type(x).__name__,

28 contents=str(x),

29 )

30

31

32 def test_numpy():

33 data = dict(

34 name="testing zanj",

35 some_array=np.random.rand(128, 128),

36 some_other_array=np.random.rand(16, 64),

37 small_array=np.random.rand(4, 4),

38 )

39 fname: Path = TEST_DATA_PATH / "test_numpy.zanj"

40 z: ZANJ = ZANJ()

41 z.save(data, fname)

42 recovered_data = z.read(fname)

43

44 print(f"{list(data.keys()) = }")

45 print(f"{list(recovered_data.keys()) = }")

46 original_vals: dict = {k: array_meta(v) for k, v in data.items()}

47 print(json.dumps(original_vals, indent=2))

48 recovered_vals: dict = {k: array_meta(v) for k, v in recovered_data.items()}

49 print(json.dumps(recovered_vals, indent=2))

50

51 assert sorted(list(data.keys())) == sorted(list(recovered_data.keys()))

52 # assert all([type(data[k]) == type(recovered_data[k]) for k in data.keys()])

53

54 assert all(

55 [

56 data["name"] == recovered_data["name"],

57 np.allclose(data["some_array"], recovered_data["some_array"]),

58 np.allclose(data["some_other_array"], recovered_data["some_other_array"]),

59 np.allclose(data["small_array"], recovered_data["small_array"]),

60 ]

61 ), f"assert failed:\n{data = }\n{recovered_data = }"

62

63

64 def test_jsonl():

65 data = dict(

66 name="testing zanj jsonl",

67 iris_data=pd.read_csv("tests/input_data/iris.csv"),

68 brain_data=pd.read_csv("tests/input_data/brain_networks.csv"),

69 some_array=np.random.rand(128, 128),

70 )

71 fname: Path = TEST_DATA_PATH / "test_jsonl.zanj"

72 z: ZANJ = ZANJ()

73 z.save(data, fname)

74 recovered_data = z.read(fname)

75

76 assert sorted(list(data.keys())) == sorted(list(recovered_data.keys()))

77 # assert all([type(data[k]) == type(recovered_data[k]) for k in data.keys()])

78

79 assert all(

80 [

81 data["name"] == recovered_data["name"],

82 np.allclose(data["some_array"], recovered_data["some_array"]),

83 data["iris_data"].equals(recovered_data["iris_data"]),

84 data["brain_data"].equals(recovered_data["brain_data"]),

85 ]

86 )

87

88

89 def test_polars_dataframe():

90 import polars as pl

91

92 # basic dataframe with various types

93 data = dict(

94 name="testing zanj polars",

95 df=pl.DataFrame(

96 {

97 "a": [1, 2, 3],

98 "b": ["x", "y", "z"],

99 "c": [1.1, 2.2, 3.3],

100 }

101 ),

102 some_array=np.random.rand(128, 128),

103 )

104 fname: Path = TEST_DATA_PATH / "test_polars.zanj"

105 z: ZANJ = ZANJ()

106 z.save(data, fname)

107 recovered_data = z.read(fname)

108

109 assert sorted(list(data.keys())) == sorted(list(recovered_data.keys()))

110

111 assert all(

112 [

113 data["name"] == recovered_data["name"],

114 np.allclose(data["some_array"], recovered_data["some_array"]),

115 data["df"].equals(recovered_data["df"]),

116 ]

117 )

118

119

120 def test_polars_dataframe_empty():

121 """Test empty polars DataFrame serialization"""

122 import polars as pl

123

124 data = dict(

125 name="testing empty polars df",

126 empty_df=pl.DataFrame({"a": [], "b": [], "c": []}),

127 )

128 fname: Path = TEST_DATA_PATH / "test_polars_empty.zanj"

129 z: ZANJ = ZANJ()

130 z.save(data, fname)

131 recovered_data = z.read(fname)

132

133 assert data["name"] == recovered_data["name"]

134 assert recovered_data["empty_df"].shape == (0, 3)

135 assert recovered_data["empty_df"].columns == ["a", "b", "c"]

136

137

138 def test_polars_dataframe_large():

139 """Test larger polars DataFrame to ensure external storage works"""

140 import polars as pl

141

142 # create a larger dataframe

143 n_rows = 1000

144 data = dict(

145 name="testing large polars df",

146 large_df=pl.DataFrame(

147 {

148 "int_col": list(range(n_rows)),

149 "float_col": [float(i) * 0.1 for i in range(n_rows)],

150 "str_col": [f"row_{i}" for i in range(n_rows)],

151 "bool_col": [i % 2 == 0 for i in range(n_rows)],

152 }

153 ),

154 )

155 fname: Path = TEST_DATA_PATH / "test_polars_large.zanj"

156 z: ZANJ = ZANJ()

157 z.save(data, fname)

158 recovered_data = z.read(fname)

159

160 assert data["name"] == recovered_data["name"]

161 assert data["large_df"].equals(recovered_data["large_df"])

162

163

164 def test_polars_with_nulls():

165 """Test polars DataFrame with null values"""

166 import polars as pl

167

168 data = dict(

169 name="testing polars with nulls",

170 df_with_nulls=pl.DataFrame(

171 {

172 "a": [1, None, 3],

173 "b": ["x", "y", None],

174 "c": [1.1, None, 3.3],

175 }

176 ),

177 )

178 fname: Path = TEST_DATA_PATH / "test_polars_nulls.zanj"

179 z: ZANJ = ZANJ()

180 z.save(data, fname)

181 recovered_data = z.read(fname)

182

183 assert data["name"] == recovered_data["name"]

184 assert data["df_with_nulls"].equals(recovered_data["df_with_nulls"])