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