-
Notifications
You must be signed in to change notification settings - Fork 233
/
test_create_dataframe.py
46 lines (32 loc) · 1.75 KB
/
test_create_dataframe.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
from optimus.helpers.functions import results_equal
from optimus.tests.base import TestBase
class TestDFPandas(TestBase):
def test_dataframe_dict_auto(self):
data = {'object': [1, "2", 3.0], "int": [1, 2, 3], "float": [1.0, 2.0, 3.5]}
df = self.op.create.dataframe(data)
result = df.cols.data_type()
expected = {'object': 'object', 'int': 'int64', 'float': 'float64'}
self.assertTrue(results_equal(result, expected))
def test_dataframe_dict_forced(self):
data = {('object', 'int64'): [1, "2", 3.0], ("int", 'float'): [1, 2, 3], ("float", 'object'): [1.0, 2.0, 3.5]}
df = self.op.create.dataframe(data, force_data_types=True)
result = df.cols.data_type()
expected = {'object': 'int64', 'int': 'float64', 'float': 'object'}
self.assertTrue(results_equal(result, expected))
def test_dataframe_pandas(self):
import pandas as pd
data = pd.util.testing.makeMixedDataFrame()
df = self.op.create.dataframe(data)
result = df.cols.data_type()
expected = {'A': 'float64', 'B': 'float64', 'C': 'object', 'D': 'datetime64[ns]'}
self.assertTrue(results_equal(result, expected))
def test_dataframe_dict_forced(self):
data = {('object', 'int64'): [1, "2", 3.0], ("int", 'float'): [1, 2, 3], ("float", 'object'): [1.0, 2.0, 3.5]}
df = self.op.create.dataframe(data, force_data_types=True)
result = df.cols.data_type()
expected = {'object': 'int64', 'int': 'float64', 'float': 'object'}
self.assertTrue(results_equal(result, expected))
class TestDFDask(TestDFPandas):
config = {'engine': 'dask', 'n_partitions': 1}
class TestDFPartitionDask(TestDFPandas):
config = {'engine': 'dask', 'n_partitions': 2}