align flows

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

Author: varun-edachali-dbx

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

@@ -6,7 +6,9 @@
 
 from __future__ import annotations
 
+import io
 import logging
+from copy import deepcopy
 from typing import (
     List,
     Optional,
@@ -20,6 +22,9 @@
     from databricks.sql.backend.sea.result_set import SeaResultSet
 
 from databricks.sql.backend.types import ExecuteResponse
+from databricks.sql.backend.sea.models.base import ResultData
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
+from databricks.sql.utils import CloudFetchQueue, ArrowQueue
 
 try:
     import pyarrow
@@ -37,32 +42,18 @@ class ResultSetFilter:
     """
 
     @staticmethod
-    def _filter_sea_result_set(
-        result_set: SeaResultSet, filter_func: Callable[[List[Any]], bool]
-    ) -> SeaResultSet:
+    def _create_execute_response(result_set: SeaResultSet) -> ExecuteResponse:
         """
-        Filter a SEA result set using the provided filter function.
+        Create an ExecuteResponse with parameters from the original result set.
 
         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
+            result_set: Original result set to copy parameters from
 
         Returns:
-            A filtered SEA result set
+            ExecuteResponse: New execute response object
         """
-
-        # Get all remaining rows
-        all_rows = result_set.results.remaining_rows()
-
-        # Filter rows
-        filtered_rows = [row for row in all_rows if filter_func(row)]
-
-        # Reuse the command_id from the original result set
-        command_id = result_set.command_id
-
-        # Create an ExecuteResponse for the filtered data
-        execute_response = ExecuteResponse(
-            command_id=command_id,
+        return ExecuteResponse(
+            command_id=result_set.command_id,
             status=result_set.status,
             description=result_set.description,
             has_been_closed_server_side=result_set.has_been_closed_server_side,
@@ -71,29 +62,99 @@ def _filter_sea_result_set(
             is_staging_operation=False,
         )
 
-        # Create a new ResultData object with filtered data
-        from databricks.sql.backend.sea.models.base import ResultData
+    @staticmethod
+    def _create_filtered_manifest(result_set: SeaResultSet, new_row_count: int):
+        """
+        Create a copy of the manifest with updated row count.
+
+        Args:
+            result_set: Original result set to copy manifest from
+            new_row_count: New total row count for filtered data
 
-        result_data = ResultData(data=filtered_rows, external_links=None)
+        Returns:
+            Updated manifest copy
+        """
+        filtered_manifest = deepcopy(result_set.manifest)
+        filtered_manifest.total_row_count = new_row_count
+        return filtered_manifest
+
+    @staticmethod
+    def _create_filtered_result_set(
+        result_set: SeaResultSet,
+        result_data: ResultData,
+        row_count: int,
+    ) -> "SeaResultSet":
+        """
+        Create a new filtered SeaResultSet with the provided data.
+
+        Args:
+            result_set: Original result set to copy parameters from
+            result_data: New result data for the filtered set
+            row_count: Number of rows in the filtered data
 
-        from databricks.sql.backend.sea.backend import SeaDatabricksClient
+        Returns:
+            New filtered SeaResultSet
+        """
         from databricks.sql.backend.sea.result_set import SeaResultSet
 
-        # Create a new SeaResultSet with the filtered data
-        manifest = result_set.manifest
-        manifest.total_row_count = len(filtered_rows)
+        execute_response = ResultSetFilter._create_execute_response(result_set)
+        filtered_manifest = ResultSetFilter._create_filtered_manifest(
+            result_set, row_count
+        )
 
-        filtered_result_set = SeaResultSet(
+        return SeaResultSet(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
             result_data=result_data,
-            manifest=manifest,
+            manifest=filtered_manifest,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
         )
 
-        return filtered_result_set
+    @staticmethod
+    def _validate_column_index(result_set: SeaResultSet, column_index: int) -> str:
+        """
+        Validate column index and return the column name.
+
+        Args:
+            result_set: Result set to validate against
+            column_index: Index of the column to validate
+
+        Returns:
+            str: Column name at the specified index
+
+        Raises:
+            ValueError: If column index is out of bounds
+        """
+        if column_index >= len(result_set.description):
+            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(
@@ -112,7 +173,6 @@ def _filter_arrow_table(
         Returns:
             A filtered PyArrow table
         """
-
         if not pyarrow:
             raise ImportError("PyArrow is required for Arrow table filtering")
 
@@ -143,78 +203,34 @@ def _filter_arrow_result_set(
         Returns:
             A filtered SEA result set
         """
+        # Validate column index and get column name
+        column_name = ResultSetFilter._validate_column_index(result_set, column_index)
 
-        # Get all remaining rows as Arrow table
+        # Get all remaining rows as Arrow table and filter it
         arrow_table = result_set.results.remaining_rows()
-
-        # Get the column name from the description
-        if column_index >= len(result_set.description):
-            raise ValueError(f"Column index {column_index} is out of bounds")
-
-        column_name = result_set.description[column_index][0]
-
-        # Filter the Arrow table
         filtered_table = ResultSetFilter._filter_arrow_table(
             arrow_table, column_name, allowed_values
         )
 
-        # Create a new result set with filtered data
-        command_id = result_set.command_id
-
-        # Create an ExecuteResponse for the filtered data
-        execute_response = ExecuteResponse(
-            command_id=command_id,
-            status=result_set.status,
-            description=result_set.description,
-            has_been_closed_server_side=result_set.has_been_closed_server_side,
-            lz4_compressed=result_set.lz4_compressed,
-            arrow_schema_bytes=result_set._arrow_schema_bytes,
-            is_staging_operation=False,
-        )
-
-        # Create ResultData with the filtered arrow table as attachment
-        # This mimics the hybrid disposition flow in build_queue
-        from databricks.sql.backend.sea.models.base import ResultData
-        from databricks.sql.backend.sea.result_set import SeaResultSet
-        from databricks.sql.backend.sea.backend import SeaDatabricksClient
-        import io
-
-        # Convert the filtered table to Arrow stream format
+        # Convert the filtered table to Arrow stream format for ResultData
         sink = io.BytesIO()
         with pyarrow.ipc.new_stream(sink, filtered_table.schema) as writer:
             writer.write_table(filtered_table)
         arrow_stream_bytes = sink.getvalue()
 
         # Create ResultData with attachment containing the filtered data
-        filtered_result_data = ResultData(
+        result_data = ResultData(
             data=None,  # No JSON data
             external_links=None,  # No external links
             attachment=arrow_stream_bytes,  # Arrow data as attachment
         )
 
-        # Update manifest to reflect new row count
-        manifest = result_set.manifest
-        # Create a copy of the manifest to avoid modifying the original
-        from copy import deepcopy
-
-        filtered_manifest = deepcopy(manifest)
-        filtered_manifest.total_row_count = filtered_table.num_rows
-
-        # Create a new SeaResultSet with the filtered data
-        filtered_result_set = SeaResultSet(
-            connection=result_set.connection,
-            execute_response=execute_response,
-            sea_client=cast(SeaDatabricksClient, result_set.backend),
-            result_data=filtered_result_data,
-            manifest=filtered_manifest,
-            buffer_size_bytes=result_set.buffer_size_bytes,
-            arraysize=result_set.arraysize,
+        return ResultSetFilter._create_filtered_result_set(
+            result_set, result_data, filtered_table.num_rows
         )
 
-        return filtered_result_set
-
     @staticmethod
-    def filter_by_column_values(
+    def _filter_json_result_set(
         result_set: SeaResultSet,
         column_index: int,
         allowed_values: List[str],
@@ -237,7 +253,7 @@ def filter_by_column_values(
         if not case_sensitive:
             allowed_values = [v.upper() for v in allowed_values]
 
-        return ResultSetFilter._filter_sea_result_set(
+        return ResultSetFilter._filter_json_table(
             result_set,
             lambda row: (
                 len(row) > column_index
@@ -268,24 +284,19 @@ def filter_tables_by_type(
         Returns:
             A filtered result set containing only tables of the specified types
         """
-
         # Default table types if none specified
         DEFAULT_TABLE_TYPES = ["TABLE", "VIEW", "SYSTEM TABLE"]
-        valid_types = (
-            table_types if table_types and len(table_types) > 0 else DEFAULT_TABLE_TYPES
-        )
+        valid_types = table_types if table_types else DEFAULT_TABLE_TYPES
 
         # Check if we have an Arrow table (cloud fetch) or JSON data
-        from databricks.sql.utils import CloudFetchQueue, ArrowQueue
-
+        # Table type is the 6th column (index 5)
         if isinstance(result_set.results, (CloudFetchQueue, ArrowQueue)):
             # For Arrow tables, we need to handle filtering differently
             return ResultSetFilter._filter_arrow_result_set(
                 result_set, column_index=5, allowed_values=valid_types
             )
         else:
             # For JSON data, use the existing filter method
-            # Table type is the 6th column (index 5)
-            return ResultSetFilter.filter_by_column_values(
+            return ResultSetFilter._filter_json_result_set(
                 result_set, 5, valid_types, case_sensitive=True
             )

tests/unit/test_filters.py

@@ -68,7 +68,7 @@ def setUp(self):
         self.mock_sea_result_set.has_been_closed_server_side = False
         self.mock_sea_result_set._arrow_schema_bytes = None
 
-    def test_filter_by_column_values(self):
+    def test__filter_json_result_set(self):
         """Test filtering by column values with various options."""
         # Case 1: Case-sensitive filtering
         allowed_values = ["table1", "table3"]
@@ -82,8 +82,8 @@ def test_filter_by_column_values(self):
                 mock_instance = MagicMock()
                 mock_sea_result_set_class.return_value = mock_instance
 
-                # Call filter_by_column_values on the table_name column (index 2)
-                result = ResultSetFilter.filter_by_column_values(
+                # Call _filter_json_result_set on the table_name column (index 2)
+                result = ResultSetFilter._filter_json_result_set(
                     self.mock_sea_result_set, 2, allowed_values, case_sensitive=True
                 )
 
@@ -109,8 +109,8 @@ def test_filter_by_column_values(self):
                 mock_instance = MagicMock()
                 mock_sea_result_set_class.return_value = mock_instance
 
-                # Call filter_by_column_values with case-insensitive matching
-                result = ResultSetFilter.filter_by_column_values(
+                # Call _filter_json_result_set with case-insensitive matching
+                result = ResultSetFilter._filter_json_result_set(
                     self.mock_sea_result_set,
                     2,
                     ["TABLE1", "TABLE3"],
@@ -128,7 +128,7 @@ def test_filter_tables_by_type(self):
 
         self.mock_sea_result_set.results = JsonQueue([])
 
-        with patch.object(ResultSetFilter, "filter_by_column_values") as mock_filter:
+        with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter:
             ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, table_types)
             args, kwargs = mock_filter.call_args
             self.assertEqual(args[0], self.mock_sea_result_set)
@@ -137,7 +137,7 @@ def test_filter_tables_by_type(self):
             self.assertEqual(kwargs.get("case_sensitive"), True)
 
         # Case 2: Default table types (None or empty list)
-        with patch.object(ResultSetFilter, "filter_by_column_values") as mock_filter:
+        with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter:
             # Test with None
             ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
             args, kwargs = mock_filter.call_args