@@ -144,16 +144,19 @@ impl ParquetFile {
144
144
let mut row_count = 0 ;
145
145
for i in 0 ..self . column_readers . len ( ) {
146
146
let array: Arc < Array > = match self . column_readers [ i] {
147
- ColumnReader :: BoolColumnReader ( ref mut r ) => {
147
+ ColumnReader :: BoolColumnReader ( ref mut _r ) => {
148
148
return Err ( ExecutionError :: NotImplemented (
149
149
"unsupported column reader type (BOOL)" . to_string ( ) ,
150
150
) ) ;
151
151
}
152
152
ColumnReader :: Int32ColumnReader ( ref mut r) => {
153
- let mut builder = Int32Builder :: new ( self . batch_size ) ;
154
153
let mut read_buffer: Vec < i32 > =
155
154
Vec :: with_capacity ( self . batch_size ) ;
156
155
156
+ for _ in 0 ..self . batch_size {
157
+ read_buffer. push ( 0 ) ;
158
+ }
159
+
157
160
match r. read_batch (
158
161
self . batch_size ,
159
162
None ,
@@ -163,7 +166,7 @@ impl ParquetFile {
163
166
//TODO this isn't handling null values
164
167
Ok ( ( count, _) ) => {
165
168
println ! ( "Read {} rows" , count) ;
166
-
169
+ let mut builder = Int32Builder :: new ( self . batch_size ) ;
167
170
builder. append_slice ( & read_buffer) . unwrap ( ) ;
168
171
row_count = count;
169
172
Arc :: new ( builder. finish ( ) )
@@ -176,17 +179,17 @@ impl ParquetFile {
176
179
}
177
180
}
178
181
}
179
- ColumnReader :: Int64ColumnReader ( ref mut r ) => {
182
+ ColumnReader :: Int64ColumnReader ( ref mut _r ) => {
180
183
return Err ( ExecutionError :: NotImplemented (
181
184
"unsupported column reader type (INT64)" . to_string ( ) ,
182
185
) ) ;
183
186
}
184
- ColumnReader :: Int96ColumnReader ( ref mut r ) => {
187
+ ColumnReader :: Int96ColumnReader ( ref mut _r ) => {
185
188
return Err ( ExecutionError :: NotImplemented (
186
189
"unsupported column reader type (INT96)" . to_string ( ) ,
187
190
) ) ;
188
191
}
189
- ColumnReader :: FloatColumnReader ( ref mut r ) => {
192
+ ColumnReader :: FloatColumnReader ( ref mut _r ) => {
190
193
return Err ( ExecutionError :: NotImplemented (
191
194
"unsupported column reader type (FLOAT)" . to_string ( ) ,
192
195
) ) ;
@@ -215,7 +218,7 @@ impl ParquetFile {
215
218
}
216
219
}
217
220
}
218
- ColumnReader :: FixedLenByteArrayColumnReader ( ref mut r ) => {
221
+ ColumnReader :: FixedLenByteArrayColumnReader ( ref mut _r ) => {
219
222
return Err ( ExecutionError :: NotImplemented (
220
223
"unsupported column reader type (FixedLenByteArray)"
221
224
. to_string ( ) ,
@@ -261,7 +264,15 @@ impl ParquetFile {
261
264
if row_count == 0 {
262
265
Ok ( None )
263
266
} else {
264
- Ok ( Some ( RecordBatch :: try_new ( self . schema . clone ( ) , batch) ?) )
267
+ match & self . projection {
268
+ Some ( proj) => Ok ( Some ( RecordBatch :: try_new (
269
+ self . schema . projection ( proj) ?,
270
+ batch,
271
+ ) ?) ) ,
272
+ None => {
273
+ Ok ( Some ( RecordBatch :: try_new ( self . schema . clone ( ) , batch) ?) )
274
+ }
275
+ }
265
276
}
266
277
}
267
278
_ => Ok ( None ) ,
@@ -337,10 +348,11 @@ impl RecordBatchIterator for ParquetFile {
337
348
#[ cfg( test) ]
338
349
mod tests {
339
350
use super :: * ;
351
+ use arrow:: array:: Int32Array ;
340
352
use std:: env;
341
353
342
354
#[ test]
343
- fn read_parquet_file ( ) {
355
+ fn read_read_i32_column ( ) {
344
356
let testdata = env:: var ( "PARQUET_TEST_DATA" ) . unwrap ( ) ;
345
357
let filename = format ! ( "{}/alltypes_plain.parquet" , testdata) ;
346
358
@@ -354,6 +366,21 @@ mod tests {
354
366
let batch = it. next ( ) . unwrap ( ) . unwrap ( ) ;
355
367
356
368
assert_eq ! ( 1 , batch. num_columns( ) ) ;
357
- assert_eq ! ( 1 , batch. num_rows( ) ) ;
369
+ assert_eq ! ( 64 * 1024 , batch. num_rows( ) ) ;
370
+
371
+ let array = batch
372
+ . column ( 0 )
373
+ . as_any ( )
374
+ . downcast_ref :: < Int32Array > ( )
375
+ . unwrap ( ) ;
376
+ let mut values: Vec < i32 > = vec ! [ ] ;
377
+ for i in 0 ..16 {
378
+ values. push ( array. value ( i) ) ;
379
+ }
380
+
381
+ assert_eq ! (
382
+ "[4, 5, 6, 7, 2, 3, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0]" ,
383
+ format!( "{:?}" , values)
384
+ ) ;
358
385
}
359
386
}
0 commit comments