12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
14
15
- import db_dtypes # type: ignore
16
15
import geopandas as gpd # type: ignore
17
16
import pandas as pd
18
17
import pyarrow as pa
19
18
import pytest
20
19
21
20
import bigframes .bigquery as bbq
22
- import bigframes .dtypes
21
+ import bigframes .dtypes as dtypes
23
22
import bigframes .pandas as bpd
24
23
25
24
26
25
@pytest .mark .parametrize (
27
26
("json_path" , "expected_json" ),
28
27
[
29
- pytest .param ("$.a" , [{"a" : 10 }], id = "simple" ),
30
- pytest .param ("$.a.b.c" , [{"a" : {"b" : {"c" : 10 , "d" : []}}}], id = "nested" ),
28
+ pytest .param ("$.a" , [' {"a": 10}' ], id = "simple" ),
29
+ pytest .param ("$.a.b.c" , [' {"a": {"b": {"c": 10, "d": []}}}' ], id = "nested" ),
31
30
],
32
31
)
33
32
def test_json_set_at_json_path (json_path , expected_json ):
34
- original_json = [{"a" : {"b" : {"c" : "tester" , "d" : []}}}]
35
- s = bpd .Series (original_json , dtype = db_dtypes . JSONDtype () )
33
+ original_json = [' {"a": {"b": {"c": "tester", "d": []}}}' ]
34
+ s = bpd .Series (original_json , dtype = dtypes . JSON_DTYPE )
36
35
actual = bbq .json_set (s , json_path_value_pairs = [(json_path , 10 )])
37
36
38
- expected = bpd .Series (expected_json , dtype = db_dtypes . JSONDtype () )
37
+ expected = bpd .Series (expected_json , dtype = dtypes . JSON_DTYPE )
39
38
pd .testing .assert_series_equal (
40
39
actual .to_pandas (),
41
40
expected .to_pandas (),
@@ -45,47 +44,49 @@ def test_json_set_at_json_path(json_path, expected_json):
45
44
@pytest .mark .parametrize (
46
45
("json_value" , "expected_json" ),
47
46
[
48
- pytest .param (10 , [{"a" : {"b" : 10 }}, {"a" : {"b" : 10 }}], id = "int" ),
49
- pytest .param (0.333 , [{"a" : {"b" : 0.333 }}, {"a" : {"b" : 0.333 }}], id = "float" ),
50
- pytest .param ("eng" , [{"a" : {"b" : "eng" }}, {"a" : {"b" : "eng" }}], id = "string" ),
51
- pytest .param ([1 , 2 ], [{"a" : {"b" : 1 }}, {"a" : {"b" : 2 }}], id = "series" ),
47
+ pytest .param (10 , ['{"a": {"b": 10}}' , '{"a": {"b": 10}}' ], id = "int" ),
48
+ pytest .param (0.333 , ['{"a": {"b": 0.333}}' , '{"a": {"b": 0.333}}' ], id = "float" ),
49
+ pytest .param (
50
+ "eng" , ['{"a": {"b": "eng"}}' , '{"a": {"b": "eng"}}' ], id = "string"
51
+ ),
52
+ pytest .param ([1 , 2 ], ['{"a": {"b": 1}}' , '{"a": {"b": 2}}' ], id = "series" ),
52
53
],
53
54
)
54
55
def test_json_set_at_json_value_type (json_value , expected_json ):
55
- original_json = [{"a" : {"b" : "dev" }}, {"a" : {"b" : [1 , 2 ]}}]
56
- s = bpd .Series (original_json , dtype = db_dtypes . JSONDtype () )
56
+ original_json = [' {"a": {"b": "dev"}}' , ' {"a": {"b": [1, 2]}}' ]
57
+ s = bpd .Series (original_json , dtype = dtypes . JSON_DTYPE )
57
58
actual = bbq .json_set (s , json_path_value_pairs = [("$.a.b" , json_value )])
58
59
59
- expected = bpd .Series (expected_json , dtype = db_dtypes . JSONDtype () )
60
+ expected = bpd .Series (expected_json , dtype = dtypes . JSON_DTYPE )
60
61
pd .testing .assert_series_equal (
61
62
actual .to_pandas (),
62
63
expected .to_pandas (),
63
64
)
64
65
65
66
66
67
def test_json_set_w_more_pairs ():
67
- original_json = [{"a" : 2 }, {"b" : 5 }, {"c" : 1 }]
68
- s = bpd .Series (original_json , dtype = db_dtypes . JSONDtype () )
68
+ original_json = [' {"a": 2}' , ' {"b": 5}' , ' {"c": 1}' ]
69
+ s = bpd .Series (original_json , dtype = dtypes . JSON_DTYPE )
69
70
actual = bbq .json_set (
70
71
s , json_path_value_pairs = [("$.a" , 1 ), ("$.b" , 2 ), ("$.a" , [3 , 4 , 5 ])]
71
72
)
72
73
73
- expected_json = [{"a" : 3 , "b" : 2 }, {"a" : 4 , "b" : 2 }, {"a" : 5 , "b" : 2 , "c" : 1 }]
74
- expected = bpd .Series (expected_json , dtype = db_dtypes . JSONDtype () )
74
+ expected_json = [' {"a": 3, "b": 2}' , ' {"a": 4, "b": 2}' , ' {"a": 5, "b": 2, "c": 1}' ]
75
+ expected = bpd .Series (expected_json , dtype = dtypes . JSON_DTYPE )
75
76
pd .testing .assert_series_equal (
76
77
actual .to_pandas (),
77
78
expected .to_pandas (),
78
79
)
79
80
80
81
81
82
def test_json_set_w_invalid_json_path_value_pairs ():
82
- s = bpd .Series ([{"a" : 10 }], dtype = db_dtypes . JSONDtype () )
83
+ s = bpd .Series ([' {"a": 10}' ], dtype = dtypes . JSON_DTYPE )
83
84
with pytest .raises (ValueError ):
84
85
bbq .json_set (s , json_path_value_pairs = [("$.a" , 1 , 100 )]) # type: ignore
85
86
86
87
87
88
def test_json_set_w_invalid_value_type ():
88
- s = bpd .Series ([{"a" : 10 }], dtype = db_dtypes . JSONDtype () )
89
+ s = bpd .Series ([' {"a": 10}' ], dtype = dtypes . JSON_DTYPE )
89
90
with pytest .raises (TypeError ):
90
91
bbq .json_set (
91
92
s ,
@@ -101,17 +102,18 @@ def test_json_set_w_invalid_value_type():
101
102
102
103
103
104
def test_json_set_w_invalid_series_type ():
105
+ s = bpd .Series ([1 , 2 ])
104
106
with pytest .raises (TypeError ):
105
- bbq .json_set (bpd . Series ([ 1 , 2 ]) , json_path_value_pairs = [("$.a" , 1 )])
107
+ bbq .json_set (s , json_path_value_pairs = [("$.a" , 1 )])
106
108
107
109
108
110
def test_json_extract_from_json ():
109
111
s = bpd .Series (
110
- [{"a" : {"b" : [1 , 2 ]}}, {"a" : {"c" : 1 }}, {"a" : {"b" : 0 }}],
111
- dtype = db_dtypes . JSONDtype () ,
112
+ [' {"a": {"b": [1, 2]}}' , ' {"a": {"c": 1}}' , ' {"a": {"b": 0}}' ],
113
+ dtype = dtypes . JSON_DTYPE ,
112
114
)
113
115
actual = bbq .json_extract (s , "$.a.b" ).to_pandas ()
114
- expected = bpd .Series ([[1 , 2 ], None , 0 ], dtype = db_dtypes . JSONDtype () ).to_pandas ()
116
+ expected = bpd .Series ([" [1, 2]" , None , "0" ], dtype = dtypes . JSON_DTYPE ).to_pandas ()
115
117
pd .testing .assert_series_equal (
116
118
actual ,
117
119
expected ,
@@ -132,14 +134,15 @@ def test_json_extract_from_string():
132
134
133
135
134
136
def test_json_extract_w_invalid_series_type ():
137
+ s = bpd .Series ([1 , 2 ])
135
138
with pytest .raises (TypeError ):
136
- bbq .json_extract (bpd . Series ([ 1 , 2 ]) , "$.a" )
139
+ bbq .json_extract (s , "$.a" )
137
140
138
141
139
142
def test_json_extract_array_from_json ():
140
143
s = bpd .Series (
141
- [{"a" : ["ab" , "2" , "3 xy" ]}, {"a" : []}, {"a" : ["4" , "5" ]}, {} ],
142
- dtype = db_dtypes . JSONDtype () ,
144
+ [' {"a": ["ab", "2", "3 xy"]}' , ' {"a": []}' , ' {"a": ["4", "5"]}' , "{}" ],
145
+ dtype = dtypes . JSON_DTYPE ,
143
146
)
144
147
actual = bbq .json_extract_array (s , "$.a" )
145
148
@@ -225,7 +228,7 @@ def test_json_extract_string_array_from_array_strings():
225
228
226
229
def test_json_extract_string_array_as_float_array_from_array_strings ():
227
230
s = bpd .Series (["[1, 2.5, 3]" , "[]" , "[4,5]" ])
228
- actual = bbq .json_extract_string_array (s , value_dtype = bigframes . dtypes .FLOAT_DTYPE )
231
+ actual = bbq .json_extract_string_array (s , value_dtype = dtypes .FLOAT_DTYPE )
229
232
expected = bpd .Series ([[1 , 2.5 , 3 ], [], [4 , 5 ]])
230
233
pd .testing .assert_series_equal (
231
234
actual .to_pandas (),
0 commit comments