From 194e6ab9fff2da56428fdc27d971f738ef18ddb2 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Thu, 13 Aug 2026 10:58:25 +0530 Subject: [PATCH 1/4] Support Arrow view types in bulk copy --- CHANGELOG.md | 3 + mssql-py-core/src/arrow_bulkcopy.rs | 99 ++++++++++++++----- .../tests/test_bulkcopy_arrow_binary.py | 24 +++++ .../tests/test_bulkcopy_arrow_varchar.py | 24 +++++ 4 files changed, 127 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fe59e327..450eeebc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -62,6 +62,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this ### Fixed +- `mssql-py-core`: Arrow bulk copy now accepts `Utf8View` and `BinaryView` + columns, enabling zero-copy string ingestion from Polars DataFrames. + - `mssql-tds`: reading a fixed-width value that straddles a TDS packet boundary could return bytes from the wrong place or panic. The readers checked for sufficient buffered data with an `if` and read a single further packet, but a diff --git a/mssql-py-core/src/arrow_bulkcopy.rs b/mssql-py-core/src/arrow_bulkcopy.rs index c0204048..8bfa359b 100644 --- a/mssql-py-core/src/arrow_bulkcopy.rs +++ b/mssql-py-core/src/arrow_bulkcopy.rs @@ -11,15 +11,15 @@ 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, + LargeStringArray, RecordBatch, StringArray, StringViewArray, Time32MillisecondArray, + Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, }; +use arrow::array::{BinaryArray, BinaryViewArray}; use arrow::datatypes::{DataType, Schema, TimeUnit}; use async_trait::async_trait; use mssql_tds::connection::bulk_copy::{BulkLoadRow, ResolvedColumnMapping}; @@ -79,8 +79,11 @@ pub enum ColumnPlanKind { Utf8VarChar, LargeUtf8Nvarchar, LargeUtf8VarChar, + Utf8ViewNvarchar, + Utf8ViewVarChar, Binary, LargeBinary, + BinaryView, /// Arrow `date32` → SQL `date`. Date32, /// Arrow `date64` (ms) → SQL `date`. @@ -207,9 +210,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 +260,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 +298,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 +422,13 @@ 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 => { - let s = downcast::(arr)?.value(row_idx).to_owned(); - Ok(ColumnValues::String(SqlString::from_utf8_string(s))) - } - ColumnPlanKind::LargeUtf8Nvarchar => { - 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::Utf8Nvarchar + | ColumnPlanKind::Utf8VarChar + | ColumnPlanKind::LargeUtf8Nvarchar + | ColumnPlanKind::LargeUtf8VarChar + | ColumnPlanKind::Utf8ViewNvarchar + | ColumnPlanKind::Utf8ViewVarChar => { + let s = read_utf8(arr, row_idx)?.to_owned(); Ok(ColumnValues::String(SqlString::from_utf8_string(s))) } ColumnPlanKind::Binary => { @@ -442,6 +445,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,12 +669,14 @@ 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)) } @@ -1165,6 +1174,42 @@ mod tests { } } + #[test] + fn utf8_view_varchar_and_nvarchar() { + let array: ArrayRef = Arc::new(StringViewArray::from(vec![Some("hello"), 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(), "hello"), + other => panic!("expected String, got {other:?}"), + } + assert_eq!( + plan.extract_value(array.as_ref(), 1, &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); + one_col_plan(DataType::Utf8View, &dest) + .extract_value(array.as_ref(), 0, &dest) + .unwrap(); + } + } + #[test] fn binary_and_large_binary() { let dest = meta("b", SqlDbType::VarBinary, true); @@ -1182,6 +1227,14 @@ mod tests { .unwrap(), ColumnValues::Bytes(vec![10, 11]) ); + + let view: ArrayRef = Arc::new(BinaryViewArray::from(vec![Some(&b"\x0c\x0d"[..])])); + assert_eq!( + one_col_plan(DataType::BinaryView, &dest) + .extract_value(view.as_ref(), 0, &dest) + .unwrap(), + ColumnValues::Bytes(vec![12, 13]) + ); } #[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..3ba097c5 100644 --- a/mssql-py-core/tests/test_bulkcopy_arrow_binary.py +++ b/mssql-py-core/tests/test_bulkcopy_arrow_binary.py @@ -113,6 +113,30 @@ 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"\x0a\x0b"], 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"\x0a\x0b")] + + 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..18394ef5 100644 --- a/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py +++ b/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py @@ -83,6 +83,30 @@ 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(["alpha", None, "gamma"], 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, "alpha"), (2, None), (3, "gamma")] + + 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).""" From 218678c4acb16cde95a27522b88dca7dd69556ca Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Thu, 13 Aug 2026 11:41:01 +0530 Subject: [PATCH 2/4] Strengthen Arrow view type coverage --- CHANGELOG.md | 3 +- mssql-py-core/src/arrow_bulkcopy.rs | 41 ++++++++++++++----- .../tests/test_bulkcopy_arrow_binary.py | 11 ++++- .../tests/test_bulkcopy_arrow_varchar.py | 11 ++++- 4 files changed, 51 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 450eeebc..2dc3f92f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,7 +63,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this ### Fixed - `mssql-py-core`: Arrow bulk copy now accepts `Utf8View` and `BinaryView` - columns, enabling zero-copy string ingestion from Polars DataFrames. + columns, allowing Polars DataFrames to load without first converting their + string columns to a PyArrow table. - `mssql-tds`: reading a fixed-width value that straddles a TDS packet boundary could return bytes from the wrong place or panic. The readers checked for diff --git a/mssql-py-core/src/arrow_bulkcopy.rs b/mssql-py-core/src/arrow_bulkcopy.rs index 8bfa359b..2f4e106a 100644 --- a/mssql-py-core/src/arrow_bulkcopy.rs +++ b/mssql-py-core/src/arrow_bulkcopy.rs @@ -128,11 +128,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, @@ -1176,16 +1176,25 @@ mod tests { #[test] fn utf8_view_varchar_and_nvarchar() { - let array: ArrayRef = Arc::new(StringViewArray::from(vec![Some("hello"), None])); + 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(), "hello"), + 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(), 1, &dest).unwrap(), + plan.extract_value(array.as_ref(), 2, &dest).unwrap(), ColumnValues::Null ); } @@ -1228,13 +1237,25 @@ mod tests { ColumnValues::Bytes(vec![10, 11]) ); - let view: ArrayRef = Arc::new(BinaryViewArray::from(vec![Some(&b"\x0c\x0d"[..])])); + 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!( - one_col_plan(DataType::BinaryView, &dest) - .extract_value(view.as_ref(), 0, &dest) - .unwrap(), + 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 3ba097c5..81366add 100644 --- a/mssql-py-core/tests/test_bulkcopy_arrow_binary.py +++ b/mssql-py-core/tests/test_bulkcopy_arrow_binary.py @@ -124,7 +124,10 @@ def test_cursor_bulkcopy_arrow_varbinary_binary_view(client_context): source = pa.table( { "id": pa.array([1, 2, 3], type=pa.int32()), - "data": pa.array([b"\x01\x02", None, b"\x0a\x0b"], type=pa.binary_view()), + "data": pa.array( + [b"\x01\x02", None, b"more than twelve bytes"], + type=pa.binary_view(), + ), } ) @@ -132,7 +135,11 @@ def test_cursor_bulkcopy_arrow_varbinary_binary_view(client_context): 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"\x0a\x0b")] + assert cursor.fetchall() == [ + (1, b"\x01\x02"), + (2, None), + (3, b"more than twelve bytes"), + ] conn.close() diff --git a/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py b/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py index 18394ef5..99535597 100644 --- a/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py +++ b/mssql-py-core/tests/test_bulkcopy_arrow_varchar.py @@ -94,7 +94,10 @@ def test_cursor_bulkcopy_arrow_varchar_string_view(client_context): source = pa.table( { "id": pa.array([1, 2, 3], type=pa.int32()), - "name": pa.array(["alpha", None, "gamma"], type=pa.string_view()), + "name": pa.array( + ["inline", None, "this value exceeds twelve bytes"], + type=pa.string_view(), + ), } ) @@ -102,7 +105,11 @@ def test_cursor_bulkcopy_arrow_varchar_string_view(client_context): assert result["rows_copied"] == 3 cursor.execute(f"SELECT id, name FROM {table_name} ORDER BY id") - assert cursor.fetchall() == [(1, "alpha"), (2, None), (3, "gamma")] + assert cursor.fetchall() == [ + (1, "inline"), + (2, None), + (3, "this value exceeds twelve bytes"), + ] conn.close() From 496d0c8ee78cf130e2a5ed904a8fcefa0ca07e41 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Thu, 13 Aug 2026 12:25:50 +0530 Subject: [PATCH 3/4] Clarify Polars conversion wording --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2dc3f92f..cb1565bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,8 +63,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this ### Fixed - `mssql-py-core`: Arrow bulk copy now accepts `Utf8View` and `BinaryView` - columns, allowing Polars DataFrames to load without first converting their - string columns to a PyArrow table. + columns, allowing Polars DataFrames to load without first converting the + DataFrame to a PyArrow table. - `mssql-tds`: reading a fixed-width value that straddles a TDS packet boundary could return bytes from the wrong place or panic. The readers checked for From c9305031bc23a7ef86024a86fafb7e910e3fcce3 Mon Sep 17 00:00:00 2001 From: gargsaumya Date: Fri, 14 Aug 2026 10:57:00 +0530 Subject: [PATCH 4/4] Address review feedback for Arrow view bulk copy --- CHANGELOG.md | 8 +++--- mssql-py-core/src/arrow_bulkcopy.rs | 43 +++++++++++++++++------------ 2 files changed, 30 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cb1565bb..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 @@ -62,10 +66,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this ### Fixed -- `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. - - `mssql-tds`: reading a fixed-width value that straddles a TDS packet boundary could return bytes from the wrong place or panic. The readers checked for sufficient buffered data with an `if` and read a single further packet, but a diff --git a/mssql-py-core/src/arrow_bulkcopy.rs b/mssql-py-core/src/arrow_bulkcopy.rs index 2f4e106a..93c7c078 100644 --- a/mssql-py-core/src/arrow_bulkcopy.rs +++ b/mssql-py-core/src/arrow_bulkcopy.rs @@ -12,14 +12,13 @@ use std::sync::Arc; use arrow::array::{ - Array, 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, + 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::array::{BinaryArray, BinaryViewArray}; use arrow::datatypes::{DataType, Schema, TimeUnit}; use async_trait::async_trait; use mssql_tds::connection::bulk_copy::{BulkLoadRow, ResolvedColumnMapping}; @@ -422,13 +421,16 @@ impl ColumnPlan { check_finite(v, dest)?; Ok(ColumnValues::Float(v)) } - ColumnPlanKind::Utf8Nvarchar - | ColumnPlanKind::Utf8VarChar - | ColumnPlanKind::LargeUtf8Nvarchar - | ColumnPlanKind::LargeUtf8VarChar - | ColumnPlanKind::Utf8ViewNvarchar - | ColumnPlanKind::Utf8ViewVarChar => { - let s = read_utf8(arr, row_idx)?.to_owned(); + ColumnPlanKind::Utf8Nvarchar | ColumnPlanKind::Utf8VarChar => { + let s = downcast::(arr)?.value(row_idx).to_owned(); + Ok(ColumnValues::String(SqlString::from_utf8_string(s))) + } + ColumnPlanKind::LargeUtf8Nvarchar | ColumnPlanKind::LargeUtf8VarChar => { + let s = downcast::(arr)?.value(row_idx).to_owned(); + Ok(ColumnValues::String(SqlString::from_utf8_string(s))) + } + ColumnPlanKind::Utf8ViewNvarchar | ColumnPlanKind::Utf8ViewVarChar => { + let s = downcast::(arr)?.value(row_idx).to_owned(); Ok(ColumnValues::String(SqlString::from_utf8_string(s))) } ColumnPlanKind::Binary => { @@ -678,7 +680,10 @@ fn read_utf8(arr: &dyn Array, row_idx: usize) -> TdsResult<&str> { } 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() + ))) } } @@ -1213,9 +1218,13 @@ mod tests { 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); - one_col_plan(DataType::Utf8View, &dest) + match one_col_plan(DataType::Utf8View, &dest) .extract_value(array.as_ref(), 0, &dest) - .unwrap(); + .unwrap() + { + ColumnValues::Xml(_) | ColumnValues::Json(_) => {} + other => panic!("expected Xml/Json, got {other:?}"), + } } }