preliminary large metadata results

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

Author: varun-edachali-dbx

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

@@ -157,6 +157,7 @@ def __init__(
         )
 
         self.use_hybrid_disposition = kwargs.get("use_hybrid_disposition", True)
+        self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
 
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
@@ -688,7 +689,7 @@ def get_catalogs(
             max_bytes=max_bytes,
             lz4_compression=False,
             cursor=cursor,
-            use_cloud_fetch=False,
+            use_cloud_fetch=self.use_cloud_fetch,
             parameters=[],
             async_op=False,
             enforce_embedded_schema_correctness=False,
@@ -721,7 +722,7 @@ def get_schemas(
             max_bytes=max_bytes,
             lz4_compression=False,
             cursor=cursor,
-            use_cloud_fetch=False,
+            use_cloud_fetch=self.use_cloud_fetch,
             parameters=[],
             async_op=False,
             enforce_embedded_schema_correctness=False,
@@ -762,7 +763,7 @@ def get_tables(
             max_bytes=max_bytes,
             lz4_compression=False,
             cursor=cursor,
-            use_cloud_fetch=False,
+            use_cloud_fetch=self.use_cloud_fetch,
             parameters=[],
             async_op=False,
             enforce_embedded_schema_correctness=False,
@@ -809,7 +810,7 @@ def get_columns(
             max_bytes=max_bytes,
             lz4_compression=False,
             cursor=cursor,
-            use_cloud_fetch=False,
+            use_cloud_fetch=self.use_cloud_fetch,
             parameters=[],
             async_op=False,
             enforce_embedded_schema_correctness=False,

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

@@ -21,6 +21,13 @@
 
 from databricks.sql.backend.types import ExecuteResponse
 
+try:
+    import pyarrow
+    import pyarrow.compute as pc
+except ImportError:
+    pyarrow = None
+    pc = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -88,6 +95,121 @@ def _filter_sea_result_set(
 
         return filtered_result_set
 
+    @staticmethod
+    def _filter_arrow_table(
+        table: Any,  # pyarrow.Table
+        column_name: str,
+        allowed_values: List[str],
+    ) -> Any:  # returns pyarrow.Table
+        """
+        Filter a PyArrow table by column values.
+
+        Args:
+            table: The PyArrow table to filter
+            column_name: The name of the column to filter on
+            allowed_values: List of allowed values for the column
+
+        Returns:
+            A filtered PyArrow table
+        """
+
+        if not pyarrow:
+            raise ImportError("PyArrow is required for Arrow table filtering")
+
+        # Convert allowed_values to PyArrow Array for better performance
+        allowed_array = pyarrow.array(allowed_values)
+
+        # Construct a boolean mask: True where column is in allowed_list
+        mask = pc.is_in(table[column_name], value_set=allowed_array)
+        return table.filter(mask)
+
+    @staticmethod
+    def _filter_arrow_result_set(
+        result_set: SeaResultSet,
+        column_index: int,
+        allowed_values: List[str],
+    ) -> SeaResultSet:
+        """
+        Filter a SEA result set that contains Arrow tables.
+
+        Args:
+            result_set: The SEA result set to filter (containing Arrow data)
+            column_index: The index of the column to filter on
+            allowed_values: List of allowed values for the column
+
+        Returns:
+            A filtered SEA result set
+        """
+
+        # Get all remaining rows as Arrow table
+        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
+        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(
+            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 filtered_result_set
+
     @staticmethod
     def filter_by_column_values(
         result_set: SeaResultSet,
@@ -150,7 +272,17 @@ def filter_tables_by_type(
             table_types if table_types and len(table_types) > 0 else DEFAULT_TABLE_TYPES
         )
 
-        # Table type is the 6th column (index 5)
-        return ResultSetFilter.filter_by_column_values(
-            result_set, 5, valid_types, case_sensitive=True
-        )
+        # Check if we have an Arrow table (cloud fetch) or JSON data
+        from databricks.sql.utils import CloudFetchQueue, ArrowQueue
+
+        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(
+                result_set, 5, valid_types, case_sensitive=True
+            )

tests/unit/test_filters.py

@@ -123,37 +123,30 @@ def test_filter_tables_by_type(self):
         # Case 1: Specific table types
         table_types = ["TABLE", "VIEW"]
 
-        with patch(
-            "databricks.sql.backend.sea.utils.filters.isinstance", return_value=True
-        ):
-            with patch.object(
-                ResultSetFilter, "filter_by_column_values"
-            ) 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)
-                self.assertEqual(args[1], 5)  # Table type column index
-                self.assertEqual(args[2], table_types)
-                self.assertEqual(kwargs.get("case_sensitive"), True)
+        # Mock results as JsonQueue (not CloudFetchQueue or ArrowQueue)
+        from databricks.sql.backend.sea.queue import JsonQueue
+
+        self.mock_sea_result_set.results = JsonQueue([])
+
+        with patch.object(ResultSetFilter, "filter_by_column_values") 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)
+            self.assertEqual(args[1], 5)  # Table type column index
+            self.assertEqual(args[2], table_types)
+            self.assertEqual(kwargs.get("case_sensitive"), True)
 
         # Case 2: Default table types (None or empty list)
-        with patch(
-            "databricks.sql.backend.sea.utils.filters.isinstance", return_value=True
-        ):
-            with patch.object(
-                ResultSetFilter, "filter_by_column_values"
-            ) as mock_filter:
-                # Test with None
-                ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
-                args, kwargs = mock_filter.call_args
-                self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
-
-                # Test with empty list
-                ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, [])
-                args, kwargs = mock_filter.call_args
-                self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
+        with patch.object(ResultSetFilter, "filter_by_column_values") as mock_filter:
+            # Test with None
+            ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
+            args, kwargs = mock_filter.call_args
+            self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
+
+            # Test with empty list
+            ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, [])
+            args, kwargs = mock_filter.call_args
+            self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
 
 
 if __name__ == "__main__":

tests/unit/test_sea_backend.py

@@ -56,6 +56,29 @@ def sea_client(self, mock_http_client):
             http_headers=http_headers,
             auth_provider=auth_provider,
             ssl_options=ssl_options,
+            use_cloud_fetch=False,
+        )
+
+        return client
+
+    @pytest.fixture
+    def sea_client_cloud_fetch(self, mock_http_client):
+        """Create a SeaDatabricksClient instance with cloud fetch enabled."""
+        server_hostname = "test-server.databricks.com"
+        port = 443
+        http_path = "/sql/warehouses/abc123"
+        http_headers = [("header1", "value1"), ("header2", "value2")]
+        auth_provider = AuthProvider()
+        ssl_options = SSLOptions()
+
+        client = SeaDatabricksClient(
+            server_hostname=server_hostname,
+            port=port,
+            http_path=http_path,
+            http_headers=http_headers,
+            auth_provider=auth_provider,
+            ssl_options=ssl_options,
+            use_cloud_fetch=True,
         )
 
         return client
@@ -884,3 +907,74 @@ def test_get_columns(self, sea_client, sea_session_id, mock_cursor):
                     cursor=mock_cursor,
                 )
             assert "Catalog name is required for get_columns" in str(excinfo.value)
+
+    def test_get_tables_with_cloud_fetch(
+        self, sea_client_cloud_fetch, sea_session_id, mock_cursor
+    ):
+        """Test the get_tables method with cloud fetch enabled."""
+        # Mock the execute_command method and ResultSetFilter
+        mock_result_set = Mock()
+
+        with patch.object(
+            sea_client_cloud_fetch, "execute_command", return_value=mock_result_set
+        ) as mock_execute:
+            with patch(
+                "databricks.sql.backend.sea.utils.filters.ResultSetFilter"
+            ) as mock_filter:
+                mock_filter.filter_tables_by_type.return_value = mock_result_set
+
+                # Call get_tables
+                result = sea_client_cloud_fetch.get_tables(
+                    session_id=sea_session_id,
+                    max_rows=100,
+                    max_bytes=1000,
+                    cursor=mock_cursor,
+                    catalog_name="test_catalog",
+                )
+
+                # Verify execute_command was called with use_cloud_fetch=True
+                mock_execute.assert_called_with(
+                    operation="SHOW TABLES IN CATALOG test_catalog",
+                    session_id=sea_session_id,
+                    max_rows=100,
+                    max_bytes=1000,
+                    lz4_compression=False,
+                    cursor=mock_cursor,
+                    use_cloud_fetch=True,  # Should use True since client was created with use_cloud_fetch=True
+                    parameters=[],
+                    async_op=False,
+                    enforce_embedded_schema_correctness=False,
+                )
+                assert result == mock_result_set
+
+    def test_get_schemas_with_cloud_fetch(
+        self, sea_client_cloud_fetch, sea_session_id, mock_cursor
+    ):
+        """Test the get_schemas method with cloud fetch enabled."""
+        # Mock the execute_command method
+        mock_result_set = Mock()
+        with patch.object(
+            sea_client_cloud_fetch, "execute_command", return_value=mock_result_set
+        ) as mock_execute:
+            # Test with catalog name
+            result = sea_client_cloud_fetch.get_schemas(
+                session_id=sea_session_id,
+                max_rows=100,
+                max_bytes=1000,
+                cursor=mock_cursor,
+                catalog_name="test_catalog",
+            )
+
+            mock_execute.assert_called_with(
+                operation="SHOW SCHEMAS IN test_catalog",
+                session_id=sea_session_id,
+                max_rows=100,
+                max_bytes=1000,
+                lz4_compression=False,
+                cursor=mock_cursor,
+                use_cloud_fetch=True,  # Should use True since client was created with use_cloud_fetch=True
+                parameters=[],
+                async_op=False,
+                enforce_embedded_schema_correctness=False,
+            )
+            assert result == mock_result_set