Complete Fetch Phase (for INLINE disposition and JSON_ARRAY format)

examples/experimental/tests/test_sea_async_query.py

@@ -51,12 +51,20 @@ def test_sea_async_query_with_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query asynchronously
+        # Execute a query that generates large rows to force multiple chunks
+        requested_row_count = 5000
         cursor = connection.cursor()
+        query = f"""
+        SELECT 
+            id, 
+            concat('value_', repeat('a', 10000)) as test_value
+        FROM range(1, {requested_row_count} + 1) AS t(id)
+        """
+
         logger.info(
-            "Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
+            f"Executing asynchronous query with cloud fetch to generate {requested_row_count} rows"
         )
-        cursor.execute_async("SELECT 1 as test_value")
+        cursor.execute_async(query)
         logger.info(
             "Asynchronous query submitted successfully with cloud fetch enabled"
         )
@@ -69,8 +77,25 @@ def test_sea_async_query_with_cloud_fetch():
 
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
+
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        actual_row_count = len(results)
+
+        logger.info(
+            f"Requested {requested_row_count} rows, received {actual_row_count} rows"
+        )
+
+        # Verify total row count
+        if actual_row_count != requested_row_count:
+            logger.error(
+                f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
+            return False
+
         logger.info(
-            "Successfully retrieved asynchronous query results with cloud fetch enabled"
+            "PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly"
         )
 
         # Close resources
@@ -130,12 +155,20 @@ def test_sea_async_query_without_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query asynchronously
+        # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
+        requested_row_count = 100
         cursor = connection.cursor()
+        query = f"""
+        SELECT 
+            id, 
+            concat('value_', repeat('a', 100)) as test_value
+        FROM range(1, {requested_row_count} + 1) AS t(id)
+        """
+
         logger.info(
-            "Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
+            f"Executing asynchronous query without cloud fetch to generate {requested_row_count} rows"
         )
-        cursor.execute_async("SELECT 1 as test_value")
+        cursor.execute_async(query)
         logger.info(
             "Asynchronous query submitted successfully with cloud fetch disabled"
         )
@@ -148,8 +181,24 @@ def test_sea_async_query_without_cloud_fetch():
 
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        actual_row_count = len(results)
+
+        logger.info(
+            f"Requested {requested_row_count} rows, received {actual_row_count} rows"
+        )
+
+        # Verify total row count
+        if actual_row_count != requested_row_count:
+            logger.error(
+                f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
+            return False
+
         logger.info(
-            "Successfully retrieved asynchronous query results with cloud fetch disabled"
+            "PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly"
         )
 
         # Close resources

examples/experimental/tests/test_sea_sync_query.py

@@ -49,13 +49,27 @@ def test_sea_sync_query_with_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query
+        # Execute a query that generates large rows to force multiple chunks
+        requested_row_count = 10000
         cursor = connection.cursor()
+        query = f"""
+        SELECT 
+            id, 
+            concat('value_', repeat('a', 10000)) as test_value
+        FROM range(1, {requested_row_count} + 1) AS t(id)
+        """
+
+        logger.info(
+            f"Executing synchronous query with cloud fetch to generate {requested_row_count} rows"
+        )
+        cursor.execute(query)
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        actual_row_count = len(results)
         logger.info(
-            "Executing synchronous query with cloud fetch: SELECT 1 as test_value"
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
         )
-        cursor.execute("SELECT 1 as test_value")
-        logger.info("Query executed successfully with cloud fetch enabled")
 
         # Close resources
         cursor.close()
@@ -114,13 +128,18 @@ def test_sea_sync_query_without_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query
+        # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
+        requested_row_count = 100
         cursor = connection.cursor()
-        logger.info(
-            "Executing synchronous query without cloud fetch: SELECT 1 as test_value"
+        logger.info("Executing synchronous query without cloud fetch: SELECT 100 rows")
+        cursor.execute(
+            "SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
         )
-        cursor.execute("SELECT 1 as test_value")
-        logger.info("Query executed successfully with cloud fetch disabled")
+
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        logger.info(f"{len(results)} rows retrieved against 100 requested")
 
         # Close resources
         cursor.close()

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

@@ -616,10 +616,10 @@ def get_execution_result(
             connection=cursor.connection,
             execute_response=execute_response,
             sea_client=self,
-            buffer_size_bytes=cursor.buffer_size_bytes,
-            arraysize=cursor.arraysize,
             result_data=response.result,
             manifest=response.manifest,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
         )
 
     # == Metadata Operations ==

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

@@ -0,0 +1,135 @@
+"""
+Type conversion utilities for the Databricks SQL Connector.
+
+This module provides functionality to convert string values from SEA Inline results
+to appropriate Python types based on column metadata.
+"""
+
+import datetime
+import decimal
+import logging
+from dateutil import parser
+from typing import Any, Callable, Dict, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+class SqlType:
+    """
+    SQL type constants
+
+    The list of types can be found in the SEA REST API Reference:
+    https://docs.databricks.com/api/workspace/statementexecution/executestatement
+    """
+
+    # Numeric types
+    BYTE = "byte"
+    SHORT = "short"
+    INT = "int"
+    LONG = "long"
+    FLOAT = "float"
+    DOUBLE = "double"
+    DECIMAL = "decimal"
+
+    # Boolean type
+    BOOLEAN = "boolean"
+
+    # Date/Time types
+    DATE = "date"
+    TIMESTAMP = "timestamp"
+    INTERVAL = "interval"
+
+    # String types
+    CHAR = "char"
+    STRING = "string"
+
+    # Binary type
+    BINARY = "binary"
+
+    # Complex types
+    ARRAY = "array"
+    MAP = "map"
+    STRUCT = "struct"
+
+    # Other types
+    NULL = "null"
+    USER_DEFINED_TYPE = "user_defined_type"
+
+
+class SqlTypeConverter:
+    """
+    Utility class for converting SQL types to Python types.
+    Based on the types supported by the Databricks SDK.
+    """
+
+    # SQL type to conversion function mapping
+    # TODO: complex types
+    TYPE_MAPPING: Dict[str, Callable] = {
+        # Numeric types
+        SqlType.BYTE: lambda v: int(v),
+        SqlType.SHORT: lambda v: int(v),
+        SqlType.INT: lambda v: int(v),
+        SqlType.LONG: lambda v: int(v),
+        SqlType.FLOAT: lambda v: float(v),
+        SqlType.DOUBLE: lambda v: float(v),
+        SqlType.DECIMAL: lambda v, p=None, s=None: (
+            decimal.Decimal(v).quantize(
+                decimal.Decimal(f'0.{"0" * s}'), context=decimal.Context(prec=p)
+            )
+            if p is not None and s is not None
+            else decimal.Decimal(v)
+        ),
+        # Boolean type
+        SqlType.BOOLEAN: lambda v: v.lower() in ("true", "t", "1", "yes", "y"),
+        # Date/Time types
+        SqlType.DATE: lambda v: datetime.date.fromisoformat(v),
+        SqlType.TIMESTAMP: lambda v: parser.parse(v),
+        SqlType.INTERVAL: lambda v: v,  # Keep as string for now
+        # String types - no conversion needed
+        SqlType.CHAR: lambda v: v,
+        SqlType.STRING: lambda v: v,
+        # Binary type
+        SqlType.BINARY: lambda v: bytes.fromhex(v),
+        # Other types
+        SqlType.NULL: lambda v: None,
+        # Complex types and user-defined types return as-is
+        SqlType.USER_DEFINED_TYPE: lambda v: v,
+    }
+
+    @staticmethod
+    def convert_value(
+        value: Any,
+        sql_type: str,
+        precision: Optional[int] = None,
+        scale: Optional[int] = None,
+    ) -> Any:
+        """
+        Convert a string value to the appropriate Python type based on SQL type.
+
+        Args:
+            value: The string value to convert
+            sql_type: The SQL type (e.g., 'int', 'decimal')
+            precision: Optional precision for decimal types
+            scale: Optional scale for decimal types
+
+        Returns:
+            The converted value in the appropriate Python type
+        """
+
+        if value is None:
+            return None
+
+        sql_type = sql_type.lower().strip()
+
+        if sql_type not in SqlTypeConverter.TYPE_MAPPING:
+            return value
+
+        converter_func = SqlTypeConverter.TYPE_MAPPING[sql_type]
+        try:
+            if sql_type == SqlType.DECIMAL:
+                return converter_func(value, precision, scale)
+            else:
+                return converter_func(value)
+        except (ValueError, TypeError, decimal.InvalidOperation) as e:
+            logger.warning(f"Error converting value '{value}' to {sql_type}: {e}")
+            return value

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

@@ -77,9 +77,9 @@ def _filter_sea_result_set(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
+            result_data=result_data,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
-            result_data=result_data,
         )
 
         return filtered_result_set

src/databricks/sql/backend/thrift_backend.py

@@ -40,11 +40,11 @@
 )
 
 from databricks.sql.utils import (
-    ResultSetQueueFactory,
+    ThriftResultSetQueueFactory,
     _bound,
     RequestErrorInfo,
     NoRetryReason,
-    ResultSetQueueFactory,
+    ThriftResultSetQueueFactory,
     convert_arrow_based_set_to_arrow_table,
     convert_decimals_in_arrow_table,
     convert_column_based_set_to_arrow_table,
@@ -1232,7 +1232,7 @@ def fetch_results(
                 )
             )
 
-        queue = ResultSetQueueFactory.build_queue(
+        queue = ThriftResultSetQueueFactory.build_queue(
             row_set_type=resp.resultSetMetadata.resultFormat,
             t_row_set=resp.results,
             arrow_schema_bytes=arrow_schema_bytes,