25
25
26
26
import pyarrow as pa
27
27
28
- from ._internal import (
29
- AggregateUDF ,
30
- Config ,
31
- DataFrame ,
28
+ from .context import (
32
29
SessionContext ,
33
30
SessionConfig ,
34
31
RuntimeConfig ,
35
- ScalarUDF ,
36
32
SQLOptions ,
37
33
)
38
34
35
+ # The following imports are okay to remain as opaque to the user.
36
+ from ._internal import Config
37
+
38
+ from .udf import ScalarUDF , AggregateUDF
39
+
39
40
from .common import (
40
41
DFSchema ,
41
42
)
42
43
44
+ from .dataframe import DataFrame
45
+
43
46
from .expr import (
44
- Alias ,
45
- Analyze ,
47
+ # Alias,
48
+ # Analyze,
46
49
Expr ,
47
- Filter ,
48
- Limit ,
49
- Like ,
50
- ILike ,
51
- Projection ,
52
- SimilarTo ,
53
- ScalarVariable ,
54
- Sort ,
55
- TableScan ,
56
- Not ,
57
- IsNotNull ,
58
- IsTrue ,
59
- IsFalse ,
60
- IsUnknown ,
61
- IsNotTrue ,
62
- IsNotFalse ,
63
- IsNotUnknown ,
64
- Negative ,
65
- InList ,
66
- Exists ,
67
- Subquery ,
68
- InSubquery ,
69
- ScalarSubquery ,
70
- GroupingSet ,
71
- Placeholder ,
72
- Case ,
73
- Cast ,
74
- TryCast ,
75
- Between ,
76
- Explain ,
77
- CreateMemoryTable ,
78
- SubqueryAlias ,
79
- Extension ,
80
- CreateView ,
81
- Distinct ,
82
- DropTable ,
83
- Repartition ,
84
- Partitioning ,
85
- Window ,
50
+ # Filter,
51
+ # Limit,
52
+ # Like,
53
+ # ILike,
54
+ # Projection,
55
+ # SimilarTo,
56
+ # ScalarVariable,
57
+ # Sort,
58
+ # TableScan,
59
+ # Not,
60
+ # IsNotNull,
61
+ # IsTrue,
62
+ # IsFalse,
63
+ # IsUnknown,
64
+ # IsNotTrue,
65
+ # IsNotFalse,
66
+ # IsNotUnknown,
67
+ # Negative,
68
+ # InList,
69
+ # Exists,
70
+ # Subquery,
71
+ # InSubquery,
72
+ # ScalarSubquery,
73
+ # GroupingSet,
74
+ # Placeholder,
75
+ # Case,
76
+ # Cast,
77
+ # TryCast,
78
+ # Between,
79
+ # Explain,
80
+ # CreateMemoryTable,
81
+ # SubqueryAlias,
82
+ # Extension,
83
+ # CreateView,
84
+ # Distinct,
85
+ # DropTable,
86
+ # Repartition,
87
+ # Partitioning,
88
+ # Window,
86
89
WindowFrame ,
87
90
)
88
91
96
99
"SQLOptions" ,
97
100
"RuntimeConfig" ,
98
101
"Expr" ,
99
- "AggregateUDF" ,
100
102
"ScalarUDF" ,
101
103
"Window" ,
102
104
"WindowFrame" ,
@@ -175,8 +177,6 @@ def column(value):
175
177
176
178
177
179
def literal (value ):
178
- if not isinstance (value , pa .Scalar ):
179
- value = pa .scalar (value )
180
180
return Expr .literal (value )
181
181
182
182
@@ -200,20 +200,20 @@ def udf(func, input_types, return_type, volatility, name=None):
200
200
)
201
201
202
202
203
- def udaf (accum , input_type , return_type , state_type , volatility , name = None ):
203
+ def udaf (accum , input_types , return_type , state_type , volatility , name = None ):
204
204
"""
205
205
Create a new User Defined Aggregate Function
206
206
"""
207
207
if not issubclass (accum , Accumulator ):
208
208
raise TypeError ("`accum` must implement the abstract base class Accumulator" )
209
209
if name is None :
210
210
name = accum .__qualname__ .lower ()
211
- if isinstance (input_type , pa .lib .DataType ):
212
- input_type = [input_type ]
211
+ if isinstance (input_types , pa .lib .DataType ):
212
+ input_types = [input_types ]
213
213
return AggregateUDF (
214
214
name = name ,
215
215
accumulator = accum ,
216
- input_type = input_type ,
216
+ input_types = input_types ,
217
217
return_type = return_type ,
218
218
state_type = state_type ,
219
219
volatility = volatility ,
0 commit comments