Merge branch 'main' into col-normalisation

Author: varun-edachali-dbx

src/databricks/sql/auth/retry.py

@@ -127,7 +127,7 @@ def __init__(
             total=_attempts_remaining,
             respect_retry_after_header=True,
             backoff_factor=self.delay_min,
-            allowed_methods=["POST"],
+            allowed_methods=["POST", "GET", "DELETE"],
             status_forcelist=[429, 503, *self.force_dangerous_codes],
         )
 

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

@@ -159,6 +159,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)
@@ -695,7 +696,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,
@@ -731,7 +732,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,
@@ -775,7 +776,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,
@@ -825,7 +826,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

@@ -6,12 +6,12 @@
 
 from __future__ import annotations
 
+import io
 import logging
 from typing import (
     List,
     Optional,
     Any,
-    Callable,
     cast,
     TYPE_CHECKING,
 )
@@ -20,6 +20,16 @@
     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
+    import pyarrow.compute as pc
+except ImportError:
+    pyarrow = None
+    pc = None
 
 logger = logging.getLogger(__name__)
 
@@ -30,32 +40,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,
@@ -64,39 +60,145 @@ 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 _update_manifest(result_set: SeaResultSet, new_row_count: int):
+        """
+        Create a copy of the manifest with updated row count.
 
-        result_data = ResultData(data=filtered_rows, external_links=None)
+        Args:
+            result_set: Original result set to copy manifest from
+            new_row_count: New total row count for filtered data
+
+        Returns:
+            Updated manifest copy
+        """
+        filtered_manifest = 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.
 
-        from databricks.sql.backend.sea.backend import SeaDatabricksClient
+        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
+
+        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._update_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,
         )
 
-        # Preserve metadata columns setup from original result set
-        if hasattr(result_set, "_metadata_columns") and result_set._metadata_columns:
-            filtered_result_set._metadata_columns = result_set._metadata_columns
-            filtered_result_set._column_index_mapping = result_set._column_index_mapping
-            # Update the description to match the original prepared description
-            filtered_result_set.description = result_set.description
+    @staticmethod
+    def _filter_arrow_table(
+        table: Any,  # pyarrow.Table
+        column_name: str,
+        allowed_values: List[str],
+        case_sensitive: bool = True,
+    ) -> 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
+            case_sensitive: Whether to perform case-sensitive comparison
+
+        Returns:
+            A filtered PyArrow table
+        """
+        if not pyarrow:
+            raise ImportError("PyArrow is required for Arrow table filtering")
+
+        if table.num_rows == 0:
+            return table
+
+        # Handle case-insensitive filtering by normalizing both column and allowed values
+        if not case_sensitive:
+            # Convert allowed values to uppercase
+            allowed_values = [v.upper() for v in allowed_values]
+            # Get column values as uppercase
+            column = pc.utf8_upper(table[column_name])
+        else:
+            # Use column as-is
+            column = table[column_name]
+
+        # Convert allowed_values to PyArrow Array
+        allowed_array = pyarrow.array(allowed_values)
+
+        # Construct a boolean mask: True where column is in allowed_list
+        mask = pc.is_in(column, value_set=allowed_array)
+        return table.filter(mask)
+
+    @staticmethod
+    def _filter_arrow_result_set(
+        result_set: SeaResultSet,
+        column_index: int,
+        allowed_values: List[str],
+        case_sensitive: bool = True,
+    ) -> 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
+            case_sensitive: Whether to perform case-sensitive comparison
+
+        Returns:
+            A filtered SEA result set
+        """
+        # Validate column index and get column name
+        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]
+
+        # Get all remaining rows as Arrow table and filter it
+        arrow_table = result_set.results.remaining_rows()
+        filtered_table = ResultSetFilter._filter_arrow_table(
+            arrow_table, column_name, allowed_values, case_sensitive
+        )
+
+        # 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
+        result_data = ResultData(
+            data=None,  # No JSON data
+            external_links=None,  # No external links
+            attachment=arrow_stream_bytes,  # Arrow data as attachment
+        )
 
-        return filtered_result_set
+        return ResultSetFilter._create_filtered_result_set(
+            result_set, result_data, filtered_table.num_rows
+        )
 
     @staticmethod
-    def filter_by_column_values(
+    def _filter_json_result_set(
         result_set: SeaResultSet,
         column_index: int,
         allowed_values: List[str],
@@ -114,22 +216,35 @@ def filter_by_column_values(
         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")
 
-        # Convert to uppercase for case-insensitive comparison if needed
+        # Extract rows
+        all_rows = result_set.results.remaining_rows()
+
+        # 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_sea_result_set(
-            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 ResultSetFilter._create_filtered_result_set(
+            result_set, result_data, len(filtered_rows)
         )
 
     @staticmethod
@@ -150,14 +265,25 @@ 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
         # Table type is the 6th column (index 5)
-        return ResultSetFilter.filter_by_column_values(
-            result_set, 5, valid_types, case_sensitive=True
-        )
+        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,
+                case_sensitive=True,
+            )
+        else:
+            # For JSON data, use the existing filter method
+            return ResultSetFilter._filter_json_result_set(
+                result_set,
+                column_index=5,
+                allowed_values=valid_types,
+                case_sensitive=True,
+            )