def testApplyMap(self):
from odps.df.expr.collections import ProjectCollectionExpr, Column
from odps.df.expr.element import MappedExpr
schema = Schema.from_lists(['idx', 'f1', 'f2', 'f3'],
[types.int64] + [types.float64] * 3)
expr = CollectionExpr(_source_data=None, _schema=schema)
self.assertRaises(
ValueError, lambda: expr.applymap(
lambda v: v + 1, columns='idx', excludes='f1'))
mapped = expr.applymap(lambda v: v + 1)
self.assertIsInstance(mapped, ProjectCollectionExpr)
for c in mapped._fields:
self.assertIsInstance(c, MappedExpr)
mapped = expr.applymap(lambda v: v + 1, columns='f1')
self.assertIsInstance(mapped, ProjectCollectionExpr)
for c in mapped._fields:
self.assertIsInstance(c, MappedExpr if c.name == 'f1' else Column)
map_cols = set(['f1', 'f2', 'f3'])
mapped = expr.applymap(lambda v: v + 1, columns=map_cols)
self.assertIsInstance(mapped, ProjectCollectionExpr)
for c in mapped._fields:
self.assertIsInstance(c,
MappedExpr if c.name in map_cols else Column)
mapped = expr.applymap(lambda v: v + 1, excludes='idx')
self.assertIsInstance(mapped, ProjectCollectionExpr)
for c in mapped._fields:
self.assertIsInstance(c, Column if c.name == 'idx' else MappedExpr)
exc_cols = set(['idx', 'f1'])
mapped = expr.applymap(lambda v: v + 1, excludes=exc_cols)
self.assertIsInstance(mapped, ProjectCollectionExpr)
for c in mapped._fields:
self.assertIsInstance(c,
Column if c.name in exc_cols else MappedExpr)
