@@ -82,6 +82,7 @@ pub fn basic_parse() {
82
82
#[ cfg( test) ]
83
83
mod test {
84
84
use super :: * ;
85
+ use datafusion:: execution:: options:: ParquetReadOptions ;
85
86
use datafusion:: {
86
87
arrow:: {
87
88
array:: { ArrayRef , Int32Array , RecordBatch , StringArray } ,
@@ -90,12 +91,16 @@ mod test {
90
91
datasource:: MemTable ,
91
92
execution:: context:: SessionContext ,
92
93
} ;
94
+ use datafusion_common:: test_util:: batches_to_string;
93
95
use datafusion_execution:: {
94
96
config:: SessionConfig , disk_manager:: DiskManagerConfig ,
95
97
runtime_env:: RuntimeEnvBuilder ,
96
98
} ;
97
99
use datafusion_physical_plan:: collect;
98
100
use datafusion_sql:: parser:: DFParser ;
101
+ use insta:: assert_snapshot;
102
+ use object_store:: { memory:: InMemory , path:: Path , ObjectStore } ;
103
+ use url:: Url ;
99
104
use wasm_bindgen_test:: wasm_bindgen_test;
100
105
101
106
wasm_bindgen_test:: wasm_bindgen_test_configure!( run_in_browser) ;
@@ -115,6 +120,22 @@ mod test {
115
120
let session_config = SessionConfig :: new ( ) . with_target_partitions ( 1 ) ;
116
121
Arc :: new ( SessionContext :: new_with_config_rt ( session_config, rt) )
117
122
}
123
+
124
+ fn create_test_data ( ) -> ( Arc < Schema > , RecordBatch ) {
125
+ let schema = Arc :: new ( Schema :: new ( vec ! [
126
+ Field :: new( "id" , DataType :: Int32 , false ) ,
127
+ Field :: new( "value" , DataType :: Utf8 , false ) ,
128
+ ] ) ) ;
129
+
130
+ let data: Vec < ArrayRef > = vec ! [
131
+ Arc :: new( Int32Array :: from( vec![ 1 , 2 , 3 ] ) ) ,
132
+ Arc :: new( StringArray :: from( vec![ "a" , "b" , "c" ] ) ) ,
133
+ ] ;
134
+
135
+ let batch = RecordBatch :: try_new ( schema. clone ( ) , data) . unwrap ( ) ;
136
+ ( schema, batch)
137
+ }
138
+
118
139
#[ wasm_bindgen_test( unsupported = tokio:: test) ]
119
140
async fn basic_execute ( ) {
120
141
let sql = "SELECT 2 + 2;" ;
@@ -185,26 +206,56 @@ mod test {
185
206
186
207
#[ wasm_bindgen_test( unsupported = tokio:: test) ]
187
208
async fn test_parquet_write ( ) {
188
- let schema = Arc :: new ( Schema :: new ( vec ! [
189
- Field :: new( "id" , DataType :: Int32 , false ) ,
190
- Field :: new( "value" , DataType :: Utf8 , false ) ,
191
- ] ) ) ;
209
+ let ( schema, batch) = create_test_data ( ) ;
210
+ let mut buffer = Vec :: new ( ) ;
211
+ let mut writer = datafusion:: parquet:: arrow:: ArrowWriter :: try_new (
212
+ & mut buffer,
213
+ schema. clone ( ) ,
214
+ None ,
215
+ )
216
+ . unwrap ( ) ;
192
217
193
- let data: Vec < ArrayRef > = vec ! [
194
- Arc :: new( Int32Array :: from( vec![ 1 ] ) ) ,
195
- Arc :: new( StringArray :: from( vec![ "a" ] ) ) ,
196
- ] ;
218
+ writer. write ( & batch) . unwrap ( ) ;
219
+ writer. close ( ) . unwrap ( ) ;
220
+ }
197
221
198
- let batch = RecordBatch :: try_new ( schema. clone ( ) , data) . unwrap ( ) ;
222
+ #[ wasm_bindgen_test( unsupported = tokio:: test) ]
223
+ async fn test_parquet_read_and_write ( ) {
224
+ let ( schema, batch) = create_test_data ( ) ;
199
225
let mut buffer = Vec :: new ( ) ;
200
226
let mut writer = datafusion:: parquet:: arrow:: ArrowWriter :: try_new (
201
227
& mut buffer,
202
228
schema. clone ( ) ,
203
229
None ,
204
230
)
205
231
. unwrap ( ) ;
206
-
207
232
writer. write ( & batch) . unwrap ( ) ;
208
233
writer. close ( ) . unwrap ( ) ;
234
+
235
+ let session_ctx = SessionContext :: new ( ) ;
236
+ let store = InMemory :: new ( ) ;
237
+
238
+ let path = Path :: from ( "a.parquet" ) ;
239
+ store. put ( & path, buffer. into ( ) ) . await . unwrap ( ) ;
240
+
241
+ let url = Url :: parse ( "memory://" ) . unwrap ( ) ;
242
+ session_ctx. register_object_store ( & url, Arc :: new ( store) ) ;
243
+
244
+ let df = session_ctx
245
+ . read_parquet ( "memory:///" , ParquetReadOptions :: new ( ) )
246
+ . await
247
+ . unwrap ( ) ;
248
+
249
+ let result = df. collect ( ) . await . unwrap ( ) ;
250
+
251
+ assert_snapshot ! ( batches_to_string( & result) , @r"
252
+ +----+-------+
253
+ | id | value |
254
+ +----+-------+
255
+ | 1 | a |
256
+ | 2 | b |
257
+ | 3 | c |
258
+ +----+-------+
259
+ " ) ;
209
260
}
210
261
}
0 commit comments