()
| 3707 | |
| 3708 | |
| 3709 | def create_sample_expressions(): |
| 3710 | # We need a schema for substrait conversion |
| 3711 | schema = pa.schema([pa.field("i64", pa.int64()), pa.field( |
| 3712 | "foo", pa.struct([pa.field("bar", pa.string())]))]) |
| 3713 | |
| 3714 | # Creates a bunch of sample expressions for testing |
| 3715 | # serialization and deserialization. The expressions are categorized |
| 3716 | # to reflect certain nuances in Substrait conversion. |
| 3717 | a = pc.scalar(1) |
| 3718 | b = pc.scalar(1.1) |
| 3719 | c = pc.scalar(True) |
| 3720 | d = pc.scalar("string") |
| 3721 | e = pc.scalar(None) |
| 3722 | f = pc.scalar({'a': 1}) |
| 3723 | g = pc.scalar(pa.scalar(1)) |
| 3724 | h = pc.scalar(np.int64(2)) |
| 3725 | j = pc.scalar(False) |
| 3726 | k = pc.scalar(0) |
| 3727 | |
| 3728 | # These expression consist entirely of literals |
| 3729 | literal_exprs = [a, b, c, d, e, g, h, j, k] |
| 3730 | |
| 3731 | # These expressions include at least one function call |
| 3732 | exprs_with_call = [a == b, a != b, a > b, c & j, c | j, ~c, d.is_valid(), |
| 3733 | a + b, a - b, a * b, a / b, pc.negate(a), |
| 3734 | pc.add(a, b), pc.subtract(a, b), pc.divide(a, b), |
| 3735 | pc.multiply(a, b), pc.power(a, a), pc.sqrt(a), |
| 3736 | pc.exp(b), pc.cos(b), pc.sin(b), pc.tan(b), |
| 3737 | pc.acos(b), pc.atan(b), pc.asin(b), pc.atan2(b, b), |
| 3738 | pc.sinh(a), pc.cosh(a), pc.tanh(a), |
| 3739 | pc.asinh(a), pc.acosh(b), pc.atanh(k), |
| 3740 | pc.abs(b), pc.sign(a), pc.bit_wise_not(a), |
| 3741 | pc.bit_wise_and(a, a), pc.bit_wise_or(a, a), |
| 3742 | pc.bit_wise_xor(a, a), pc.is_nan(b), pc.is_finite(b), |
| 3743 | pc.coalesce(a, b), |
| 3744 | a.cast(pa.int32(), safe=False)] |
| 3745 | |
| 3746 | # These expressions test out various reference styles and may include function |
| 3747 | # calls. Named references are used here. |
| 3748 | exprs_with_ref = [pc.field('i64') > 5, pc.field('i64') == 5, |
| 3749 | pc.field('i64') == 7, |
| 3750 | pc.field(('foo', 'bar')) == 'value', |
| 3751 | pc.field('foo', 'bar') == 'value'] |
| 3752 | |
| 3753 | # Similar to above but these use numeric references instead of string refs |
| 3754 | exprs_with_numeric_refs = [pc.field(0) > 5, pc.field(0) == 5, |
| 3755 | pc.field(0) == 7, |
| 3756 | pc.field((1, 0)) == 'value', |
| 3757 | pc.field(1, 0) == 'value'] |
| 3758 | |
| 3759 | # Expressions that behave uniquely when converting to/from substrait |
| 3760 | special_cases = [ |
| 3761 | f, # Struct literals lose their field names |
| 3762 | a.isin([1, 2, 3]), # isin converts to an or list |
| 3763 | pc.field('i64').is_null() # pyarrow always specifies a FunctionOptions |
| 3764 | # for is_null which, being the default, is |
| 3765 | # dropped on serialization |
| 3766 | ] |
no test coverage detected