Aug 4, 2025 · May 28, 2025 · May 30, 2025 · Jun 3, 2025 · Jun 4, 2025 · Jun 9, 2025
diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py
        )

        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)
            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,
            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,
            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,
            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,
diff --git a/src/databricks/sql/backend/sea/utils/filters.py b/src/databricks/sql/backend/sea/utils/filters.py

 from __future__ import annotations

 import io
 import logging
 from typing import (
    List,
    Optional,
    Any,
    Callable,
    cast,
    TYPE_CHECKING,
 )
    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__)

    """

    @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 theprovided filter function.
 Create an ExecuteResponse with parameters from theoriginal 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,
            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.

        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 = result_set.manifest
        filtered_manifest.total_row_count = new_row_count
        return filtered_manifest

        from databricks.sql.backend.sea.backend import SeaDatabricksClient
    @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

        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,
        )

        return filtered_result_set
    @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 ResultSetFilter._create_filtered_result_set(
            result_set, result_data, filtered_table.num_rows
        )

    @staticmethod
    deffilter_by_column_values(
    def_filter_json_result_set(
        result_set: SeaResultSet,
        column_index: int,
        allowed_values: List[str],
        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
        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,
            )
Original file line number	Diff line number	Diff line change
Expand Up		@@ -158,6 +158,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)
Expand DownExpand Up		@@ -694,7 +695,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,
Expand DownExpand Up		@@ -727,7 +728,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,
Expand DownExpand Up		@@ -768,7 +769,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,
Expand DownExpand Up		@@ -815,7 +816,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,
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,12 +6,12 @@

		from __future__ import annotations

		import io
		import logging
		from typing import (
		List,
		Optional,
		Any,
		Callable,
		cast,
		TYPE_CHECKING,
		)
Expand All		@@ -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__)

Expand All		@@ -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 theprovided filter function.
		Create an ExecuteResponse with parameters from theoriginal 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,
Expand All		@@ -64,32 +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.

		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 = result_set.manifest
		filtered_manifest.total_row_count = new_row_count
		return filtered_manifest

		from databricks.sql.backend.sea.backend import SeaDatabricksClient
		@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

		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,
		)

		return filtered_result_set
		@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 ResultSetFilter._create_filtered_result_set(
		result_set, result_data, filtered_table.num_rows
		)

		@staticmethod
		deffilter_by_column_values(
		def_filter_json_result_set(
		result_set: SeaResultSet,
		column_index: int,
		allowed_values: List[str],
Expand All		@@ -107,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
Expand All		@@ -143,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,
		)