def test_rowadapter_cloudpickled_callable(self):
from oarphpy.spark import CloudpickeledCallable
def moof():
return 'moof'
cc_moof = CloudpickeledCallable(moof)
assert cc_moof() == moof()
cc_empty = CloudpickeledCallable()
assert cc_empty == CloudpickeledCallable.empty()
with pytest.raises(Exception):
cc_empty()
rows = [
Row(id=0, f=cc_empty),
Row(id=1, f=cc_moof),
]
df = self._check_serialization(rows)
assert _select_distinct(df, 'f.__pyclass__') == [
'oarphpy.spark.CloudpickeledCallableData'
]
# Ensuring invoking the functions still works
from oarphpy.spark import RowAdapter
decoded = [RowAdapter.from_row(r) for r in df.orderBy('id').collect()]
with pytest.raises(Exception):
decoded[0].f()
assert decoded[1].f() == moof()
def _check_serialization(spark, rows, testname, schema=None):
from oarphpy import util
from oarphpy.spark import RowAdapter
TEST_TEMPDIR = testutil.test_tempdir('spark_row_adapter_test')
adapted_rows = [RowAdapter.to_row(r) for r in rows]
if schema:
df = spark.createDataFrame(adapted_rows,
schema=schema,
verifySchema=False)
# verifySchema is expensive and improperly erros on mostly empty rows
else:
df = spark.createDataFrame(adapted_rows)
# Automatically samples both rows to get schema
outpath = os.path.join(TEST_TEMPDIR, 'rowdata_%s' % testname)
df.write.parquet(outpath)
df2 = spark.read.parquet(outpath)
decoded_wrapped_rows = df2.collect()
decoded_rows = [RowAdapter.from_row(row) for row in decoded_wrapped_rows]
# We can't do assert sorted(rows) == sorted(decoded_rows)
# because numpy syntatic sugar breaks ==
import pprint
def sorted_row_str(rowz):
return pprint.pformat(sorted(rowz, key=lambda row: row['id']))
assert sorted_row_str(rows) == sorted_row_str(decoded_rows)
def test_rowadapter_numpy_unpacked(self):
import numpy as np
# RowAdapter translates Numpy arrays to a oarphy Tensor object that affords
# SQL-based inspection for small arrays (and uses a more efficient row- or
# column-major packed encoding for large arrays; see next test)
rows = [
Row(id=0, x=np.array([1, 2, 3])),
Row(id=1, x=np.array([4, 5, 6])),
]
df = self._check_serialization(rows)
assert _select_distinct(df,
'x.__pyclass__') == ['oarphpy.spark.Tensor']
EXPECTED = """
id x
0 0 (oarphpy.spark.Tensor, [3], int64, C, [1, 2, 3], [])
1 1 (oarphpy.spark.Tensor, [3], int64, C, [4, 5, 6], [])
"""
self._pandas_compare_str(df.orderBy('id').toPandas(), EXPECTED)
# _check_serialization() verifies that the data gets decoded as numpy
# arrays, but just to make things visible in this test:
from oarphpy.spark import RowAdapter
decoded = [RowAdapter.from_row(r) for r in df.orderBy('id').collect()]
np.testing.assert_equal(decoded[0].x, np.array([1, 2, 3]))
np.testing.assert_equal(decoded[1].x, np.array([4, 5, 6]))
def test_rowadapter_numpy_packed(self):
import numpy as np
from oarphpy.spark import TENSOR_AUTO_PACK_MIN_KBYTES
N = int(TENSOR_AUTO_PACK_MIN_KBYTES *
(2**10) / np.dtype(int).itemsize) + 1
expect_packed = np.reshape(np.array(range(2 * N)), (2, N))
rows = [
Row(id=0, x=expect_packed),
Row(id=1, x=(expect_packed + 1)),
]
df = self._check_serialization(rows)
assert _select_distinct(df,
'x.__pyclass__') == ['oarphpy.spark.Tensor']
# Verify that we actually have a column of packed values
bin_data = df.select('*').first().x.values_packed
assert len(
bin_data) == expect_packed.size * expect_packed.dtype.itemsize
# For ints, usually 8 bytes per int * 2 * N
# _check_serialization() verifies that the data gets decoded as numpy
# arrays, but just to make things visible in this test:
from oarphpy.spark import RowAdapter
decoded = [RowAdapter.from_row(r) for r in df.orderBy('id').collect()]
np.testing.assert_equal(decoded[0].x, expect_packed)
np.testing.assert_equal(decoded[1].x, expect_packed + 1)
def test_rowadapter_unslotted(self):
# RowAdapter will deserialize and re-create Python objects, using the
# `Unslotted` class defined at runtime.
rows = [
Row(id=0, x=Unslotted(v=4)),
Row(id=1, x=Unslotted(v=5)),
]
df = self._check_serialization(rows)
# NB: The check above also checks that `Unslotted` instances are created
# and checks equality via __eq__ or pprint.pformat(). But just to make
# things visible in this test:
from oarphpy.spark import RowAdapter
decoded = sorted(RowAdapter.from_row(r) for r in df.collect())
assert [r.x for r in decoded] == [Unslotted(v=4), Unslotted(v=5)]
# RowAdapter records the (full) class name in the table as the
# `__pyclass__` attribute of each value of the `x` column.
assert _select_distinct(
df, 'x.__pyclass__') == ['oarphpy_test.test_spark.Unslotted']
# RowAdapter encodes objects as structs (even though in Python objects are
# very dict-like).
self._check_schema(rows, [
('id', 'bigint'),
('x', 'struct<__pyclass__:string,v:bigint>'),
])
def test_nonadapted_input(self):
from oarphpy.spark import RowAdapter
# RowAdapter leaves bare data input unchanged
BARE_VALUES = True, 1, 1.0, "moof", bytes(b"moof")
for datum in BARE_VALUES:
assert RowAdapter.to_row(datum) == datum
assert RowAdapter.from_row(datum) == datum
def _check_raw_adaption(self, raw_expected):
from oarphpy.spark import RowAdapter
for raw_data, expected_row in raw_expected:
actual_row = RowAdapter.to_row(raw_data)
assert actual_row == expected_row
actual_data = RowAdapter.from_row(expected_row)
assert actual_data == raw_data
def test_pesky_numpy(self):
import numpy as np
from oarphpy.spark import RowAdapter
# RowAdapter translates pesky numpy-boxed numbers ...
assert RowAdapter.to_row(Row(x=np.float32(1.))) == Row(x=1.)
# ... but only one way! In practice, just don't save boxed numbers in rows.
assert RowAdapter.from_row(
Row(x=np.float32(1.))) == Row(x=np.float32(1.))
def _check_serialization(self, rows, schema=None, do_adaption=True):
import inspect
from oarphpy import util
from oarphpy.spark import RowAdapter
test_name = inspect.stack()[1][3]
TEST_TEMPDIR = testutil.test_tempdir('TestRowAdapter.' + test_name)
if do_adaption:
adapted_rows = [RowAdapter.to_row(r) for r in rows]
else:
adapted_rows = rows
with testutil.LocalSpark.sess() as spark:
if schema:
df = spark.createDataFrame(adapted_rows,
schema=schema,
verifySchema=False)
# verifySchema is expensive and improperly errors on mostly
# empty rows
else:
df = spark.createDataFrame(adapted_rows)
# Automatically samples rows to get schema
outpath = os.path.join(TEST_TEMPDIR, 'rowdata_%s' % test_name)
df.write.parquet(outpath)
df2 = spark.read.parquet(outpath)
decoded_wrapped_rows = df2.collect()
if do_adaption:
decoded_rows = [
RowAdapter.from_row(row) for row in decoded_wrapped_rows
]
# We can't do assert sorted(rows) == sorted(decoded_rows)
# because numpy syntatic sugar breaks __eq__, so use pprint,
# which is safe for our tests
import pprint
def sorted_row_str(rowz):
if self._is_spark_2x():
# Spark 2.x has non-stable sorting semantics for Row
if len(rowz) > 1:
rowz = sorted(rowz, key=lambda r: r.id)
return pprint.pformat(rowz)
else:
return pprint.pformat(sorted(rowz))
assert sorted_row_str(rows) == sorted_row_str(decoded_rows)
return df