Skip to content

Commit d5251c1

Browse files
fix: VARIANT follow-ups for SchemaTransform etc (delta-io#1106)
## What changes are proposed in this pull request? Follow-up from VARIANT work in delta-io#1015. chiefly: (1) modify `transform_variant` to take/return structtypes, and default to recursing through children, and (2) track a couple follow ups in TODOs. the actual bug that must be fixed is that when a VARIANT went in to `transform_variant`, it came back out as a regular STRUCT if transformed, since DataType::from on a struct is just a struct (note that if not transformed, we return the original data_type which is still variant) original comment kept below for posterity ### Original exploration in earlier commits Originally prompted as follow-up from delta-io#1015, this serves as an exploration of introducing a new `VariantType` struct (similar to `MapType`/`ArrayType`/etc.) which just wraps `StructType`. The original question pertaining to `SchemaTransform` was basically: can we implement `transform_variant` to take+return a `StructType` instead of `DataType`? This loses the information that says the thing is a variant and not a normal struct. This prompted the exploration of how we actually want to represent Variant in our schemas. In general we need a way of distinguishing variants from plain structs I think we have two options: we either rely on `DataType` which communicates `DataType::Variant(struct_type)` OR we introduce a new `StructType` much like MapType/ArrayType/etc. This PR serves as exploration of the former - though i'm realizing if we _do_ take `DataType` as the thing that 'tags' a struct as variant then it's probably also okay? And would just have the schema transform still retain taking/returning DataType? (or perhaps could consume a StructType but just always return a DataType) I remember @nicklan had even more/better thoughts around this - please comment and we can discuss :) ## How was this change tested? Existing UT
1 parent 9325ec6 commit d5251c1

4 files changed

Lines changed: 22 additions & 11 deletions

File tree

kernel/src/schema/mod.rs

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -874,9 +874,10 @@ pub trait SchemaTransform<'a> {
874874
self.transform(etype)
875875
}
876876

877-
/// Called for each variant values encountered. By default does nothing
878-
fn transform_variant(&mut self, etype: &'a DataType) -> Option<Cow<'a, DataType>> {
879-
Some(Cow::Borrowed(etype))
877+
/// Called for each variant value encountered. By default, recurses into the fields of the
878+
/// variant struct type.
879+
fn transform_variant(&mut self, stype: &'a StructType) -> Option<Cow<'a, StructType>> {
880+
self.recurse_into_struct(stype)
880881
}
881882

882883
/// General entry point for a recursive traversal over any data type. Also invoked internally to
@@ -896,9 +897,9 @@ pub trait SchemaTransform<'a> {
896897
Map(mtype) => self
897898
.transform_map(mtype)?
898899
.map_owned_or_else(data_type, DataType::from),
899-
Variant(_) => self
900-
.transform_variant(data_type)?
901-
.map_owned_or_else(data_type, DataType::from),
900+
Variant(stype) => self
901+
.transform_variant(stype)?
902+
.map_owned_or_else(data_type, |s| DataType::Variant(Box::new(s))),
902903
};
903904
Some(result)
904905
}

kernel/src/schema/variant_utils.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! Utility functions for the variant type and variant-related table features.
22
33
use crate::actions::Protocol;
4-
use crate::schema::{DataType, Schema, SchemaTransform};
4+
use crate::schema::{Schema, SchemaTransform, StructType};
55
use crate::table_features::{ReaderFeature, WriterFeature};
66
use crate::utils::require;
77
use crate::{DeltaResult, Error};
@@ -12,7 +12,7 @@ use std::borrow::Cow;
1212
pub(crate) struct UsesVariant(pub(crate) bool);
1313

1414
impl<'a> SchemaTransform<'a> for UsesVariant {
15-
fn transform_variant(&mut self, _: &'a DataType) -> Option<Cow<'a, DataType>> {
15+
fn transform_variant(&mut self, _: &'a StructType) -> Option<Cow<'a, StructType>> {
1616
self.0 = true;
1717
None
1818
}

kernel/src/table_features/column_mapping.rs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
//! Code to handle column mapping, including modes and schema transforms
22
use super::ReaderFeature;
33
use crate::actions::Protocol;
4-
use crate::schema::{ColumnName, DataType, MetadataValue, Schema, SchemaTransform, StructField};
4+
use crate::schema::{
5+
ColumnName, DataType, MetadataValue, Schema, SchemaTransform, StructField, StructType,
6+
};
57
use crate::table_properties::TableProperties;
68
use crate::{DeltaResult, Error};
79

@@ -155,6 +157,12 @@ impl<'a> SchemaTransform<'a> for ValidateColumnMappings<'a> {
155157
}
156158
None
157159
}
160+
fn transform_variant(&mut self, _: &'a StructType) -> Option<Cow<'a, StructType>> {
161+
// don't recurse into variant's fields, as they are not expected to have column mapping
162+
// annotations
163+
// TODO: this changes with icebergcompat right? see issue#1125 for icebergcompat.
164+
None
165+
}
158166
}
159167

160168
#[cfg(test)]

test-utils/src/lib.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,10 @@ pub async fn create_table(
248248
}
249249
if enable_column_mapping {
250250
reader_features.push("columnMapping");
251+
// TODO: (#1124) we don't actually support column mapping writes yet, but have some
252+
// tests that do column mapping on writes. for now omit the writer feature to let tests
253+
// run, but after actual support this should be enabled.
254+
// writer_features.push("columnMapping");
251255
}
252256
(reader_features, writer_features)
253257
};
@@ -357,7 +361,6 @@ pub fn to_arrow(data: Box<dyn EngineData>) -> DeltaResult<RecordBatch> {
357361
.into())
358362
}
359363

360-
// TODO (zach): this is listed as unused for acceptance crate
361364
pub fn read_scan(scan: &Scan, engine: Arc<dyn Engine>) -> DeltaResult<Vec<RecordBatch>> {
362365
let scan_results = scan.execute(engine)?;
363366
scan_results
@@ -375,7 +378,6 @@ pub fn read_scan(scan: &Scan, engine: Arc<dyn Engine>) -> DeltaResult<Vec<Record
375378
.try_collect()
376379
}
377380

378-
// TODO (zach): this is listed as unused for acceptance crate
379381
pub fn test_read(
380382
expected: &ArrowEngineData,
381383
url: &Url,

0 commit comments

Comments
 (0)