11
11
Any ,
12
12
Callable ,
13
13
cast ,
14
- TYPE_CHECKING ,
15
14
)
16
15
17
16
from databricks .sql .backend .sea .backend import SeaDatabricksClient
18
17
from databricks .sql .backend .types import ExecuteResponse
19
18
20
- if TYPE_CHECKING :
21
- from databricks .sql .result_set import ResultSet , SeaResultSet
19
+ from databricks .sql .result_set import ResultSet , SeaResultSet
22
20
23
21
logger = logging .getLogger (__name__ )
24
22
@@ -30,8 +28,8 @@ class ResultSetFilter:
30
28
31
29
@staticmethod
32
30
def _filter_sea_result_set (
33
- result_set : " SeaResultSet" , filter_func : Callable [[List [Any ]], bool ]
34
- ) -> " SeaResultSet" :
31
+ result_set : SeaResultSet , filter_func : Callable [[List [Any ]], bool ]
32
+ ) -> SeaResultSet :
35
33
"""
36
34
Filter a SEA result set using the provided filter function.
37
35
@@ -49,9 +47,6 @@ def _filter_sea_result_set(
49
47
# Filter rows
50
48
filtered_rows = [row for row in all_rows if filter_func (row )]
51
49
52
- # Import SeaResultSet here to avoid circular imports
53
- from databricks .sql .result_set import SeaResultSet
54
-
55
50
# Reuse the command_id from the original result set
56
51
command_id = result_set .command_id
57
52
@@ -67,10 +62,13 @@ def _filter_sea_result_set(
67
62
)
68
63
69
64
# Create a new ResultData object with filtered data
65
+
70
66
from databricks .sql .backend .sea .models .base import ResultData
71
67
72
68
result_data = ResultData (data = filtered_rows , external_links = None )
73
69
70
+ from databricks .sql .result_set import SeaResultSet
71
+
74
72
# Create a new SeaResultSet with the filtered data
75
73
filtered_result_set = SeaResultSet (
76
74
connection = result_set .connection ,
@@ -85,11 +83,11 @@ def _filter_sea_result_set(
85
83
86
84
@staticmethod
87
85
def filter_by_column_values (
88
- result_set : " ResultSet" ,
86
+ result_set : ResultSet ,
89
87
column_index : int ,
90
88
allowed_values : List [str ],
91
89
case_sensitive : bool = False ,
92
- ) -> " ResultSet" :
90
+ ) -> ResultSet :
93
91
"""
94
92
Filter a result set by values in a specific column.
95
93
@@ -133,8 +131,8 @@ def filter_by_column_values(
133
131
134
132
@staticmethod
135
133
def filter_tables_by_type (
136
- result_set : " ResultSet" , table_types : Optional [List [str ]] = None
137
- ) -> " ResultSet" :
134
+ result_set : ResultSet , table_types : Optional [List [str ]] = None
135
+ ) -> ResultSet :
138
136
"""
139
137
Filter a result set of tables by the specified table types.
140
138
0 commit comments