@@ -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,69 +80,3 @@ impl<'a> MetadataTable<'a> {
4380 ManifestsTable :: new ( self . 0 )
4481 }
4582}
46-
47- #[ cfg( test) ]
48- pub mod tests {
49- //! Sharable tests for the metadata table.
50-
51- use expect_test:: Expect ;
52- use futures:: TryStreamExt ;
53- use itertools:: Itertools ;
54-
55- use crate :: scan:: ArrowRecordBatchStream ;
56-
57- /// Snapshot testing to check the resulting record batch.
58- ///
59- /// - `expected_schema/data`: put `expect![[""]]` as a placeholder,
60- /// and then run test with `UPDATE_EXPECT=1 cargo test` to automatically update the result,
61- /// or use rust-analyzer (see [video](https://github.com/rust-analyzer/expect-test)).
62- /// Check the doc of [`expect_test`] for more details.
63- /// - `ignore_check_columns`: Some columns are not stable, so we can skip them.
64- /// - `sort_column`: The order of the data might be non-deterministic, so we can sort it by a column.
65- pub async fn check_record_batches (
66- batch_stream : ArrowRecordBatchStream ,
67- expected_schema : Expect ,
68- expected_data : Expect ,
69- ignore_check_columns : & [ & str ] ,
70- sort_column : Option < & str > ,
71- ) {
72- let record_batches = batch_stream. try_collect :: < Vec < _ > > ( ) . await . unwrap ( ) ;
73- assert ! ( !record_batches. is_empty( ) , "Empty record batches" ) ;
74-
75- // Combine record batches using the first batch's schema
76- let first_batch = record_batches. first ( ) . unwrap ( ) ;
77- let record_batch =
78- arrow_select:: concat:: concat_batches ( & first_batch. schema ( ) , & record_batches) . unwrap ( ) ;
79-
80- let mut columns = record_batch. columns ( ) . to_vec ( ) ;
81- if let Some ( sort_column) = sort_column {
82- let column = record_batch. column_by_name ( sort_column) . unwrap ( ) ;
83- let indices = arrow_ord:: sort:: sort_to_indices ( column, None , None ) . unwrap ( ) ;
84- columns = columns
85- . iter ( )
86- . map ( |column| arrow_select:: take:: take ( column. as_ref ( ) , & indices, None ) . unwrap ( ) )
87- . collect_vec ( ) ;
88- }
89-
90- expected_schema. assert_eq ( & format ! (
91- "{}" ,
92- record_batch. schema( ) . fields( ) . iter( ) . format( ",\n " )
93- ) ) ;
94- expected_data. assert_eq ( & format ! (
95- "{}" ,
96- record_batch
97- . schema( )
98- . fields( )
99- . iter( )
100- . zip_eq( columns)
101- . map( |( field, column) | {
102- if ignore_check_columns. contains( & field. name( ) . as_str( ) ) {
103- format!( "{}: (skipped)" , field. name( ) )
104- } else {
105- format!( "{}: {:?}" , field. name( ) , column)
106- }
107- } )
108- . format( ",\n " )
109- ) ) ;
110- }
111- }
0 commit comments