@@ -27,6 +27,43 @@ use crate::table::Table;
2727#[ derive( Debug ) ]
2828pub struct MetadataTable < ' a > ( & ' a Table ) ;
2929
30+ /// Metadata table type.
31+ #[ derive( Debug , Clone , strum:: EnumIter ) ]
32+ pub enum MetadataTableType {
33+ /// [`SnapshotsTable`]
34+ Snapshots ,
35+ /// [`ManifestsTable`]
36+ Manifests ,
37+ }
38+
39+ impl MetadataTableType {
40+ /// Returns the string representation of the metadata table type.
41+ pub fn as_str ( & self ) -> & str {
42+ match self {
43+ MetadataTableType :: Snapshots => "snapshots" ,
44+ MetadataTableType :: Manifests => "manifests" ,
45+ }
46+ }
47+
48+ /// Returns all the metadata table types.
49+ pub fn all_types ( ) -> impl Iterator < Item = Self > {
50+ use strum:: IntoEnumIterator ;
51+ Self :: iter ( )
52+ }
53+ }
54+
55+ impl TryFrom < & str > for MetadataTableType {
56+ type Error = String ;
57+
58+ fn try_from ( value : & str ) -> std:: result:: Result < Self , String > {
59+ match value {
60+ "snapshots" => Ok ( Self :: Snapshots ) ,
61+ "manifests" => Ok ( Self :: Manifests ) ,
62+ _ => Err ( format ! ( "invalid metadata table type: {value}" ) ) ,
63+ }
64+ }
65+ }
66+
3067impl < ' a > MetadataTable < ' a > {
3168 /// Creates a new metadata scan.
3269pub fn new ( table : & ' a Table ) -> Self {
@@ -43,67 +80,3 @@ impl<'a> MetadataTable<'a> {
4380 ManifestsTable :: new ( self . 0 )
4481 }
4582}
46-
47- #[ cfg( test) ]
48- pub mod tests {
49- use expect_test:: Expect ;
50- use futures:: TryStreamExt ;
51- use itertools:: Itertools ;
52-
53- use crate :: scan:: ArrowRecordBatchStream ;
54-
55- /// Snapshot testing to check the resulting record batch.
56- ///
57- /// - `expected_schema/data`: put `expect![[""]]` as a placeholder,
58- /// and then run test with `UPDATE_EXPECT=1 cargo test` to automatically update the result,
59- /// or use rust-analyzer (see [video](https://github.com/rust-analyzer/expect-test)).
60- /// Check the doc of [`expect_test`] for more details.
61- /// - `ignore_check_columns`: Some columns are not stable, so we can skip them.
62- /// - `sort_column`: The order of the data might be non-deterministic, so we can sort it by a column.
63- pub async fn check_record_batches (
64- batch_stream : ArrowRecordBatchStream ,
65- expected_schema : Expect ,
66- expected_data : Expect ,
67- ignore_check_columns : & [ & str ] ,
68- sort_column : Option < & str > ,
69- ) {
70- let record_batches = batch_stream. try_collect :: < Vec < _ > > ( ) . await . unwrap ( ) ;
71- assert ! ( !record_batches. is_empty( ) , "Empty record batches" ) ;
72-
73- // Combine record batches using the first batch's schema
74- let first_batch = record_batches. first ( ) . unwrap ( ) ;
75- let record_batch =
76- arrow_select:: concat:: concat_batches ( & first_batch. schema ( ) , & record_batches) . unwrap ( ) ;
77-
78- let mut columns = record_batch. columns ( ) . to_vec ( ) ;
79- if let Some ( sort_column) = sort_column {
80- let column = record_batch. column_by_name ( sort_column) . unwrap ( ) ;
81- let indices = arrow_ord:: sort:: sort_to_indices ( column, None , None ) . unwrap ( ) ;
82- columns = columns
83- . iter ( )
84- . map ( |column| arrow_select:: take:: take ( column. as_ref ( ) , & indices, None ) . unwrap ( ) )
85- . collect_vec ( ) ;
86- }
87-
88- expected_schema. assert_eq ( & format ! (
89- "{}" ,
90- record_batch. schema( ) . fields( ) . iter( ) . format( ",\n " )
91- ) ) ;
92- expected_data. assert_eq ( & format ! (
93- "{}" ,
94- record_batch
95- . schema( )
96- . fields( )
97- . iter( )
98- . zip_eq( columns)
99- . map( |( field, column) | {
100- if ignore_check_columns. contains( & field. name( ) . as_str( ) ) {
101- format!( "{}: (skipped)" , field. name( ) )
102- } else {
103- format!( "{}: {:?}" , field. name( ) , column)
104- }
105- } )
106- . format( ",\n " )
107- ) ) ;
108- }
109- }
0 commit comments