python/pyarrow/tests/test_scalars.py - arrow - Git at Google

 # -*- coding: utf-8 -*-
 # 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.

 import pytest

 import numpy as np

 from pyarrow.compat import unittest, u, unicode_type
 import pyarrow as pa


 class TestScalars(unittest.TestCase):

     def test_null_singleton(self):
         with pytest.raises(Exception):
             pa.NAType()

     def test_nulls(self):
         arr = pa.array([None, None])
         for v in arr:
             assert v is pa.NA
             assert v.as_py() is None

     def test_bool(self):
         arr = pa.array([True, None, False, None])

         v = arr[0]
         assert isinstance(v, pa.BooleanValue)
         assert repr(v) == "True"
         assert str(v) == "True"
         assert v.as_py() is True

         assert arr[1] is pa.NA

     def test_int64(self):
         arr = pa.array([1, 2, None])

         v = arr[0]
         assert isinstance(v, pa.Int64Value)
         assert repr(v) == "1"
         assert str(v) == "1"
         assert v.as_py() == 1
         assert v == 1

         assert arr[2] is pa.NA

     def test_double(self):
         arr = pa.array([1.5, None, 3])

         v = arr[0]
         assert isinstance(v, pa.DoubleValue)
         assert repr(v) == "1.5"
         assert str(v) == "1.5"
         assert v.as_py() == 1.5
         assert v == 1.5

         assert arr[1] is pa.NA

         v = arr[2]
         assert v.as_py() == 3.0

     def test_half_float(self):
         arr = pa.array([np.float16(1.5), None], type=pa.float16())
         v = arr[0]
         assert isinstance(v, pa.HalfFloatValue)
         assert repr(v) == "1.5"
         assert str(v) == "1.5"
         assert v.as_py() == 1.5
         assert v == 1.5

         assert arr[1] is pa.NA

     def test_string_unicode(self):
         arr = pa.array([u'foo', None, u'mañana'])

         v = arr[0]
         assert isinstance(v, pa.StringValue)
         assert v.as_py() == u'foo'
         assert repr(v) == repr(u"foo")
         assert str(v) == str(u"foo")
         assert v == u'foo'
         # Assert that newly created values are equal to the previously created
         # one.
         assert v == arr[0]

         assert arr[1] is pa.NA

         v = arr[2].as_py()
         assert v == u'mañana'
         assert isinstance(v, unicode_type)

     def test_bytes(self):
         arr = pa.array([b'foo', None, u('bar')])

         def check_value(v, expected):
             assert isinstance(v, pa.BinaryValue)
             assert v.as_py() == expected
             assert str(v) == str(expected)
             assert repr(v) == repr(expected)
             assert v == expected
             assert v != b'xxxxx'
             buf = v.as_buffer()
             assert isinstance(buf, pa.Buffer)
             assert buf.to_pybytes() == expected

         check_value(arr[0], b'foo')
         assert arr[1] is pa.NA
         check_value(arr[2], b'bar')

     def test_fixed_size_bytes(self):
         data = [b'foof', None, b'barb']
         arr = pa.array(data, type=pa.binary(4))

         v = arr[0]
         assert isinstance(v, pa.FixedSizeBinaryValue)
         assert v.as_py() == b'foof'

         assert arr[1] is pa.NA

         v = arr[2].as_py()
         assert v == b'barb'
         assert isinstance(v, bytes)

     def test_list(self):
         arr = pa.array([['foo', None], None, ['bar'], []])

         v = arr[0]
         assert len(v) == 2
         assert isinstance(v, pa.ListValue)
         assert repr(v) == "['foo', None]"
         assert v.as_py() == ['foo', None]
         assert v[0].as_py() == 'foo'
         assert v[1] is pa.NA
         assert v[-1] == v[1]
         assert v[-2] == v[0]
         with pytest.raises(IndexError):
             v[-3]
         with pytest.raises(IndexError):
             v[2]

         assert arr[1] is pa.NA

         v = arr[3]
         assert len(v) == 0

     @pytest.mark.pandas
     def test_timestamp(self):
         import pandas as pd
         arr = pd.date_range('2000-01-01 12:34:56', periods=10).values

         units = ['ns', 'us', 'ms', 's']

         for i, unit in enumerate(units):
             dtype = 'datetime64[{0}]'.format(unit)
             arrow_arr = pa.Array.from_pandas(arr.astype(dtype))
             expected = pd.Timestamp('2000-01-01 12:34:56')

             assert arrow_arr[0].as_py() == expected
             assert arrow_arr[0].value * 1000**i == expected.value

             tz = 'America/New_York'
             arrow_type = pa.timestamp(unit, tz=tz)

             dtype = 'datetime64[{0}]'.format(unit)
             arrow_arr = pa.Array.from_pandas(arr.astype(dtype),
                                              type=arrow_type)
             expected = (pd.Timestamp('2000-01-01 12:34:56')
                         .tz_localize('utc')
                         .tz_convert(tz))

             assert arrow_arr[0].as_py() == expected
             assert arrow_arr[0].value * 1000**i == expected.value

     @pytest.mark.pandas
     def test_dictionary(self):
         import pandas as pd
         colors = ['red', 'green', 'blue']
         colors_dict = {'red': 0, 'green': 1, 'blue': 2}
         values = pd.Series(colors * 4)

         categorical = pd.Categorical(values, categories=colors)

         v = pa.DictionaryArray.from_arrays(categorical.codes,
                                            categorical.categories)
         for i, c in enumerate(values):
             assert v[i].as_py() == c
             assert v[i].dictionary_value == c
             assert v[i].index_value == colors_dict[c]

     def test_int_hash(self):
         # ARROW-640
         int_arr = pa.array([1, 1, 2, 1])
         assert hash(int_arr[0]) == hash(1)

     def test_float_hash(self):
         # ARROW-640
         float_arr = pa.array([1.4, 1.2, 2.5, 1.8])
         assert hash(float_arr[0]) == hash(1.4)

     def test_string_hash(self):
         # ARROW-640
         str_arr = pa.array(["foo", "bar"])
         assert hash(str_arr[1]) == hash("bar")

     def test_bytes_hash(self):
         # ARROW-640
         byte_arr = pa.array([b'foo', None, b'bar'])
         assert hash(byte_arr[2]) == hash(b"bar")

     def test_array_to_set(self):
         # ARROW-640
         arr = pa.array([1, 1, 2, 1])
         set_from_array = set(arr)
         assert isinstance(set_from_array, set)
         assert set_from_array == {1, 2}

     def test_struct_value_subscripting(self):
         ty = pa.struct([pa.field('x', pa.int16()),
                         pa.field('y', pa.float32())])
         arr = pa.array([(1, 2.5), (3, 4.5), (5, 6.5)], type=ty)

         assert arr[0]['x'] == 1
         assert arr[0]['y'] == 2.5
         assert arr[1]['x'] == 3
         assert arr[1]['y'] == 4.5
         assert arr[2]['x'] == 5
         assert arr[2]['y'] == 6.5

         with pytest.raises(IndexError):
             arr[4]['non-existent']

         with pytest.raises(KeyError):
             arr[0]['non-existent']
	# -- coding: utf-8 --
	# 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.

	import pytest

	import numpy as np

	from pyarrow.compat import unittest, u, unicode_type
	import pyarrow as pa


	class TestScalars(unittest.TestCase):

	def test_null_singleton(self):
	with pytest.raises(Exception):
	pa.NAType()

	def test_nulls(self):
	arr = pa.array([None, None])
	for v in arr:
	assert v is pa.NA
	assert v.as_py() is None

	def test_bool(self):
	arr = pa.array([True, None, False, None])

	v = arr[0]
	assert isinstance(v, pa.BooleanValue)
	assert repr(v) == "True"
	assert str(v) == "True"
	assert v.as_py() is True

	assert arr[1] is pa.NA

	def test_int64(self):
	arr = pa.array([1, 2, None])

	v = arr[0]
	assert isinstance(v, pa.Int64Value)
	assert repr(v) == "1"
	assert str(v) == "1"
	assert v.as_py() == 1
	assert v == 1

	assert arr[2] is pa.NA

	def test_double(self):
	arr = pa.array([1.5, None, 3])

	v = arr[0]
	assert isinstance(v, pa.DoubleValue)
	assert repr(v) == "1.5"
	assert str(v) == "1.5"
	assert v.as_py() == 1.5
	assert v == 1.5

	assert arr[1] is pa.NA

	v = arr[2]
	assert v.as_py() == 3.0

	def test_half_float(self):
	arr = pa.array([np.float16(1.5), None], type=pa.float16())
	v = arr[0]
	assert isinstance(v, pa.HalfFloatValue)
	assert repr(v) == "1.5"
	assert str(v) == "1.5"
	assert v.as_py() == 1.5
	assert v == 1.5

	assert arr[1] is pa.NA

	def test_string_unicode(self):
	arr = pa.array([u'foo', None, u'mañana'])

	v = arr[0]
	assert isinstance(v, pa.StringValue)
	assert v.as_py() == u'foo'
	assert repr(v) == repr(u"foo")
	assert str(v) == str(u"foo")
	assert v == u'foo'
	# Assert that newly created values are equal to the previously created
	# one.
	assert v == arr[0]

	assert arr[1] is pa.NA

	v = arr[2].as_py()
	assert v == u'mañana'
	assert isinstance(v, unicode_type)

	def test_bytes(self):
	arr = pa.array([b'foo', None, u('bar')])

	def check_value(v, expected):
	assert isinstance(v, pa.BinaryValue)
	assert v.as_py() == expected
	assert str(v) == str(expected)
	assert repr(v) == repr(expected)
	assert v == expected
	assert v != b'xxxxx'
	buf = v.as_buffer()
	assert isinstance(buf, pa.Buffer)
	assert buf.to_pybytes() == expected

	check_value(arr[0], b'foo')
	assert arr[1] is pa.NA
	check_value(arr[2], b'bar')

	def test_fixed_size_bytes(self):
	data = [b'foof', None, b'barb']
	arr = pa.array(data, type=pa.binary(4))

	v = arr[0]
	assert isinstance(v, pa.FixedSizeBinaryValue)
	assert v.as_py() == b'foof'

	assert arr[1] is pa.NA

	v = arr[2].as_py()
	assert v == b'barb'
	assert isinstance(v, bytes)

	def test_list(self):
	arr = pa.array([['foo', None], None, ['bar'], []])

	v = arr[0]
	assert len(v) == 2
	assert isinstance(v, pa.ListValue)
	assert repr(v) == "['foo', None]"
	assert v.as_py() == ['foo', None]
	assert v[0].as_py() == 'foo'
	assert v[1] is pa.NA
	assert v[-1] == v[1]
	assert v[-2] == v[0]
	with pytest.raises(IndexError):
	v[-3]
	with pytest.raises(IndexError):
	v[2]

	assert arr[1] is pa.NA

	v = arr[3]
	assert len(v) == 0

	@pytest.mark.pandas
	def test_timestamp(self):
	import pandas as pd
	arr = pd.date_range('2000-01-01 12:34:56', periods=10).values

	units = ['ns', 'us', 'ms', 's']

	for i, unit in enumerate(units):
	dtype = 'datetime64[{0}]'.format(unit)
	arrow_arr = pa.Array.from_pandas(arr.astype(dtype))
	expected = pd.Timestamp('2000-01-01 12:34:56')

	assert arrow_arr[0].as_py() == expected
	assert arrow_arr[0].value * 1000**i == expected.value

	tz = 'America/New_York'
	arrow_type = pa.timestamp(unit, tz=tz)

	dtype = 'datetime64[{0}]'.format(unit)
	arrow_arr = pa.Array.from_pandas(arr.astype(dtype),
	type=arrow_type)
	expected = (pd.Timestamp('2000-01-01 12:34:56')
	.tz_localize('utc')
	.tz_convert(tz))

	assert arrow_arr[0].as_py() == expected
	assert arrow_arr[0].value * 1000**i == expected.value

	@pytest.mark.pandas
	def test_dictionary(self):
	import pandas as pd
	colors = ['red', 'green', 'blue']
	colors_dict = {'red': 0, 'green': 1, 'blue': 2}
	values = pd.Series(colors * 4)

	categorical = pd.Categorical(values, categories=colors)

	v = pa.DictionaryArray.from_arrays(categorical.codes,
	categorical.categories)
	for i, c in enumerate(values):
	assert v[i].as_py() == c
	assert v[i].dictionary_value == c
	assert v[i].index_value == colors_dict[c]

	def test_int_hash(self):
	# ARROW-640
	int_arr = pa.array([1, 1, 2, 1])
	assert hash(int_arr[0]) == hash(1)

	def test_float_hash(self):
	# ARROW-640
	float_arr = pa.array([1.4, 1.2, 2.5, 1.8])
	assert hash(float_arr[0]) == hash(1.4)

	def test_string_hash(self):
	# ARROW-640
	str_arr = pa.array(["foo", "bar"])
	assert hash(str_arr[1]) == hash("bar")

	def test_bytes_hash(self):
	# ARROW-640
	byte_arr = pa.array([b'foo', None, b'bar'])
	assert hash(byte_arr[2]) == hash(b"bar")

	def test_array_to_set(self):
	# ARROW-640
	arr = pa.array([1, 1, 2, 1])
	set_from_array = set(arr)
	assert isinstance(set_from_array, set)
	assert set_from_array == {1, 2}

	def test_struct_value_subscripting(self):
	ty = pa.struct([pa.field('x', pa.int16()),
	pa.field('y', pa.float32())])
	arr = pa.array([(1, 2.5), (3, 4.5), (5, 6.5)], type=ty)

	assert arr[0]['x'] == 1
	assert arr[0]['y'] == 2.5
	assert arr[1]['x'] == 3
	assert arr[1]['y'] == 4.5
	assert arr[2]['x'] == 5
	assert arr[2]['y'] == 6.5

	with pytest.raises(IndexError):
	arr[4]['non-existent']

	with pytest.raises(KeyError):
	arr[0]['non-existent']