@@ -47,14 +47,15 @@ mod tests {
4747 use datafusion_physical_plan:: { collect, ExecutionPlan } ;
4848
4949 use arrow:: array:: {
50- BooleanArray , Float64Array , Int32Array , RecordBatch , StringArray ,
50+ Array , BooleanArray , Float64Array , Int32Array , RecordBatch , StringArray ,
5151 } ;
5252 use arrow:: compute:: concat_batches;
5353 use arrow:: csv:: ReaderBuilder ;
5454 use arrow:: util:: pretty:: pretty_format_batches;
5555 use async_trait:: async_trait;
5656 use bytes:: Bytes ;
5757 use chrono:: DateTime ;
58+ use datafusion_common:: config:: CsvOptions ;
5859 use futures:: stream:: BoxStream ;
5960 use futures:: StreamExt ;
6061 use insta:: assert_snapshot;
@@ -1174,4 +1175,141 @@ mod tests {
11741175 . build_decoder ( ) ;
11751176 DecoderDeserializer :: new ( CsvDecoder :: new ( decoder) )
11761177 }
1178+
1179+ fn csv_deserializer_with_truncated (
1180+ batch_size : usize ,
1181+ schema : & Arc < Schema > ,
1182+ ) -> impl BatchDeserializer < Bytes > {
1183+ // using Arrow's ReaderBuilder and enabling truncated_rows
1184+ let decoder = ReaderBuilder :: new ( schema. clone ( ) )
1185+ . with_batch_size ( batch_size)
1186+ . with_truncated_rows ( true ) // <- enable runtime truncated_rows
1187+ . build_decoder ( ) ;
1188+ DecoderDeserializer :: new ( CsvDecoder :: new ( decoder) )
1189+ }
1190+
1191+ #[ tokio:: test]
1192+ async fn infer_schema_with_truncated_rows_true ( ) -> Result < ( ) > {
1193+ let session_ctx = SessionContext :: new ( ) ;
1194+ let state = session_ctx. state ( ) ;
1195+
1196+ // CSV: header has 3 columns, but first data row has only 2 columns, second row has 3
1197+ let csv_data = Bytes :: from ( "a,b,c\n 1,2\n 3,4,5\n " ) ;
1198+ let variable_object_store = Arc :: new ( VariableStream :: new ( csv_data, 1 ) ) ;
1199+ let object_meta = ObjectMeta {
1200+ location : Path :: parse ( "/" ) ?,
1201+ last_modified : DateTime :: default ( ) ,
1202+ size : u64:: MAX ,
1203+ e_tag : None ,
1204+ version : None ,
1205+ } ;
1206+
1207+ // Construct CsvFormat and enable truncated_rows via CsvOptions
1208+ let csv_options = CsvOptions :: default ( ) . with_truncated_rows ( true ) ;
1209+ let csv_format = CsvFormat :: default ( )
1210+ . with_has_header ( true )
1211+ . with_options ( csv_options)
1212+ . with_schema_infer_max_rec ( 10 ) ;
1213+
1214+ let inferred_schema = csv_format
1215+ . infer_schema (
1216+ & state,
1217+ & ( variable_object_store. clone ( ) as Arc < dyn ObjectStore > ) ,
1218+ & [ object_meta] ,
1219+ )
1220+ . await ?;
1221+
1222+ // header has 3 columns; inferred schema should also have 3
1223+ assert_eq ! ( inferred_schema. fields( ) . len( ) , 3 ) ;
1224+
1225+ // inferred columns should be nullable
1226+ for f in inferred_schema. fields ( ) {
1227+ assert ! ( f. is_nullable( ) ) ;
1228+ }
1229+
1230+ Ok ( ( ) )
1231+ }
1232+ #[ test]
1233+ fn test_decoder_truncated_rows_runtime ( ) -> Result < ( ) > {
1234+ // Synchronous test: Decoder API used here is synchronous
1235+ let schema = csv_schema ( ) ; // helper already defined in file
1236+
1237+ // Construct a decoder that enables truncated_rows at runtime
1238+ let mut deserializer = csv_deserializer_with_truncated ( 10 , & schema) ;
1239+
1240+ // Provide two rows: first row complete, second row missing last column
1241+ let input = Bytes :: from ( "0,0.0,true,0-string\n 1,1.0,true\n " ) ;
1242+ deserializer. digest ( input) ;
1243+
1244+ // Finish and collect output
1245+ deserializer. finish ( ) ;
1246+
1247+ let output = deserializer. next ( ) ?;
1248+ match output {
1249+ DeserializerOutput :: RecordBatch ( batch) => {
1250+ // ensure at least two rows present
1251+ assert ! ( batch. num_rows( ) >= 2 ) ;
1252+ // column 4 (index 3) should be a StringArray where second row is NULL
1253+ let col4 = batch
1254+ . column ( 3 )
1255+ . as_any ( )
1256+ . downcast_ref :: < StringArray > ( )
1257+ . expect ( "column 4 should be StringArray" ) ;
1258+
1259+ // first row present, second row should be null
1260+ assert ! ( !col4. is_null( 0 ) ) ;
1261+ assert ! ( col4. is_null( 1 ) ) ;
1262+ }
1263+ other => panic ! ( "expected RecordBatch but got {other:?}" ) ,
1264+ }
1265+ Ok ( ( ) )
1266+ }
1267+
1268+ #[ tokio:: test]
1269+ async fn infer_schema_truncated_rows_false_error ( ) -> Result < ( ) > {
1270+ let session_ctx = SessionContext :: new ( ) ;
1271+ let state = session_ctx. state ( ) ;
1272+
1273+ // CSV: header has 4 cols, first data row has 3 cols -> truncated at end
1274+ let csv_data = Bytes :: from ( "id,a,b,c\n 1,foo,bar\n 2,foo,bar,baz\n " ) ;
1275+ let variable_object_store = Arc :: new ( VariableStream :: new ( csv_data, 1 ) ) ;
1276+ let object_meta = ObjectMeta {
1277+ location : Path :: parse ( "/" ) ?,
1278+ last_modified : DateTime :: default ( ) ,
1279+ size : u64:: MAX ,
1280+ e_tag : None ,
1281+ version : None ,
1282+ } ;
1283+
1284+ // CsvFormat without enabling truncated_rows (default behavior = false)
1285+ let csv_format = CsvFormat :: default ( )
1286+ . with_has_header ( true )
1287+ . with_schema_infer_max_rec ( 10 ) ;
1288+
1289+ let res = csv_format
1290+ . infer_schema (
1291+ & state,
1292+ & ( variable_object_store. clone ( ) as Arc < dyn ObjectStore > ) ,
1293+ & [ object_meta] ,
1294+ )
1295+ . await ;
1296+
1297+ // Expect an error due to unequal lengths / incorrect number of fields
1298+ assert ! (
1299+ res. is_err( ) ,
1300+ "expected infer_schema to error on truncated rows when disabled"
1301+ ) ;
1302+
1303+ // Optional: check message contains indicative text (two known possibilities)
1304+ if let Err ( err) = res {
1305+ let msg = format ! ( "{err}" ) ;
1306+ assert ! (
1307+ msg. contains( "Encountered unequal lengths" )
1308+ || msg. contains( "incorrect number of fields" ) ,
1309+ "unexpected error message: {msg}" ,
1310+ ) ;
1311+ }
1312+
1313+ Ok ( ( ) )
1314+ }
11771315}
0 commit comments