align flow of json with arrow

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/utils/filters.py

@@ -13,7 +13,6 @@
     List,
     Optional,
     Any,
-    Callable,
     cast,
     TYPE_CHECKING,
 )
@@ -131,31 +130,6 @@ def _validate_column_index(result_set: SeaResultSet, column_index: int) -> str:
             raise ValueError(f"Column index {column_index} is out of bounds")
         return result_set.description[column_index][0]
 
-    @staticmethod
-    def _filter_json_table(
-        result_set: SeaResultSet, filter_func: Callable[[List[Any]], bool]
-    ) -> SeaResultSet:
-        """
-        Filter a SEA result set using the provided filter function.
-
-        Args:
-            result_set: The SEA result set to filter
-            filter_func: Function that takes a row and returns True if the row should be included
-
-        Returns:
-            A filtered SEA result set
-        """
-        # Get all remaining rows and filter them
-        all_rows = result_set.results.remaining_rows()
-        filtered_rows = [row for row in all_rows if filter_func(row)]
-
-        # Create ResultData with filtered rows
-        result_data = ResultData(data=filtered_rows, external_links=None)
-
-        return ResultSetFilter._create_filtered_result_set(
-            result_set, result_data, len(filtered_rows)
-        )
-
     @staticmethod
     def _filter_arrow_table(
         table: Any,  # pyarrow.Table
@@ -248,22 +222,36 @@ def _filter_json_result_set(
         Returns:
             A filtered result set
         """
+        # Validate column index (optional - not in arrow version but good practice)
+        if column_index >= len(result_set.description):
+            raise ValueError(f"Column index {column_index} is out of bounds")
+
+        # Extract rows
+        all_rows = result_set.results.remaining_rows()
 
-        # Convert to uppercase for case-insensitive comparison if needed
+        # Convert allowed values if case-insensitive
         if not case_sensitive:
             allowed_values = [v.upper() for v in allowed_values]
+        # Helper lambda to get column value based on case sensitivity
+        get_column_value = (
+            lambda row: row[column_index].upper()
+            if not case_sensitive
+            else row[column_index]
+        )
+
+        # Filter rows based on allowed values
+        filtered_rows = [
+            row
+            for row in all_rows
+            if len(row) > column_index and get_column_value(row) in allowed_values
+        ]
+
+        # Create filtered result set
+        result_data = ResultData(data=filtered_rows, external_links=None)
 
-        return ResultSetFilter._filter_json_table(
-            result_set,
-            lambda row: (
-                len(row) > column_index
-                and (
-                    row[column_index].upper()
-                    if not case_sensitive
-                    else row[column_index]
-                )
-                in allowed_values
-            ),
+        # Return
+        return ResultSetFilter._create_filtered_result_set(
+            result_set, result_data, len(filtered_rows)
         )
 
     @staticmethod