99 List ,
1010 Optional ,
1111 Any ,
12- Dict ,
1312 Callable ,
14- TypeVar ,
15- Generic ,
1613 cast ,
17- TYPE_CHECKING ,
1814)
1915
20- from databricks .sql .backend .types import ExecuteResponse , CommandId
21- from databricks .sql .backend .sea .models .base import ResultData
2216from databricks .sql .backend .sea .backend import SeaDatabricksClient
17+ from databricks .sql .backend .types import ExecuteResponse
2318
24- if TYPE_CHECKING :
25- from databricks .sql .result_set import ResultSet , SeaResultSet
19+ from databricks .sql .result_set import ResultSet , SeaResultSet
2620
2721logger = logging .getLogger (__name__ )
2822
2923
3024class ResultSetFilter :
3125 """
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.
26+ A general-purpose filter for result sets.
3627 """
3728
3829 @staticmethod
3930 def _filter_sea_result_set (
40- result_set : " SeaResultSet" , filter_func : Callable [[List [Any ]], bool ]
41- ) -> " SeaResultSet" :
31+ result_set : SeaResultSet , filter_func : Callable [[List [Any ]], bool ]
32+ ) -> SeaResultSet :
4233 """
4334 Filter a SEA result set using the provided filter function.
4435
@@ -49,15 +40,13 @@ def _filter_sea_result_set(
4940 Returns:
5041 A filtered SEA result set
5142 """
43+
5244 # Get all remaining rows
5345 all_rows = result_set .results .remaining_rows ()
5446
5547 # Filter rows
5648 filtered_rows = [row for row in all_rows if filter_func (row )]
5749
58- # Import SeaResultSet here to avoid circular imports
59- from databricks .sql .result_set import SeaResultSet
60-
6150 # Reuse the command_id from the original result set
6251 command_id = result_set .command_id
6352
@@ -73,10 +62,13 @@ def _filter_sea_result_set(
7362 )
7463
7564 # Create a new ResultData object with filtered data
65+
7666 from databricks .sql .backend .sea .models .base import ResultData
7767
7868 result_data = ResultData (data = filtered_rows , external_links = None )
7969
70+ from databricks .sql .result_set import SeaResultSet
71+
8072 # Create a new SeaResultSet with the filtered data
8173 filtered_result_set = SeaResultSet (
8274 connection = result_set .connection ,
@@ -91,11 +83,11 @@ def _filter_sea_result_set(
9183
9284 @staticmethod
9385 def filter_by_column_values (
94- result_set : " ResultSet" ,
86+ result_set : ResultSet ,
9587 column_index : int ,
9688 allowed_values : List [str ],
9789 case_sensitive : bool = False ,
98- ) -> " ResultSet" :
90+ ) -> ResultSet :
9991 """
10092 Filter a result set by values in a specific column.
10193
@@ -108,6 +100,7 @@ def filter_by_column_values(
108100 Returns:
109101 A filtered result set
110102 """
103+
111104 # Convert to uppercase for case-insensitive comparison if needed
112105 if not case_sensitive :
113106 allowed_values = [v .upper () for v in allowed_values ]
@@ -138,8 +131,8 @@ def filter_by_column_values(
138131
139132 @staticmethod
140133 def filter_tables_by_type (
141- result_set : " ResultSet" , table_types : Optional [List [str ]] = None
142- ) -> " ResultSet" :
134+ result_set : ResultSet , table_types : Optional [List [str ]] = None
135+ ) -> ResultSet :
143136 """
144137 Filter a result set of tables by the specified table types.
145138
@@ -154,6 +147,7 @@ def filter_tables_by_type(
154147 Returns:
155148 A filtered result set containing only tables of the specified types
156149 """
150+
157151 # Default table types if none specified
158152 DEFAULT_TABLE_TYPES = ["TABLE" , "VIEW" , "SYSTEM TABLE" ]
159153 valid_types = (
0 commit comments