def test_cupy():
t1 = mt.ones((100, 50), chunk_size=50, gpu=True)
t2 = mt.ones(50, chunk_size=50, gpu=True)
t = (t1 - t2) / mt.sqrt(t2 * (1 - t2) * len(t2))
graph = TileableGraph([t.data])
next(TileableGraphBuilder(graph).build())
context = dict()
chunk_graph_builder = ChunkGraphBuilder(graph,
fuse_enabled=False,
tile_context=context)
chunk_graph = next(chunk_graph_builder.build())
CupyRuntimeOptimizer(chunk_graph).optimize()
assert any(n.op.__class__.__name__ == 'TensorCpFuseChunk'
for n in chunk_graph)
def test_read_csv_head(gen_data1):
pdf, tempdir = gen_data1
file_path = os.path.join(tempdir, 'test.csv')
pdf.to_csv(file_path)
df1 = md.read_csv(file_path)
df2 = df1.head(5)
graph = TileableGraph([df2.data])
next(TileableGraphBuilder(graph).build())
context = dict()
chunk_graph_builder = ChunkGraphBuilder(graph,
fuse_enabled=False,
tile_context=context)
chunk_graph = next(chunk_graph_builder.build())
chunk1 = context[df1.data].chunks[0].data
chunk2 = context[df2.data].chunks[0].data
records = optimize(chunk_graph)
assert records.get_optimization_result(chunk1) is None
opt_chunk2 = records.get_optimization_result(chunk2)
assert opt_chunk2.op.nrows == 5
assert len(chunk_graph) == 1
assert opt_chunk2 in chunk_graph.results
def test_groupby_read_csv(gen_data1):
pdf, tempdir = gen_data1
file_path = os.path.join(tempdir, 'test.csv')
pdf.to_csv(file_path)
df1 = md.read_csv(file_path)
df2 = df1[['a', 'b']]
graph = TileableGraph([df2.data])
next(TileableGraphBuilder(graph).build())
context = dict()
chunk_graph_builder = ChunkGraphBuilder(graph,
fuse_enabled=False,
tile_context=context)
chunk_graph = next(chunk_graph_builder.build())
chunk1 = context[df1.data].chunks[0].data
chunk2 = context[df2.data].chunks[0].data
records = optimize(chunk_graph)
opt_chunk1 = records.get_optimization_result(chunk1)
assert opt_chunk1 is None
opt_chunk2 = records.get_optimization_result(chunk2)
assert opt_chunk2 is not None
assert opt_chunk2.op.usecols == ['a', 'b']
# original tileable should not be modified
assert chunk2.inputs[0] is chunk1