FEAT: Support for Apache Arrow in Bulk Copy API#665
Draft
gargsaumya wants to merge 1 commit into
Draft
Conversation
Adds cursor.bulkcopy_arrow(table_name, source) to bulk-load directly from Apache Arrow data (pyarrow Table/RecordBatch/RecordBatchReader, __arrow_c_stream__/__arrow_c_array__ producers, or an iterable of batches) via the mssql_py_core Rust core, streaming through the Arrow C Data Interface with the GIL released during transfer. bulkcopy() now raises TypeError to steer Arrow inputs here. Auth setup is refactored into a shared _build_pycore_context() helper (preserves SQL, pre-acquired-token, and ServicePrincipal factory paths). Adds tests/test_024_bulkcopy_arrow.py (36 unit tests at 100% coverage of the new code + 16 live round-trip/type-matrix tests) and benchmarks/bench_bulkcopy_arrow.py. Requires mssql-py-core 0.1.5+.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| def test_sql_auth_keeps_credentials(self): | ||
| cur = _cursor_with_conn( | ||
| "Server=localhost;Database=testdb;UID=sa;PWD=mypwd", auth_type=None |
| @patch("mssql_python.cursor.logger") | ||
| def test_success_returns_result_and_forwards_args(self, mock_logger): | ||
| mock_logger.is_debug_enabled = False | ||
| cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") |
| @patch("mssql_python.cursor.logger") | ||
| def test_sensitive_fields_cleared_after_success(self, mock_logger): | ||
| mock_logger.is_debug_enabled = False | ||
| cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=secret") |
| @patch("mssql_python.cursor.logger") | ||
| def test_resources_closed_on_success(self, mock_logger): | ||
| mock_logger.is_debug_enabled = False | ||
| cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") |
| @patch("mssql_python.cursor.logger") | ||
| def test_core_exception_is_reraised_and_cleaned_up(self, mock_logger): | ||
| mock_logger.is_debug_enabled = False | ||
| cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") |
| def test_cleanup_swallows_close_errors(self, mock_logger): | ||
| """A failing resource.close() during teardown must not mask the result.""" | ||
| mock_logger.is_debug_enabled = False | ||
| cur = _cursor_with_conn("Server=localhost;Database=d;UID=sa;PWD=p") |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Apache Arrow support to the Bulk Copy API by introducing a dedicated Cursor.bulkcopy_arrow() path and refactoring shared connection/auth context construction so both tuple-based and Arrow-based bulk copy use consistent handling.
Changes:
- Added
Cursor.bulkcopy_arrow(table_name, source)and Arrow-source detection/steering to prevent Arrow inputs from going through tuple bulkcopy. - Refactored bulk-copy connection/auth parsing into
Cursor._build_pycore_context()for reuse across bulk copy entry points. - Added comprehensive tests and a benchmark comparing tuple vs Arrow bulk copy throughput.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_024_bulkcopy_arrow.py | Adds unit + live-DB integration coverage for Arrow bulk copy, auth context building, and Arrow-source detection. |
| mssql_python/mssql_python.pyi | Updates public type stubs to include bulkcopy and new bulkcopy_arrow API. |
| mssql_python/cursor.py | Implements _build_pycore_context, Arrow source detection/steering, and the new bulkcopy_arrow method. |
| CHANGELOG.md | Documents the new Arrow bulk copy feature and related requirements/behavior changes. |
| benchmarks/bench_bulkcopy_arrow.py | Adds a benchmark script comparing tuple-based vs Arrow-based bulk copy performance. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+3014
to
+3018
| raise TypeError( | ||
| "bulkcopy() expects an iterable of row tuples/lists. " | ||
| "For pyarrow.Table / RecordBatch / RecordBatchReader / objects " | ||
| "implementing __arrow_c_stream__, call cursor.bulkcopy_arrow() instead." | ||
| ) |
Comment on lines
+3293
to
+3300
| except Exception as e: | ||
| logger.debug( | ||
| "bulkcopy_arrow failed for table '%s': %s: %s", | ||
| table_name, | ||
| type(e).__name__, | ||
| str(e), | ||
| ) | ||
| raise type(e)(str(e)) from None |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Work Item / Issue Reference
Summary
This pull request introduces a new high-performance bulk copy method for Apache Arrow data and refactors the existing bulk copy logic to improve maintainability and error handling. It also adds a comprehensive benchmark for comparing tuple-based and Arrow-based bulk copy performance.
Summary of most important changes:
New Features
Cursor.bulkcopy_arrow(table_name, source)method for high-performance bulk loading of data directly from Apache Arrow sources. This method streams data using the Arrow C Data Interface, bypassing the need to materialize Python row objects and releasing the GIL during transfer. It supports multiple Arrow-compatible sources and is significantly faster for Arrow-originating data. The classicbulkcopy()now raises aTypeErrorif given Arrow input, steering users to the new method.Benchmarks
benchmarks/bench_bulkcopy_arrow.py, a benchmark script that compares the performance ofbulkcopy()(tuple-based) andbulkcopy_arrow()(Arrow-based) using identical data. The script supports multiple data profiles and row counts, and prints detailed throughput and speedup statistics.Code Refactoring and Error Handling
_build_pycore_context. This ensures bothbulkcopyandbulkcopy_arrowuse consistent connection and authentication handling, improving code maintainability and reducing duplication. [1] [2]bulkcopy(): now immediately raises aTypeErrorif the input is an Arrow source, with a clear message directing users tobulkcopy_arrow().