def _check_schema(self, rows, expected_schema): from oarphpy.spark import RowAdapter schema = RowAdapter.to_schema(rows[0]) adapted_rows = [RowAdapter.to_row(r) for r in rows] with testutil.LocalSpark.sess() as spark: df = spark.createDataFrame(adapted_rows, schema=schema, verifySchema=False) # verifySchema is expensive and improperly errors on mostly empty rows if self._is_spark_2x(): # Spark 2.x returns schema values in a different order, so we do a more # flexible test def tokenize(s): import re return sorted(re.split('[<>,]+', s)) actual = dict((col, tokenize(s)) for col, s in df.dtypes) expected = dict( (col, tokenize(s)) for col, s in expected_schema) assert actual == expected else: # Tests are written for Spark 3.x assert df.dtypes == expected_schema return schema
def test_row_adapter(): import numpy as np from pyspark.sql import Row from oarphpy.spark import RowAdapter rows = [ Row( id=1, np_number=np.float32(1.), a=np.array([1]), b={'foo': np.array([[1]], dtype=np.uint8)}, c=[np.array([[[1.]], [[2.]], [[3.]]])], d=Slotted(foo=5, bar="abc", _not_hidden=1), e=[Slotted(foo=6, bar="def", _not_hidden=1)], f=Unslotted(meow=4, _not_hidden=1, __hidden=2), g=Unslotted(), # Intentionally empty; adapter should set nothing h=Row(i=1, j=2), ), # Include a mostly empty row below to exercise Spark type validation. # Spark will ensure the row below and row above have the same schema; # note that `None` (or 'null') is only allowed for Struct / Row types. Row( id=2, np_number=np.float32(2.), a=np.array([]), b={}, c=[], d=None, e=[], f=None, g=None, h=Row(i=3, j=3), ), ] with testutil.LocalSpark.sess() as spark: ## Test basic round-trip serialization and adaptation _check_serialization(spark, rows, 'basic') ## Test Schema Deduction mostly_empty = Row( id=2, np_number=None, a=None, b={}, c=[], d=None, e=[], f=None, g=None, h=None, ) mostly_empty_adapted = RowAdapter.to_row(mostly_empty) # Spark can't deduce schema from the empty-ish row ... with pytest.raises(ValueError) as excinfo: df = spark.createDataFrame([mostly_empty_adapted]) assert "Some of types cannot be determined" in str(excinfo.value) # ... but this works if we tell it the schema! schema = RowAdapter.to_schema(rows[0]) df = spark.createDataFrame([mostly_empty_adapted], schema=schema, verifySchema=False) EXPECTED_SCHEMA = [ ('a', 'struct<__pyclass__:string,dtype:string,order:string,shape:array<bigint>,values:array<bigint>>' ), ('b', 'map<string,struct<__pyclass__:string,dtype:string,order:string,shape:array<bigint>,values:array<bigint>>>' ), ('c', 'array<struct<__pyclass__:string,dtype:string,order:string,shape:array<bigint>,values:array<double>>>' ), ('d', 'struct<__pyclass__:string,_not_hidden:bigint,bar:string,foo:bigint>' ), ('e', 'array<struct<__pyclass__:string,_not_hidden:bigint,bar:string,foo:bigint>>' ), ('f', 'struct<__pyclass__:string,_not_hidden:bigint,meow:bigint>'), ('g', 'struct<__pyclass__:string>'), ('h', 'struct<i:bigint,j:bigint>'), ('id', 'bigint'), ('np_number', 'double'), ] assert df.dtypes == EXPECTED_SCHEMA # Check that pyspark retains the empty values in `mostly_empty` for colname in sorted(df.columns): values = df.select(colname).collect() assert len(values) == 1 assert mostly_empty[colname] == values[0][colname] # ... and we can also read/write the empty-ish row! _check_serialization(spark, [mostly_empty], 'with_schema', schema=schema)
def test_rowadapter_complex(self): from oarphpy.spark import RowAdapter # A large-ish example that covers the above cases in aggregate rows = [ Row( id=1, np_number=np.float32(1.), a=np.array([1]), b={'foo': np.array([[1]], dtype=np.uint8)}, c=[np.array([[[1.]], [[2.]], [[3.]]])], d=Slotted(foo=5, bar="abc", _not_hidden=1), e=[Slotted(foo=6, bar="def", _not_hidden=1)], f=Unslotted(meow=4, _not_hidden=1, __hidden=2), g=Unslotted( ), # Intentionally empty; adapter should set nothing h=Row(i=1, j=2), ), # Include a mostly empty row below to exercise Spark type validation. # Spark will ensure the row below and row above have the same schema; # note that `None` (or 'null') is only allowed for Struct / Row types. Row( id=2, np_number=np.float32(2.), a=np.array([]), b={}, c=[], d=None, e=[], f=None, g=None, h=Row(i=3, j=3), ), ] df = self._check_serialization(rows) EXPECTED_ALL = """ 0 1 id 1 2 np_number 1.0 2.0 a (oarphpy.spark.Tensor, [1], int64, C, [1], []) (oarphpy.spark.Tensor, [0], float64, C, [], []) b {'foo': ('oarphpy.spark.Tensor', [1, 1], 'uint8', 'C', [1], [])} {} c [(oarphpy.spark.Tensor, [3, 1, 1], float64, C, [1.0, 2.0, 3.0], [])] [] d (oarphpy_test.test_spark.Slotted, 5, abc, 1) None e [(oarphpy_test.test_spark.Slotted, 6, def, 1)] [] f (oarphpy_test.test_spark.Unslotted, 4, 1) None g (oarphpy_test.test_spark.Unslotted,) None h (1, 2) (3, 3) """ # DEPRECATED: pyspark 2.x is deprecated import pyspark if pyspark.__version__.startswith('2.'): EXPECTED_ALL = """ 0 1 id 1 2 np_number 1 2 a (oarphpy.spark.Tensor, [1], int64, C, [1], []) (oarphpy.spark.Tensor, [0], float64, C, [], []) b {'foo': ('oarphpy.spark.Tensor', [1, 1], 'uint8', 'C', [1], [])} {} c [(oarphpy.spark.Tensor, [3, 1, 1], float64, C, [1.0, 2.0, 3.0], [])] [] d (oarphpy_test.test_spark.Slotted, 5, abc, 1) None e [(oarphpy_test.test_spark.Slotted, 6, def, 1)] [] f (oarphpy_test.test_spark.Unslotted, 4, 1) None g (oarphpy_test.test_spark.Unslotted,) None h (1, 2) (3, 3) """ self._pandas_compare_str(df.orderBy('id').toPandas().T, EXPECTED_ALL) # Test Schema Deduction mostly_empty = Row( id=2, np_number=None, a=None, b={}, c=[], d=None, e=[], f=None, g=None, h=None, ) mostly_empty_adapted = RowAdapter.to_row(mostly_empty) # Spark can't deduce schema from the empty-ish row ... with pytest.raises(ValueError) as excinfo: self._check_serialization([mostly_empty_adapted], do_adaption=False) assert "Some of types cannot be determined" in str(excinfo.value) # ... but this works if we tell it the schema! schema = RowAdapter.to_schema(rows[0]) self._check_serialization([mostly_empty_adapted], schema=schema) # Let's check that RowAdapter schema deduction works as expected EXPECTED_SCHEMA = [ ('id', 'bigint'), ('np_number', 'double'), ('a', 'struct<__pyclass__:string,shape:array<bigint>,dtype:string,order:string,values:array<bigint>,values_packed:binary>' ), ('b', 'map<string,struct<__pyclass__:string,shape:array<bigint>,dtype:string,order:string,values:array<bigint>,values_packed:binary>>' ), ('c', 'array<struct<__pyclass__:string,shape:array<bigint>,dtype:string,order:string,values:array<double>,values_packed:binary>>' ), ('d', 'struct<__pyclass__:string,foo:bigint,bar:string,_not_hidden:bigint>' ), ('e', 'array<struct<__pyclass__:string,foo:bigint,bar:string,_not_hidden:bigint>>' ), ('f', 'struct<__pyclass__:string,meow:bigint,_not_hidden:bigint>'), ('g', 'struct<__pyclass__:string>'), ('h', 'struct<i:bigint,j:bigint>'), ] self._check_schema(rows, EXPECTED_SCHEMA)