mirror of
https://github.com/aykhans/AzSuicideDataVisualization.git
synced 2025-04-22 10:28:02 +00:00
327 lines
12 KiB
Python
327 lines
12 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
from collections import OrderedDict
|
|
import io
|
|
import itertools
|
|
import json
|
|
import pickle
|
|
import string
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import pyarrow as pa
|
|
from pyarrow.json import read_json, ReadOptions, ParseOptions
|
|
|
|
|
|
def generate_col_names():
|
|
# 'a', 'b'... 'z', then 'aa', 'ab'...
|
|
letters = string.ascii_lowercase
|
|
yield from letters
|
|
for first in letters:
|
|
for second in letters:
|
|
yield first + second
|
|
|
|
|
|
def make_random_json(num_cols=2, num_rows=10, linesep='\r\n'):
|
|
arr = np.random.RandomState(42).randint(0, 1000, size=(num_cols, num_rows))
|
|
col_names = list(itertools.islice(generate_col_names(), num_cols))
|
|
lines = []
|
|
for row in arr.T:
|
|
json_obj = OrderedDict([(k, int(v)) for (k, v) in zip(col_names, row)])
|
|
lines.append(json.dumps(json_obj))
|
|
data = linesep.join(lines).encode()
|
|
columns = [pa.array(col, type=pa.int64()) for col in arr]
|
|
expected = pa.Table.from_arrays(columns, col_names)
|
|
return data, expected
|
|
|
|
|
|
def check_options_class_pickling(cls, **attr_values):
|
|
opts = cls(**attr_values)
|
|
new_opts = pickle.loads(pickle.dumps(opts,
|
|
protocol=pickle.HIGHEST_PROTOCOL))
|
|
for name, value in attr_values.items():
|
|
assert getattr(new_opts, name) == value
|
|
|
|
|
|
def test_read_options():
|
|
cls = ReadOptions
|
|
opts = cls()
|
|
|
|
assert opts.block_size > 0
|
|
opts.block_size = 12345
|
|
assert opts.block_size == 12345
|
|
|
|
assert opts.use_threads is True
|
|
opts.use_threads = False
|
|
assert opts.use_threads is False
|
|
|
|
opts = cls(block_size=1234, use_threads=False)
|
|
assert opts.block_size == 1234
|
|
assert opts.use_threads is False
|
|
|
|
check_options_class_pickling(cls, block_size=1234,
|
|
use_threads=False)
|
|
|
|
|
|
def test_parse_options():
|
|
cls = ParseOptions
|
|
opts = cls()
|
|
assert opts.newlines_in_values is False
|
|
assert opts.explicit_schema is None
|
|
|
|
opts.newlines_in_values = True
|
|
assert opts.newlines_in_values is True
|
|
|
|
schema = pa.schema([pa.field('foo', pa.int32())])
|
|
opts.explicit_schema = schema
|
|
assert opts.explicit_schema == schema
|
|
|
|
assert opts.unexpected_field_behavior == "infer"
|
|
for value in ["ignore", "error", "infer"]:
|
|
opts.unexpected_field_behavior = value
|
|
assert opts.unexpected_field_behavior == value
|
|
|
|
with pytest.raises(ValueError):
|
|
opts.unexpected_field_behavior = "invalid-value"
|
|
|
|
check_options_class_pickling(cls, explicit_schema=schema,
|
|
newlines_in_values=False,
|
|
unexpected_field_behavior="ignore")
|
|
|
|
|
|
class BaseTestJSONRead:
|
|
|
|
def read_bytes(self, b, **kwargs):
|
|
return self.read_json(pa.py_buffer(b), **kwargs)
|
|
|
|
def check_names(self, table, names):
|
|
assert table.num_columns == len(names)
|
|
assert [c.name for c in table.columns] == names
|
|
|
|
def test_file_object(self):
|
|
data = b'{"a": 1, "b": 2}\n'
|
|
expected_data = {'a': [1], 'b': [2]}
|
|
bio = io.BytesIO(data)
|
|
table = self.read_json(bio)
|
|
assert table.to_pydict() == expected_data
|
|
# Text files not allowed
|
|
sio = io.StringIO(data.decode())
|
|
with pytest.raises(TypeError):
|
|
self.read_json(sio)
|
|
|
|
def test_block_sizes(self):
|
|
rows = b'{"a": 1}\n{"a": 2}\n{"a": 3}'
|
|
read_options = ReadOptions()
|
|
parse_options = ParseOptions()
|
|
|
|
for data in [rows, rows + b'\n']:
|
|
for newlines_in_values in [False, True]:
|
|
parse_options.newlines_in_values = newlines_in_values
|
|
read_options.block_size = 4
|
|
with pytest.raises(ValueError,
|
|
match="try to increase block size"):
|
|
self.read_bytes(data, read_options=read_options,
|
|
parse_options=parse_options)
|
|
|
|
# Validate reader behavior with various block sizes.
|
|
# There used to be bugs in this area.
|
|
for block_size in range(9, 20):
|
|
read_options.block_size = block_size
|
|
table = self.read_bytes(data, read_options=read_options,
|
|
parse_options=parse_options)
|
|
assert table.to_pydict() == {'a': [1, 2, 3]}
|
|
|
|
def test_no_newline_at_end(self):
|
|
rows = b'{"a": 1,"b": 2, "c": 3}\n{"a": 4,"b": 5, "c": 6}'
|
|
table = self.read_bytes(rows)
|
|
assert table.to_pydict() == {
|
|
'a': [1, 4],
|
|
'b': [2, 5],
|
|
'c': [3, 6],
|
|
}
|
|
|
|
def test_simple_ints(self):
|
|
# Infer integer columns
|
|
rows = b'{"a": 1,"b": 2, "c": 3}\n{"a": 4,"b": 5, "c": 6}\n'
|
|
table = self.read_bytes(rows)
|
|
schema = pa.schema([('a', pa.int64()),
|
|
('b', pa.int64()),
|
|
('c', pa.int64())])
|
|
assert table.schema == schema
|
|
assert table.to_pydict() == {
|
|
'a': [1, 4],
|
|
'b': [2, 5],
|
|
'c': [3, 6],
|
|
}
|
|
|
|
def test_simple_varied(self):
|
|
# Infer various kinds of data
|
|
rows = (b'{"a": 1,"b": 2, "c": "3", "d": false}\n'
|
|
b'{"a": 4.0, "b": -5, "c": "foo", "d": true}\n')
|
|
table = self.read_bytes(rows)
|
|
schema = pa.schema([('a', pa.float64()),
|
|
('b', pa.int64()),
|
|
('c', pa.string()),
|
|
('d', pa.bool_())])
|
|
assert table.schema == schema
|
|
assert table.to_pydict() == {
|
|
'a': [1.0, 4.0],
|
|
'b': [2, -5],
|
|
'c': ["3", "foo"],
|
|
'd': [False, True],
|
|
}
|
|
|
|
def test_simple_nulls(self):
|
|
# Infer various kinds of data, with nulls
|
|
rows = (b'{"a": 1, "b": 2, "c": null, "d": null, "e": null}\n'
|
|
b'{"a": null, "b": -5, "c": "foo", "d": null, "e": true}\n'
|
|
b'{"a": 4.5, "b": null, "c": "nan", "d": null,"e": false}\n')
|
|
table = self.read_bytes(rows)
|
|
schema = pa.schema([('a', pa.float64()),
|
|
('b', pa.int64()),
|
|
('c', pa.string()),
|
|
('d', pa.null()),
|
|
('e', pa.bool_())])
|
|
assert table.schema == schema
|
|
assert table.to_pydict() == {
|
|
'a': [1.0, None, 4.5],
|
|
'b': [2, -5, None],
|
|
'c': [None, "foo", "nan"],
|
|
'd': [None, None, None],
|
|
'e': [None, True, False],
|
|
}
|
|
|
|
def test_empty_lists(self):
|
|
# ARROW-10955: Infer list(null)
|
|
rows = b'{"a": []}'
|
|
table = self.read_bytes(rows)
|
|
schema = pa.schema([('a', pa.list_(pa.null()))])
|
|
assert table.schema == schema
|
|
assert table.to_pydict() == {'a': [[]]}
|
|
|
|
def test_empty_rows(self):
|
|
rows = b'{}\n{}\n'
|
|
table = self.read_bytes(rows)
|
|
schema = pa.schema([])
|
|
assert table.schema == schema
|
|
assert table.num_columns == 0
|
|
assert table.num_rows == 2
|
|
|
|
def test_reconcile_accross_blocks(self):
|
|
# ARROW-12065: reconciling inferred types across blocks
|
|
first_row = b'{ }\n'
|
|
read_options = ReadOptions(block_size=len(first_row))
|
|
for next_rows, expected_pylist in [
|
|
(b'{"a": 0}', [None, 0]),
|
|
(b'{"a": []}', [None, []]),
|
|
(b'{"a": []}\n{"a": [[1]]}', [None, [], [[1]]]),
|
|
(b'{"a": {}}', [None, {}]),
|
|
(b'{"a": {}}\n{"a": {"b": {"c": 1}}}',
|
|
[None, {"b": None}, {"b": {"c": 1}}]),
|
|
]:
|
|
table = self.read_bytes(first_row + next_rows,
|
|
read_options=read_options)
|
|
expected = {"a": expected_pylist}
|
|
assert table.to_pydict() == expected
|
|
# Check that the issue was exercised
|
|
assert table.column("a").num_chunks > 1
|
|
|
|
def test_explicit_schema_with_unexpected_behaviour(self):
|
|
# infer by default
|
|
rows = (b'{"foo": "bar", "num": 0}\n'
|
|
b'{"foo": "baz", "num": 1}\n')
|
|
schema = pa.schema([
|
|
('foo', pa.binary())
|
|
])
|
|
|
|
opts = ParseOptions(explicit_schema=schema)
|
|
table = self.read_bytes(rows, parse_options=opts)
|
|
assert table.schema == pa.schema([
|
|
('foo', pa.binary()),
|
|
('num', pa.int64())
|
|
])
|
|
assert table.to_pydict() == {
|
|
'foo': [b'bar', b'baz'],
|
|
'num': [0, 1],
|
|
}
|
|
|
|
# ignore the unexpected fields
|
|
opts = ParseOptions(explicit_schema=schema,
|
|
unexpected_field_behavior="ignore")
|
|
table = self.read_bytes(rows, parse_options=opts)
|
|
assert table.schema == pa.schema([
|
|
('foo', pa.binary()),
|
|
])
|
|
assert table.to_pydict() == {
|
|
'foo': [b'bar', b'baz'],
|
|
}
|
|
|
|
# raise error
|
|
opts = ParseOptions(explicit_schema=schema,
|
|
unexpected_field_behavior="error")
|
|
with pytest.raises(pa.ArrowInvalid,
|
|
match="JSON parse error: unexpected field"):
|
|
self.read_bytes(rows, parse_options=opts)
|
|
|
|
def test_small_random_json(self):
|
|
data, expected = make_random_json(num_cols=2, num_rows=10)
|
|
table = self.read_bytes(data)
|
|
assert table.schema == expected.schema
|
|
assert table.equals(expected)
|
|
assert table.to_pydict() == expected.to_pydict()
|
|
|
|
def test_stress_block_sizes(self):
|
|
# Test a number of small block sizes to stress block stitching
|
|
data_base, expected = make_random_json(num_cols=2, num_rows=100)
|
|
read_options = ReadOptions()
|
|
parse_options = ParseOptions()
|
|
|
|
for data in [data_base, data_base.rstrip(b'\r\n')]:
|
|
for newlines_in_values in [False, True]:
|
|
parse_options.newlines_in_values = newlines_in_values
|
|
for block_size in [22, 23, 37]:
|
|
read_options.block_size = block_size
|
|
table = self.read_bytes(data, read_options=read_options,
|
|
parse_options=parse_options)
|
|
assert table.schema == expected.schema
|
|
if not table.equals(expected):
|
|
# Better error output
|
|
assert table.to_pydict() == expected.to_pydict()
|
|
|
|
|
|
class TestSerialJSONRead(BaseTestJSONRead, unittest.TestCase):
|
|
|
|
def read_json(self, *args, **kwargs):
|
|
read_options = kwargs.setdefault('read_options', ReadOptions())
|
|
read_options.use_threads = False
|
|
table = read_json(*args, **kwargs)
|
|
table.validate(full=True)
|
|
return table
|
|
|
|
|
|
class TestParallelJSONRead(BaseTestJSONRead, unittest.TestCase):
|
|
|
|
def read_json(self, *args, **kwargs):
|
|
read_options = kwargs.setdefault('read_options', ReadOptions())
|
|
read_options.use_threads = True
|
|
table = read_json(*args, **kwargs)
|
|
table.validate(full=True)
|
|
return table
|