def test_expression_ergonomics(): zero = ds.scalar(0) one = ds.scalar(1) true = ds.scalar(True) false = ds.scalar(False) string = ds.scalar("string") field = ds.field("field") assert one.equals(ds.ScalarExpression(1)) assert zero.equals(ds.ScalarExpression(0)) assert true.equals(ds.ScalarExpression(True)) assert false.equals(ds.ScalarExpression(False)) assert string.equals(ds.ScalarExpression("string")) assert field.equals(ds.FieldExpression("field")) expected = ds.AndExpression(ds.ScalarExpression(1), ds.ScalarExpression(0)) for expr in [one & zero, 1 & zero, one & 0]: assert expr.equals(expected) expected = ds.OrExpression(ds.ScalarExpression(1), ds.ScalarExpression(0)) for expr in [one | zero, 1 | zero, one | 0]: assert expr.equals(expected) comparison_ops = [ (operator.eq, ds.CompareOperator.Equal), (operator.ne, ds.CompareOperator.NotEqual), (operator.ge, ds.CompareOperator.GreaterEqual), (operator.le, ds.CompareOperator.LessEqual), (operator.lt, ds.CompareOperator.Less), (operator.gt, ds.CompareOperator.Greater), ] for op, compare_op in comparison_ops: expr = op(zero, one) expected = ds.ComparisonExpression(compare_op, zero, one) assert expr.equals(expected) expr = ~true == false expected = ds.ComparisonExpression( ds.CompareOperator.Equal, ds.NotExpression(ds.ScalarExpression(True)), ds.ScalarExpression(False) ) assert expr.equals(expected) for typ in ("bool", pa.bool_()): expr = field.cast(typ) == true expected = ds.ComparisonExpression( ds.CompareOperator.Equal, ds.CastExpression(ds.FieldExpression("field"), pa.bool_()), ds.ScalarExpression(True) ) assert expr.equals(expected) expr = field.isin([1, 2]) expected = ds.InExpression(ds.FieldExpression("field"), pa.array([1, 2])) assert expr.equals(expected) with pytest.raises(TypeError): field.isin(1)
def test_expression(): a = ds.ScalarExpression(1) b = ds.ScalarExpression(1.1) c = ds.ScalarExpression(True) d = ds.ScalarExpression("string") e = ds.ScalarExpression(None) equal = ds.ComparisonExpression(ds.CompareOperator.Equal, a, b) greater = a > b assert equal.op == ds.CompareOperator.Equal and_ = ds.AndExpression(a, b) assert and_.left_operand.equals(a) assert and_.right_operand.equals(b) assert and_.equals(ds.AndExpression(a, b)) assert and_.equals(and_) or_ = ds.OrExpression(a, b) not_ = ds.NotExpression(ds.OrExpression(a, b)) is_valid = ds.IsValidExpression(a) cast_safe = ds.CastExpression(a, pa.int32()) cast_unsafe = ds.CastExpression(a, pa.int32(), safe=False) in_ = ds.InExpression(a, pa.array([1, 2, 3])) assert is_valid.operand == a assert in_.set_.equals(pa.array([1, 2, 3])) assert cast_unsafe.to == pa.int32() assert cast_unsafe.safe is False assert cast_safe.safe is True condition = ds.ComparisonExpression(ds.CompareOperator.Greater, ds.FieldExpression('i64'), ds.ScalarExpression(5)) schema = pa.schema( [pa.field('i64', pa.int64()), pa.field('f64', pa.float64())]) assert condition.validate(schema) == pa.bool_() i64_is_5 = ds.ComparisonExpression(ds.CompareOperator.Equal, ds.FieldExpression('i64'), ds.ScalarExpression(5)) i64_is_7 = ds.ComparisonExpression(ds.CompareOperator.Equal, ds.FieldExpression('i64'), ds.ScalarExpression(7)) assert condition.assume(i64_is_5).equals(ds.ScalarExpression(False)) assert condition.assume(i64_is_7).equals(ds.ScalarExpression(True)) assert str(condition) == "(i64 > 5:int64)" assert "(i64 > 5:int64)" in repr(condition) all_exprs = [ a, b, c, d, e, equal, greater, and_, or_, not_, is_valid, cast_unsafe, cast_safe, in_, condition, i64_is_5, i64_is_7 ] for expr in all_exprs: restored = pickle.loads(pickle.dumps(expr)) assert expr.equals(restored)
def test_expression(): a = ds.ScalarExpression(1) b = ds.ScalarExpression(1.1) c = ds.ScalarExpression(True) d = ds.ScalarExpression("string") equal = ds.ComparisonExpression(ds.CompareOperator.Equal, a, b) assert equal.op() == ds.CompareOperator.Equal and_ = ds.AndExpression(a, b) assert and_.left_operand.equals(a) assert and_.right_operand.equals(b) assert and_.equals(ds.AndExpression(a, b)) assert and_.equals(and_) ds.AndExpression(a, b, c) ds.OrExpression(a, b) ds.OrExpression(a, b, c, d) ds.NotExpression(ds.OrExpression(a, b, c)) ds.IsValidExpression(a) ds.CastExpression(a, pa.int32()) ds.CastExpression(a, pa.int32(), safe=True) ds.InExpression(a, pa.array([1, 2, 3])) condition = ds.ComparisonExpression( ds.CompareOperator.Greater, ds.FieldExpression('i64'), ds.ScalarExpression(5) ) schema = pa.schema([ pa.field('i64', pa.int64()), pa.field('f64', pa.float64()) ]) assert condition.validate(schema) == pa.bool_() i64_is_5 = ds.ComparisonExpression( ds.CompareOperator.Equal, ds.FieldExpression('i64'), ds.ScalarExpression(5) ) i64_is_7 = ds.ComparisonExpression( ds.CompareOperator.Equal, ds.FieldExpression('i64'), ds.ScalarExpression(7) ) assert condition.assume(i64_is_5).equals(ds.ScalarExpression(False)) assert condition.assume(i64_is_7).equals(ds.ScalarExpression(True)) assert str(condition) == "(i64 > 5:int64)" assert "(i64 > 5:int64)" in repr(condition)