arrow_sql_server/schema/
table_mapping.rs

1//! Bidirectional Arrow/MSSQL schema mapping.
2//!
3//! The initial mapping function starts from an Arrow schema because the first
4//! operation is Arrow-to-SQL Server writing. The resulting `SchemaMapping`
5//! values keep Arrow field metadata and MSSQL column metadata as peer concepts
6//! so future SQL Server-to-Arrow read planning can reuse the shared
7//! representation instead of inheriting a write-only column model.
8
9use arrow_schema::{Field, Schema};
10
11use crate::observability::schema::SchemaPlanningTrace;
12use crate::schema::type_conversion::plan_arrow_data_type_as_mssql_type;
13use crate::write::PlanOptions;
14use crate::{
15    ArrowFieldRef, Diagnostic, DiagnosticCode, DiagnosticSet, FieldRef, Identifier, MssqlColumn,
16    MssqlProfile, PlanOutcome, Result, SchemaMapping, TableName, create_table_sql,
17};
18
19/// Planned Arrow/MSSQL table schema for one SQL Server profile.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct PlannedSchema {
22    profile: MssqlProfile,
23    plan_options: PlanOptions,
24    mappings: Vec<SchemaMapping>,
25}
26
27impl PlannedSchema {
28    /// Creates a planned schema.
29    pub fn new(
30        profile: MssqlProfile,
31        plan_options: PlanOptions,
32        mappings: Vec<SchemaMapping>,
33    ) -> Self {
34        Self {
35            profile,
36            plan_options,
37            mappings,
38        }
39    }
40
41    /// Returns the SQL Server profile used for planning.
42    pub const fn profile(&self) -> MssqlProfile {
43        self.profile
44    }
45
46    /// Returns the conversion policies used for planning.
47    pub const fn plan_options(&self) -> PlanOptions {
48        self.plan_options
49    }
50
51    /// Returns planned column mappings.
52    pub fn mappings(&self) -> &[SchemaMapping] {
53        &self.mappings
54    }
55
56    /// Consumes the planned schema into its mappings.
57    pub fn into_mappings(self) -> Vec<SchemaMapping> {
58        self.mappings
59    }
60}
61
62impl AsRef<[SchemaMapping]> for PlannedSchema {
63    fn as_ref(&self) -> &[SchemaMapping] {
64        self.mappings()
65    }
66}
67
68impl PlanOutcome<PlannedSchema> {
69    /// Returns planned column mappings.
70    pub fn mappings(&self) -> &[SchemaMapping] {
71        self.value().mappings()
72    }
73}
74
75/// Plans an Arrow schema into a profile-bound MSSQL schema.
76pub fn plan_arrow_schema_to_mssql_schema(
77    schema: impl AsRef<Schema>,
78    profile: MssqlProfile,
79    options: PlanOptions,
80) -> Result<PlanOutcome<PlannedSchema>> {
81    let schema = schema.as_ref();
82    let field_count = schema.fields().len();
83    let trace = SchemaPlanningTrace::start(field_count, profile, options);
84
85    trace.trace_planned_schema_result(plan_arrow_schema_to_mssql_schema_inner(
86        schema, profile, &options,
87    ))
88}
89
90/// Plans Arrow/MSSQL column mappings from an Arrow schema.
91#[cfg(test)]
92pub(crate) fn plan_arrow_schema_to_mssql_mappings(
93    schema: impl AsRef<Schema>,
94    profile: MssqlProfile,
95    options: PlanOptions,
96) -> Result<PlanOutcome<Vec<SchemaMapping>>> {
97    let outcome = plan_arrow_schema_to_mssql_schema(schema, profile, options)?;
98    let (planned_schema, diagnostics) = outcome.into_parts();
99
100    Ok(PlanOutcome::new(
101        planned_schema.into_mappings(),
102        diagnostics,
103    ))
104}
105
106fn plan_arrow_schema_to_mssql_schema_inner(
107    schema: &Schema,
108    profile: MssqlProfile,
109    options: &PlanOptions,
110) -> Result<PlanOutcome<PlannedSchema>> {
111    let mut mappings = Vec::with_capacity(schema.fields().len());
112    let mut diagnostics = DiagnosticSet::new();
113
114    for (index, field) in schema.fields().iter().enumerate() {
115        match plan_arrow_field_to_mssql_column_mapping(index, field, options) {
116            Ok(mapping) => mappings.push(mapping),
117            Err(diagnostic) => diagnostics.push(diagnostic),
118        }
119    }
120
121    if diagnostics.has_errors() {
122        return Err(crate::Error::Planning { diagnostics });
123    }
124
125    Ok(PlanOutcome::new(
126        PlannedSchema::new(profile, *options, mappings),
127        diagnostics,
128    ))
129}
130
131/// Returns the planned MSSQL columns in mapping order.
132pub fn mssql_columns_from_mappings(mappings: impl AsRef<[SchemaMapping]>) -> Vec<MssqlColumn> {
133    let mappings = mappings.as_ref();
134
135    mappings
136        .iter()
137        .map(|mapping| mapping.mssql().clone())
138        .collect()
139}
140
141/// Renders deterministic `CREATE TABLE` SQL from mapping metadata.
142pub fn create_table_sql_from_mappings(
143    table: &TableName,
144    mappings: impl AsRef<[SchemaMapping]>,
145) -> String {
146    create_table_sql(
147        table,
148        &mssql_columns_from_mappings(mappings.as_ref()),
149        crate::CreateTableOptions,
150    )
151}
152
153fn plan_arrow_field_to_mssql_column_mapping(
154    index: usize,
155    field: &Field,
156    options: &PlanOptions,
157) -> std::result::Result<SchemaMapping, Diagnostic> {
158    let name = Identifier::new(field.name()).map_err(|err| {
159        Diagnostic::error(DiagnosticCode::IdentifierInvalid, err.to_string())
160            .with_field(FieldRef::new(index, field.name()))
161    })?;
162
163    let ty = plan_arrow_data_type_as_mssql_type(index, field, options)?;
164
165    let arrow = ArrowFieldRef::new(
166        index,
167        field.name().clone(),
168        field.is_nullable(),
169        field.data_type().clone(),
170    );
171    let mssql = MssqlColumn::new(name, ty, field.is_nullable());
172
173    Ok(SchemaMapping::new(arrow, mssql))
174}
175
176#[cfg(test)]
177mod tests {
178    use std::sync::Arc;
179
180    use crate::{
181        DiagnosticCode, Error, MssqlProfile, MssqlType, PlanOptions, TableName,
182        create_table_sql_from_mappings, mssql_columns_from_mappings,
183        plan_arrow_schema_to_mssql_mappings, plan_arrow_schema_to_mssql_schema,
184    };
185    use arrow_schema::{DataType, Field, Schema, UnionFields, UnionMode};
186
187    #[test]
188    fn plans_boolean_and_int32_mappings() {
189        let schema = Arc::new(Schema::new(vec![
190            Field::new("is_active", DataType::Boolean, false),
191            Field::new("quantity", DataType::Int32, true),
192        ]));
193
194        let outcome = plan_arrow_schema_to_mssql_mappings(
195            Arc::clone(&schema),
196            MssqlProfile::sql_server_2016_compat_100(),
197            PlanOptions::default(),
198        )
199        .unwrap();
200        let mappings = outcome.value();
201
202        assert_eq!(mappings.len(), 2);
203
204        let is_active = &mappings[0];
205        assert_eq!(is_active.arrow().index(), 0);
206        assert_eq!(is_active.arrow().name(), "is_active");
207        assert_eq!(is_active.arrow().data_type(), &DataType::Boolean);
208        assert!(!is_active.arrow().nullable());
209        assert_eq!(is_active.mssql().name().quoted_sql(), "[is_active]");
210        assert!(!is_active.mssql().nullable());
211        assert_eq!(is_active.mssql().ty(), &MssqlType::Bit);
212
213        let quantity = &mappings[1];
214        assert_eq!(quantity.arrow().index(), 1);
215        assert_eq!(quantity.arrow().name(), "quantity");
216        assert_eq!(quantity.arrow().data_type(), &DataType::Int32);
217        assert!(quantity.arrow().nullable());
218        assert_eq!(quantity.mssql().name().quoted_sql(), "[quantity]");
219        assert!(quantity.mssql().nullable());
220        assert_eq!(quantity.mssql().ty(), &MssqlType::Int);
221    }
222
223    #[test]
224    fn renders_create_table_sql_from_mssql_side() {
225        let schema = Schema::new(vec![
226            Field::new("is_active", DataType::Boolean, false),
227            Field::new("quantity", DataType::Int32, true),
228        ]);
229        let outcome = plan_arrow_schema_to_mssql_mappings(
230            Arc::new(schema),
231            MssqlProfile::sql_server_2016_compat_100(),
232            PlanOptions::default(),
233        )
234        .unwrap();
235        let table = TableName::new("dbo", "target").unwrap();
236
237        let sql = create_table_sql_from_mappings(&table, outcome.value());
238
239        assert_eq!(
240            sql,
241            "CREATE TABLE [dbo].[target] (\n    [is_active] bit NOT NULL,\n    [quantity] int NULL\n);"
242        );
243    }
244
245    #[test]
246    fn exposes_mssql_columns_without_arrow_identity() {
247        let schema = Schema::new(vec![Field::new("is_active", DataType::Boolean, false)]);
248        let outcome = plan_arrow_schema_to_mssql_mappings(
249            Arc::new(schema),
250            MssqlProfile::sql_server_2016_compat_100(),
251            PlanOptions::default(),
252        )
253        .unwrap();
254
255        let columns = mssql_columns_from_mappings(outcome.value());
256
257        assert_eq!(columns.len(), 1);
258        assert_eq!(columns[0].name().as_str(), "is_active");
259        assert_eq!(columns[0].ty(), &MssqlType::Bit);
260        assert!(!columns[0].nullable());
261    }
262
263    #[test]
264    fn planned_schema_preserves_profile() {
265        let profile = MssqlProfile::sql_server_2017_compat_140();
266        let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
267        let options = PlanOptions::default();
268        let outcome =
269            plan_arrow_schema_to_mssql_schema(Arc::new(schema), profile, options).unwrap();
270        let planned_schema = outcome.value();
271
272        assert_eq!(planned_schema.profile(), profile);
273        assert_eq!(planned_schema.plan_options(), options);
274        assert_eq!(planned_schema.mappings().len(), 1);
275        assert_eq!(planned_schema.mappings()[0].mssql().ty(), &MssqlType::Int);
276    }
277
278    #[test]
279    fn profile_method_plans_schema() {
280        let profile = MssqlProfile::sql_server_2017_compat_100();
281        let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
282        let outcome = profile
283            .plan_arrow_schema(Arc::new(schema), PlanOptions::default())
284            .unwrap();
285
286        assert_eq!(outcome.value().profile(), profile);
287        assert_eq!(outcome.mappings().len(), 1);
288    }
289
290    #[test]
291    fn unsupported_nested_and_encoded_types_collect_schema_order_diagnostics() {
292        let union_fields = UnionFields::try_new(
293            [1_i8, 2],
294            [
295                Field::new("left", DataType::Int32, true),
296                Field::new("right", DataType::Utf8, true),
297            ],
298        )
299        .unwrap();
300        let schema = Schema::new(vec![
301            Field::new("ok", DataType::Int32, false),
302            Field::new("list_col", DataType::new_list(DataType::Int64, true), true),
303            Field::new(
304                "struct_col",
305                DataType::Struct(
306                    vec![Field::new("child", DataType::Boolean, true)]
307                        .into_iter()
308                        .collect(),
309                ),
310                true,
311            ),
312            Field::new(
313                "union_col",
314                DataType::Union(union_fields, UnionMode::Sparse),
315                true,
316            ),
317            Field::new(
318                "run_end_col",
319                DataType::RunEndEncoded(
320                    Arc::new(Field::new("run_ends", DataType::Int32, false)),
321                    Arc::new(Field::new("values", DataType::Utf8, true)),
322                ),
323                true,
324            ),
325        ]);
326
327        let err = plan_arrow_schema_to_mssql_mappings(
328            Arc::new(schema),
329            MssqlProfile::sql_server_2016_compat_100(),
330            PlanOptions::default(),
331        )
332        .expect_err("unsupported fields should produce diagnostics");
333
334        let Error::Planning { diagnostics } = err else {
335            panic!("expected planning error");
336        };
337
338        assert_eq!(diagnostics.len(), 4);
339        assert!(
340            diagnostics
341                .all()
342                .iter()
343                .all(|diagnostic| diagnostic.code() == DiagnosticCode::UnsupportedArrowType)
344        );
345
346        let field_refs = diagnostics
347            .all()
348            .iter()
349            .map(|diagnostic| {
350                let field = diagnostic.field().unwrap();
351                (field.index(), field.name())
352            })
353            .collect::<Vec<_>>();
354
355        assert_eq!(
356            field_refs,
357            vec![
358                (1, "list_col"),
359                (2, "struct_col"),
360                (3, "union_col"),
361                (4, "run_end_col"),
362            ]
363        );
364
365        let messages = diagnostics
366            .all()
367            .iter()
368            .map(crate::Diagnostic::message)
369            .collect::<Vec<_>>();
370        assert!(messages[0].contains("nested"));
371        assert!(messages[1].contains("nested"));
372        assert!(messages[2].contains("nested"));
373        assert!(messages[3].contains("encoded"));
374    }
375
376    #[test]
377    fn invalid_identifier_returns_structured_planning_diagnostic() {
378        let schema = Schema::new(vec![Field::new("", DataType::Boolean, false)]);
379
380        let err = plan_arrow_schema_to_mssql_mappings(
381            Arc::new(schema),
382            MssqlProfile::sql_server_2016_compat_100(),
383            PlanOptions::default(),
384        )
385        .expect_err("empty field name should be rejected");
386
387        let Error::Planning { diagnostics } = err else {
388            panic!("expected planning error");
389        };
390
391        assert_eq!(diagnostics.len(), 1);
392
393        let diagnostic = &diagnostics.all()[0];
394        assert_eq!(diagnostic.code(), DiagnosticCode::IdentifierInvalid);
395        assert_eq!(diagnostic.field().unwrap().index(), 0);
396        assert_eq!(diagnostic.field().unwrap().name(), "");
397    }
398}
arrow_sql_server/schema/table_mapping.rs

arrow_sql_server/schema/
table_mapping.rs
