diff --git a/CHANGELOG.md b/CHANGELOG.md index fe59e327..1fd52a0d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this close) so it surfaces with a `SQL_SUCCESS_WITH_INFO` hint instead of being posted under `SQL_NO_DATA`, which many applications never inspect. +- `mssql-py-core`: Arrow bulk copy now accepts `Utf8View` and `BinaryView` + columns, allowing Polars DataFrames to load without first converting the + DataFrame to a PyArrow table. + - Initial public release of the mssql-rs workspace. ### Changed diff --git a/mssql-py-core/src/arrow_bulkcopy.rs b/mssql-py-core/src/arrow_bulkcopy.rs index c0204048..93c7c078 100644 --- a/mssql-py-core/src/arrow_bulkcopy.rs +++ b/mssql-py-core/src/arrow_bulkcopy.rs @@ -11,14 +11,13 @@ use std::sync::Arc; -use arrow::array::BinaryArray; use arrow::array::{ - Array, BooleanArray, Date32Array, Date64Array, Decimal128Array, FixedSizeBinaryArray, - Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, LargeBinaryArray, - LargeStringArray, RecordBatch, StringArray, Time32MillisecondArray, Time32SecondArray, - Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray, - TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt8Array, - UInt16Array, UInt32Array, UInt64Array, + Array, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, Decimal128Array, + FixedSizeBinaryArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, + Int64Array, LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StringViewArray, + Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, }; use arrow::datatypes::{DataType, Schema, TimeUnit}; use async_trait::async_trait; @@ -79,8 +78,11 @@ pub enum ColumnPlanKind { Utf8VarChar, LargeUtf8Nvarchar, LargeUtf8VarChar, + Utf8ViewNvarchar, + Utf8ViewVarChar, Binary, LargeBinary, + BinaryView, /// Arrow `date32` → SQL `date`. Date32, /// Arrow `date64` (ms) → SQL `date`. @@ -125,11 +127,11 @@ pub enum ColumnPlanKind { SmallDateTime { unit: TimeUnit, }, - /// Arrow `utf8`/`large_utf8` → SQL `xml`. + /// Arrow `utf8`/`large_utf8`/`utf8_view` → SQL `xml`. Xml, - /// Arrow `utf8`/`large_utf8` → SQL `json`. + /// Arrow `utf8`/`large_utf8`/`utf8_view` → SQL `json`. Json, - /// Arrow `utf8`/`large_utf8` GUID text → SQL `uniqueidentifier`. + /// Arrow `utf8`/`large_utf8`/`utf8_view` GUID text → SQL `uniqueidentifier`. Utf8Uuid, /// All-null arrow column. Null, @@ -207,9 +209,12 @@ fn resolve_kind(arrow_ty: &DataType, dest: &BulkCopyColumnMetadata) -> TdsResult ColumnPlanKind::LargeUtf8Nvarchar } (DataType::LargeUtf8, S::VarChar | S::Char | S::Text) => ColumnPlanKind::LargeUtf8VarChar, + (DataType::Utf8View, S::NVarChar | S::NChar | S::NText) => ColumnPlanKind::Utf8ViewNvarchar, + (DataType::Utf8View, S::VarChar | S::Char | S::Text) => ColumnPlanKind::Utf8ViewVarChar, (DataType::Binary, S::VarBinary | S::Binary | S::Image) => ColumnPlanKind::Binary, (DataType::LargeBinary, S::VarBinary | S::Binary | S::Image) => ColumnPlanKind::LargeBinary, + (DataType::BinaryView, S::VarBinary | S::Binary | S::Image) => ColumnPlanKind::BinaryView, (DataType::Date32, S::Date) => ColumnPlanKind::Date32, (DataType::Date64, S::Date) => ColumnPlanKind::Date64, @@ -254,11 +259,13 @@ fn resolve_kind(arrow_ty: &DataType, dest: &BulkCopyColumnMetadata) -> TdsResult }, (DataType::FixedSizeBinary(16), S::UniqueIdentifier) => ColumnPlanKind::FixedBin16Uuid, - // Arrow utf8/large_utf8 GUID text → uniqueidentifier. mssql-python's + // Arrow UTF-8 GUID text → uniqueidentifier. mssql-python's // cursor.arrow() reads a GUID column as a string, so this enables a // read-then-bulkload roundtrip; the binary FixedSizeBinary(16) form is // still accepted above. - (DataType::Utf8 | DataType::LargeUtf8, S::UniqueIdentifier) => ColumnPlanKind::Utf8Uuid, + (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, S::UniqueIdentifier) => { + ColumnPlanKind::Utf8Uuid + } // Any-width fixed-size binary loads into BINARY/VARBINARY/IMAGE (the // extractor reads FixedSizeBinaryArray of any width; the server enforces // the column length), mirroring the variable-width binary arm above. @@ -290,9 +297,11 @@ fn resolve_kind(arrow_ty: &DataType, dest: &BulkCopyColumnMetadata) -> TdsResult )); } - // utf8 → xml / json (Arrow already guarantees valid UTF-8). - (DataType::Utf8 | DataType::LargeUtf8, S::Xml) => ColumnPlanKind::Xml, - (DataType::Utf8 | DataType::LargeUtf8, S::Json) => ColumnPlanKind::Json, + // UTF-8 → xml / json (Arrow already guarantees valid UTF-8). + (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, S::Xml) => ColumnPlanKind::Xml, + (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, S::Json) => { + ColumnPlanKind::Json + } _ => { return Err(Error::UsageError( @@ -412,20 +421,16 @@ impl ColumnPlan { check_finite(v, dest)?; Ok(ColumnValues::Float(v)) } - ColumnPlanKind::Utf8Nvarchar => { - let s = downcast::(arr)?.value(row_idx).to_owned(); - Ok(ColumnValues::String(SqlString::from_utf8_string(s))) - } - ColumnPlanKind::Utf8VarChar => { + ColumnPlanKind::Utf8Nvarchar | ColumnPlanKind::Utf8VarChar => { let s = downcast::(arr)?.value(row_idx).to_owned(); Ok(ColumnValues::String(SqlString::from_utf8_string(s))) } - ColumnPlanKind::LargeUtf8Nvarchar => { + ColumnPlanKind::LargeUtf8Nvarchar | ColumnPlanKind::LargeUtf8VarChar => { let s = downcast::(arr)?.value(row_idx).to_owned(); Ok(ColumnValues::String(SqlString::from_utf8_string(s))) } - ColumnPlanKind::LargeUtf8VarChar => { - let s = downcast::(arr)?.value(row_idx).to_owned(); + ColumnPlanKind::Utf8ViewNvarchar | ColumnPlanKind::Utf8ViewVarChar => { + let s = downcast::(arr)?.value(row_idx).to_owned(); Ok(ColumnValues::String(SqlString::from_utf8_string(s))) } ColumnPlanKind::Binary => { @@ -442,6 +447,10 @@ impl ColumnPlan { let bytes = downcast::(arr)?.value(row_idx).to_vec(); Ok(ColumnValues::Bytes(bytes)) } + ColumnPlanKind::BinaryView => { + let bytes = downcast::(arr)?.value(row_idx).to_vec(); + Ok(ColumnValues::Bytes(bytes)) + } ColumnPlanKind::Date32 => { let days = downcast::(arr)?.value(row_idx); let sql_days = days @@ -662,14 +671,19 @@ fn timestamp_to_1900_components(ticks_since_epoch: i64) -> (i64, u8, u8, u8, u32 (days_since_1900, hour, minute, second, microsecond) } -/// Read a UTF-8 string cell from either a `StringArray` or `LargeStringArray`. +/// Read a UTF-8 string cell from any supported variable-width string array. fn read_utf8(arr: &dyn Array, row_idx: usize) -> TdsResult<&str> { if let Some(a) = arr.as_any().downcast_ref::() { Ok(a.value(row_idx)) } else if let Some(a) = arr.as_any().downcast_ref::() { Ok(a.value(row_idx)) + } else if let Some(a) = arr.as_any().downcast_ref::() { + Ok(a.value(row_idx)) } else { - Err(downcast_err::(arr)) + Err(Error::UsageError(format!( + "Arrow array downcast failed: expected a UTF-8 string array (utf8/large_utf8/utf8_view), got {:?}", + arr.data_type() + ))) } } @@ -1165,6 +1179,55 @@ mod tests { } } + #[test] + fn utf8_view_varchar_and_nvarchar() { + let external = "this value exceeds twelve bytes"; + let array: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("inline"), + Some(external), + None, + ])); + for sql_type in [SqlDbType::VarChar, SqlDbType::NVarChar] { + let dest = meta("name", sql_type, true); + let plan = one_col_plan(DataType::Utf8View, &dest); + match plan.extract_value(array.as_ref(), 0, &dest).unwrap() { + ColumnValues::String(value) => assert_eq!(value.to_utf8_string(), "inline"), + other => panic!("expected String, got {other:?}"), + } + match plan.extract_value(array.as_ref(), 1, &dest).unwrap() { + ColumnValues::String(value) => assert_eq!(value.to_utf8_string(), external), + other => panic!("expected String, got {other:?}"), + } + assert_eq!( + plan.extract_value(array.as_ref(), 2, &dest).unwrap(), + ColumnValues::Null + ); + } + } + + #[test] + fn utf8_view_special_text_destinations() { + let guid = "58185e0d-3a91-44d8-bc46-7107217e0a6d"; + let guid_array: ArrayRef = Arc::new(StringViewArray::from(vec![Some(guid)])); + let guid_dest = meta("guid", SqlDbType::UniqueIdentifier, true); + let guid_value = one_col_plan(DataType::Utf8View, &guid_dest) + .extract_value(guid_array.as_ref(), 0, &guid_dest) + .unwrap(); + assert!(matches!(guid_value, ColumnValues::Uuid(_))); + + for (sql_type, text) in [(SqlDbType::Xml, ""), (SqlDbType::Json, "{\"a\":1}")] { + let array: ArrayRef = Arc::new(StringViewArray::from(vec![Some(text)])); + let dest = meta("document", sql_type, true); + match one_col_plan(DataType::Utf8View, &dest) + .extract_value(array.as_ref(), 0, &dest) + .unwrap() + { + ColumnValues::Xml(_) | ColumnValues::Json(_) => {} + other => panic!("expected Xml/Json, got {other:?}"), + } + } + } + #[test] fn binary_and_large_binary() { let dest = meta("b", SqlDbType::VarBinary, true); @@ -1182,6 +1245,26 @@ mod tests { .unwrap(), ColumnValues::Bytes(vec![10, 11]) ); + + let external = b"more than twelve bytes"; + let view: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some(&b"\x0c\x0d"[..]), + Some(&external[..]), + None, + ])); + let plan = one_col_plan(DataType::BinaryView, &dest); + assert_eq!( + plan.extract_value(view.as_ref(), 0, &dest).unwrap(), + ColumnValues::Bytes(vec![12, 13]) + ); + assert_eq!( + plan.extract_value(view.as_ref(), 1, &dest).unwrap(), + ColumnValues::Bytes(external.to_vec()) + ); + assert_eq!( + plan.extract_value(view.as_ref(), 2, &dest).unwrap(), + ColumnValues::Null + ); } #[test] diff --git a/mssql-py-core/tests/test_bulkcopy_arrow_binary.py b/mssql-py-core/tests/test_bulkcopy_arrow_binary.py index c1d33f1a..81366add 100644 --- a/mssql-py-core/tests/test_bulkcopy_arrow_binary.py +++ b/mssql-py-core/tests/test_bulkcopy_arrow_binary.py @@ -113,6 +113,37 @@ def test_cursor_bulkcopy_arrow_binary_auto_mapping(client_context): conn.close() +@pytest.mark.integration +def test_cursor_bulkcopy_arrow_varbinary_binary_view(client_context): + """Arrow binary_view bulkcopy preserves values and NULLs.""" + conn = mssql_py_core.PyCoreConnection(client_context) + cursor = conn.cursor() + + table_name = "#BulkCopyArrowVarBinaryView" + cursor.execute(f"CREATE TABLE {table_name} (id INT, data VARBINARY(50))") + source = pa.table( + { + "id": pa.array([1, 2, 3], type=pa.int32()), + "data": pa.array( + [b"\x01\x02", None, b"more than twelve bytes"], + type=pa.binary_view(), + ), + } + ) + + result = cursor.bulkcopy_arrow(table_name, source, batch_size=1000, timeout=30) + assert result["rows_copied"] == 3 + + cursor.execute(f"SELECT id, data FROM {table_name} ORDER BY id") + assert cursor.fetchall() == [ + (1, b"\x01\x02"), + (2, None), + (3, b"more than twelve bytes"), + ] + + conn.close() + + @pytest.mark.integration def test_cursor_bulkcopy_arrow_binary_null_to_non_nullable_column(client_context): """A NULL value into a non-nullable VARBINARY column must raise ValueError.""" diff --git a/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py b/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py index dce652cb..99535597 100644 --- a/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py +++ b/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py @@ -83,6 +83,37 @@ def test_cursor_bulkcopy_arrow_varchar_auto_mapping(client_context): conn.close() +@pytest.mark.integration +def test_cursor_bulkcopy_arrow_varchar_string_view(client_context): + """Arrow string_view bulkcopy supports Polars-style C-stream columns.""" + conn = mssql_py_core.PyCoreConnection(client_context) + cursor = conn.cursor() + + table_name = "#BulkCopyArrowVarcharStringView" + cursor.execute(f"CREATE TABLE {table_name} (id INT, name VARCHAR(50))") + source = pa.table( + { + "id": pa.array([1, 2, 3], type=pa.int32()), + "name": pa.array( + ["inline", None, "this value exceeds twelve bytes"], + type=pa.string_view(), + ), + } + ) + + result = cursor.bulkcopy_arrow(table_name, source, batch_size=1000, timeout=30) + assert result["rows_copied"] == 3 + + cursor.execute(f"SELECT id, name FROM {table_name} ORDER BY id") + assert cursor.fetchall() == [ + (1, "inline"), + (2, None), + (3, "this value exceeds twelve bytes"), + ] + + conn.close() + + @pytest.mark.integration def test_cursor_bulkcopy_arrow_varchar_max_large(client_context): """Arrow utf8 with a large string round-trips through VARCHAR(MAX)."""