Merge branch 'main' into jahnvi/entra_cache

Author: jahnvi480

mssql_python/cursor.py

@@ -36,7 +36,10 @@
 )
 
 if TYPE_CHECKING:
+    import pyarrow  # type: ignore
     from mssql_python.connection import Connection
+else:
+    pyarrow = None
 
 # Constants for string handling
 MAX_INLINE_CHAR: int = (
@@ -775,6 +778,19 @@ def _check_closed(self) -> None:
                 ddbc_error="",
             )
 
+    def _ensure_pyarrow(self) -> Any:
+        """
+        Import and return pyarrow or raise ImportError accordingly.
+        """
+        try:
+            import pyarrow
+
+            return pyarrow
+        except ImportError as e:
+            raise ImportError(
+                "pyarrow is required for Arrow fetch methods. Please install pyarrow."
+            ) from e
+
     def setinputsizes(self, sizes: List[Union[int, tuple]]) -> None:
         """
         Sets the type information to be used for parameters in execute and executemany.
@@ -2516,6 +2532,94 @@ def fetchall(self) -> List[Row]:
             # On error, don't increment rownumber - rethrow the error
             raise e
 
+    def arrow_batch(self, batch_size: int = 8192) -> "pyarrow.RecordBatch":
+        """
+        Fetch a single pyarrow Record Batch of the specified size from the
+        query result set.
+
+        Args:
+            batch_size: Maximum number of rows to fetch in the Record Batch.
+
+        Returns:
+            A pyarrow RecordBatch object containing up to batch_size rows.
+        """
+        self._check_closed()  # Check if the cursor is closed
+        pyarrow = self._ensure_pyarrow()
+
+        if not self._has_result_set and self.description:
+            self._reset_rownumber()
+
+        capsules = []
+        ret = ddbc_bindings.DDBCSQLFetchArrowBatch(self.hstmt, capsules, max(batch_size, 0))
+        check_error(ddbc_sql_const.SQL_HANDLE_STMT.value, self.hstmt, ret)
+
+        batch = pyarrow.RecordBatch._import_from_c_capsule(*capsules)
+
+        if self.hstmt:
+            self.messages.extend(ddbc_bindings.DDBCSQLGetAllDiagRecords(self.hstmt))
+
+        # Update rownumber for the number of rows actually fetched
+        num_fetched = batch.num_rows
+        if num_fetched > 0 and self._has_result_set:
+            self._next_row_index += num_fetched
+            self._rownumber = self._next_row_index - 1
+
+        # Centralize rowcount assignment after fetch
+        if num_fetched == 0 and self._next_row_index == 0:
+            self.rowcount = 0
+        else:
+            self.rowcount = self._next_row_index
+
+        return batch
+
+    def arrow(self, batch_size: int = 8192) -> "pyarrow.Table":
+        """
+        Fetch the entire result as a pyarrow Table.
+
+        Args:
+            batch_size: Size of the Record Batches which make up the Table.
+
+        Returns:
+            A pyarrow Table containing all remaining rows from the result set.
+        """
+        self._check_closed()  # Check if the cursor is closed
+        pyarrow = self._ensure_pyarrow()
+
+        batches: list["pyarrow.RecordBatch"] = []
+        while True:
+            batch = self.arrow_batch(batch_size)
+            if batch.num_rows < batch_size or batch_size <= 0:
+                if not batches or batch.num_rows > 0:
+                    batches.append(batch)
+                break
+            batches.append(batch)
+        return pyarrow.Table.from_batches(batches, schema=batches[0].schema)
+
+    def arrow_reader(self, batch_size: int = 8192) -> "pyarrow.RecordBatchReader":
+        """
+        Fetch the result as a pyarrow RecordBatchReader, which yields Record
+        Batches of the specified size until the current result set is
+        exhausted.
+
+        Args:
+            batch_size: Size of the Record Batches produced by the reader.
+
+        Returns:
+            A pyarrow RecordBatchReader for the result set.
+        """
+        self._check_closed()  # Check if the cursor is closed
+        pyarrow = self._ensure_pyarrow()
+
+        # Fetch schema without advancing cursor
+        schema_batch = self.arrow_batch(0)
+        schema = schema_batch.schema
+
+        def batch_generator():
+            while (batch := self.arrow_batch(batch_size)).num_rows > 0:
+                yield batch
+
+        return pyarrow.RecordBatchReader.from_batches(schema, batch_generator())
+
     def nextset(self) -> Union[bool, None]:
         """
         Skip to the next available result set.

mssql_python/mssql_python.pyi

@@ -7,6 +7,7 @@ Type stubs for mssql_python package - based on actual public API
 from typing import Any, Dict, List, Optional, Union, Tuple, Sequence, Callable, Iterator
 import datetime
 import logging
+import pyarrow
 
 # GLOBALS - DB-API 2.0 Required Module Globals
 # https://www.python.org/dev/peps/pep-0249/#module-interface
@@ -199,6 +200,11 @@ class Cursor:
     def setinputsizes(self, sizes: List[Union[int, Tuple[Any, ...]]]) -> None: ...
     def setoutputsize(self, size: int, column: Optional[int] = None) -> None: ...
 
+    # Arrow Extension Methods (requires pyarrow)
+    def arrow_batch(self, batch_size: int = 8192) -> pyarrow.RecordBatch: ...
+    def arrow(self, batch_size: int = 8192) -> pyarrow.Table: ...
+    def arrow_reader(self, batch_size: int = 8192) -> pyarrow.RecordBatchReader: ...
+
 # DB-API 2.0 Connection Object
 # https://www.python.org/dev/peps/pep-0249/#connection-objects
 class Connection: