1111 Any ,
1212 Callable ,
1313 cast ,
14- TYPE_CHECKING ,
1514)
1615
1716from databricks .sql .backend .sea .backend import SeaDatabricksClient
1817from databricks .sql .backend .types import ExecuteResponse
1918
20- if TYPE_CHECKING :
21- from databricks .sql .result_set import ResultSet , SeaResultSet
19+ from databricks .sql .result_set import ResultSet , SeaResultSet
2220
2321logger = logging .getLogger (__name__ )
2422
@@ -30,8 +28,8 @@ class ResultSetFilter:
3028
3129 @staticmethod
3230 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 :
3533 """
3634 Filter a SEA result set using the provided filter function.
3735
@@ -49,9 +47,6 @@ def _filter_sea_result_set(
4947 # Filter rows
5048 filtered_rows = [row for row in all_rows if filter_func (row )]
5149
52- # Import SeaResultSet here to avoid circular imports
53- from databricks .sql .result_set import SeaResultSet
54-
5550 # Reuse the command_id from the original result set
5651 command_id = result_set .command_id
5752
@@ -67,10 +62,13 @@ def _filter_sea_result_set(
6762 )
6863
6964 # Create a new ResultData object with filtered data
65+
7066 from databricks .sql .backend .sea .models .base import ResultData
7167
7268 result_data = ResultData (data = filtered_rows , external_links = None )
7369
70+ from databricks .sql .result_set import SeaResultSet
71+
7472 # Create a new SeaResultSet with the filtered data
7573 filtered_result_set = SeaResultSet (
7674 connection = result_set .connection ,
@@ -85,11 +83,11 @@ def _filter_sea_result_set(
8583
8684 @staticmethod
8785 def filter_by_column_values (
88- result_set : " ResultSet" ,
86+ result_set : ResultSet ,
8987 column_index : int ,
9088 allowed_values : List [str ],
9189 case_sensitive : bool = False ,
92- ) -> " ResultSet" :
90+ ) -> ResultSet :
9391 """
9492 Filter a result set by values in a specific column.
9593
@@ -133,8 +131,8 @@ def filter_by_column_values(
133131
134132 @staticmethod
135133 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 :
138136 """
139137 Filter a result set of tables by the specified table types.
140138
0 commit comments