@@ -115,3 +115,132 @@ fn test_metadata_field() {
115115 "Expected METADATA FROM option in column definition"
116116 ) ;
117117}
118+
119+ #[ test]
120+ fn test_iceberg_partitioned_by ( ) {
121+ let sql = "CREATE TABLE ice (
122+ ts TIMESTAMP NOT NULL,
123+ id INT NOT NULL,
124+ favorite_color TEXT
125+ ) WITH (
126+ connector = 'iceberg',
127+ format = 'parquet',
128+ table_name = 'arroyo_test'
129+ ) PARTITIONED BY (
130+ hour(ts),
131+ bucket(32, id),
132+ truncate(8, favorite_color)
133+ )" ;
134+
135+ let parse = Parser :: parse_sql ( & ArroyoDialect { } , sql) . unwrap ( ) ;
136+ let Statement :: CreateTable ( ct) = parse. get ( 0 ) . unwrap ( ) else {
137+ panic ! ( "not create table" )
138+ } ;
139+
140+ // Verify basic structure
141+ assert_eq ! ( ct. name. to_string( ) , "ice" ) ;
142+ assert_eq ! ( ct. columns. len( ) , 3 ) ;
143+
144+ // Verify arroyo_partitions is present
145+ let partitions = ct. arroyo_partitions . as_ref ( ) . expect ( "Expected arroyo_partitions to be Some" ) ;
146+ assert_eq ! ( partitions. len( ) , 3 ) ;
147+
148+ // Check each partition transform
149+ // hour(ts)
150+ match & partitions[ 0 ] {
151+ Expr :: Function ( f) => {
152+ assert_eq ! ( f. name. to_string( ) , "hour" ) ;
153+ if let sqlparser:: ast:: FunctionArguments :: List ( list) = & f. args {
154+ assert_eq ! ( list. args. len( ) , 1 ) ;
155+ } else {
156+ panic ! ( "Expected List arguments" ) ;
157+ }
158+ }
159+ _ => panic ! ( "Expected Function for hour(ts)" ) ,
160+ }
161+
162+ // bucket(32, id)
163+ match & partitions[ 1 ] {
164+ Expr :: Function ( f) => {
165+ assert_eq ! ( f. name. to_string( ) , "bucket" ) ;
166+ if let sqlparser:: ast:: FunctionArguments :: List ( list) = & f. args {
167+ assert_eq ! ( list. args. len( ) , 2 ) ;
168+ } else {
169+ panic ! ( "Expected List arguments" ) ;
170+ }
171+ }
172+ _ => panic ! ( "Expected Function for bucket(32, id)" ) ,
173+ }
174+
175+ // truncate(8, favorite_color)
176+ match & partitions[ 2 ] {
177+ Expr :: Function ( f) => {
178+ assert_eq ! ( f. name. to_string( ) , "truncate" ) ;
179+ if let sqlparser:: ast:: FunctionArguments :: List ( list) = & f. args {
180+ assert_eq ! ( list. args. len( ) , 2 ) ;
181+ } else {
182+ panic ! ( "Expected List arguments" ) ;
183+ }
184+ }
185+ _ => panic ! ( "Expected Function for truncate(8, favorite_color)" ) ,
186+ }
187+
188+ // Test round-trip: the formatted output should parse back to the same structure
189+ let formatted = ct. to_string ( ) ;
190+ let reparsed = Parser :: parse_sql ( & ArroyoDialect { } , & formatted) . unwrap ( ) ;
191+ let Statement :: CreateTable ( ct2) = reparsed. get ( 0 ) . unwrap ( ) else {
192+ panic ! ( "not create table on reparse" )
193+ } ;
194+
195+ assert_eq ! ( ct. arroyo_partitions, ct2. arroyo_partitions) ;
196+ }
197+
198+ #[ test]
199+ fn test_iceberg_partitioned_by_single ( ) {
200+ let sql = "CREATE TABLE events (
201+ event_time TIMESTAMP
202+ ) WITH (
203+ connector = 'iceberg'
204+ ) PARTITIONED BY (day(event_time))" ;
205+
206+ let parse = Parser :: parse_sql ( & ArroyoDialect { } , sql) . unwrap ( ) ;
207+ let Statement :: CreateTable ( ct) = parse. get ( 0 ) . unwrap ( ) else {
208+ panic ! ( "not create table" )
209+ } ;
210+
211+ let partitions = ct. arroyo_partitions . as_ref ( ) . expect ( "Expected arroyo_partitions" ) ;
212+ assert_eq ! ( partitions. len( ) , 1 ) ;
213+
214+ match & partitions[ 0 ] {
215+ Expr :: Function ( f) => {
216+ assert_eq ! ( f. name. to_string( ) , "day" ) ;
217+ }
218+ _ => panic ! ( "Expected Function for day(event_time)" ) ,
219+ }
220+ }
221+
222+ #[ test]
223+ fn test_iceberg_partitioned_by_identity ( ) {
224+ // Test identity transform (just a column name)
225+ let sql = "CREATE TABLE data (
226+ region TEXT,
227+ value INT
228+ ) WITH (
229+ connector = 'iceberg'
230+ ) PARTITIONED BY (region)" ;
231+
232+ let parse = Parser :: parse_sql ( & ArroyoDialect { } , sql) . unwrap ( ) ;
233+ let Statement :: CreateTable ( ct) = parse. get ( 0 ) . unwrap ( ) else {
234+ panic ! ( "not create table" )
235+ } ;
236+
237+ let partitions = ct. arroyo_partitions . as_ref ( ) . expect ( "Expected arroyo_partitions" ) ;
238+ assert_eq ! ( partitions. len( ) , 1 ) ;
239+
240+ match & partitions[ 0 ] {
241+ Expr :: Identifier ( ident) => {
242+ assert_eq ! ( ident. value, "region" ) ;
243+ }
244+ _ => panic ! ( "Expected Identifier for region" ) ,
245+ }
246+ }
0 commit comments