def test_simple_methods(): df = ArrayDataFrame([["a", 1], ["b", "2"]], "x:str,y:double") assert 2 == df.count() assert not df.empty assert ["a", 1.0] == df.peek_array() assert dict(x="a", y=1.0) == df.peek_dict() assert [["a", 1], ["b", "2"]] == df.as_array() df = ArrayDataFrame([["a", 1], ["b", "2"]], "x:str,y:double") pdf = df.as_pandas() assert [["a", 1.0], ["b", 2.0]] == pdf.values.tolist() df = ArrayDataFrame([], "x:str,y:double") pdf = df.as_pandas() assert [] == pdf.values.tolist()
def test_transformer(): assert isinstance(t1, CoTransformer) df1 = ArrayDataFrame([[0, 2]], "a:int,b:int") df2 = ArrayDataFrame([[0, 2]], "a:int,c:int") dfs = DataFrames(df1, df2) t1._output_schema = t1.get_output_schema(dfs) assert t1.output_schema == "a:int,b:int" t2._output_schema = t2.get_output_schema(dfs) assert t2.output_schema == "b:int,a:int" assert [[0, 2, 1]] == list(t3(df1.as_array(), df2.as_pandas()))