9
9
List ,
10
10
Optional ,
11
11
Any ,
12
- Dict ,
13
12
Callable ,
14
- TypeVar ,
15
- Generic ,
16
13
cast ,
17
- TYPE_CHECKING ,
18
14
)
19
15
20
16
from databricks .sql .backend .types import ExecuteResponse , CommandId
29
25
30
26
class ResultSetFilter :
31
27
"""
32
- A general-purpose filter for result sets that can be applied to any backend.
33
-
34
- This class provides methods to filter result sets based on various criteria,
35
- similar to the client-side filtering in the JDBC connector.
28
+ A general-purpose filter for result sets.
36
29
"""
37
30
38
31
@staticmethod
39
32
def _filter_sea_result_set (
40
- result_set : " SeaResultSet" , filter_func : Callable [[List [Any ]], bool ]
41
- ) -> " SeaResultSet" :
33
+ result_set : SeaResultSet , filter_func : Callable [[List [Any ]], bool ]
34
+ ) -> SeaResultSet :
42
35
"""
43
36
Filter a SEA result set using the provided filter function.
44
37
@@ -49,15 +42,13 @@ def _filter_sea_result_set(
49
42
Returns:
50
43
A filtered SEA result set
51
44
"""
45
+
52
46
# Get all remaining rows
53
47
all_rows = result_set .results .remaining_rows ()
54
48
55
49
# Filter rows
56
50
filtered_rows = [row for row in all_rows if filter_func (row )]
57
51
58
- # Import SeaResultSet here to avoid circular imports
59
- from databricks .sql .result_set import SeaResultSet
60
-
61
52
# Reuse the command_id from the original result set
62
53
command_id = result_set .command_id
63
54
@@ -73,10 +64,13 @@ def _filter_sea_result_set(
73
64
)
74
65
75
66
# Create a new ResultData object with filtered data
67
+
76
68
from databricks .sql .backend .sea .models .base import ResultData
77
69
78
70
result_data = ResultData (data = filtered_rows , external_links = None )
79
71
72
+ from databricks .sql .result_set import SeaResultSet
73
+
80
74
# Create a new SeaResultSet with the filtered data
81
75
from databricks .sql .backend .sea .backend import SeaDatabricksClient
82
76
@@ -93,11 +87,11 @@ def _filter_sea_result_set(
93
87
94
88
@staticmethod
95
89
def filter_by_column_values (
96
- result_set : " ResultSet" ,
90
+ result_set : ResultSet ,
97
91
column_index : int ,
98
92
allowed_values : List [str ],
99
93
case_sensitive : bool = False ,
100
- ) -> " ResultSet" :
94
+ ) -> ResultSet :
101
95
"""
102
96
Filter a result set by values in a specific column.
103
97
@@ -110,6 +104,7 @@ def filter_by_column_values(
110
104
Returns:
111
105
A filtered result set
112
106
"""
107
+
113
108
# Convert to uppercase for case-insensitive comparison if needed
114
109
if not case_sensitive :
115
110
allowed_values = [v .upper () for v in allowed_values ]
@@ -140,8 +135,8 @@ def filter_by_column_values(
140
135
141
136
@staticmethod
142
137
def filter_tables_by_type (
143
- result_set : " ResultSet" , table_types : Optional [List [str ]] = None
144
- ) -> " ResultSet" :
138
+ result_set : ResultSet , table_types : Optional [List [str ]] = None
139
+ ) -> ResultSet :
145
140
"""
146
141
Filter a result set of tables by the specified table types.
147
142
@@ -156,6 +151,7 @@ def filter_tables_by_type(
156
151
Returns:
157
152
A filtered result set containing only tables of the specified types
158
153
"""
154
+
159
155
# Default table types if none specified
160
156
DEFAULT_TABLE_TYPES = ["TABLE" , "VIEW" , "SYSTEM TABLE" ]
161
157
valid_types = (
0 commit comments