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