| 1252 | } |
| 1253 | |
| 1254 | Result<std::shared_ptr<Table>> FileReaderImpl::ReadRowGroups( |
| 1255 | const std::vector<int>& row_groups, const std::vector<int>& column_indices) { |
| 1256 | RETURN_NOT_OK(BoundsCheck(row_groups, column_indices)); |
| 1257 | |
| 1258 | // PARQUET-1698/PARQUET-1820: pre-buffer row groups/column chunks if enabled |
| 1259 | if (reader_properties_.pre_buffer()) { |
| 1260 | BEGIN_PARQUET_CATCH_EXCEPTIONS |
| 1261 | parquet_reader()->PreBuffer(row_groups, column_indices, |
| 1262 | reader_properties_.io_context(), |
| 1263 | reader_properties_.cache_options()); |
| 1264 | END_PARQUET_CATCH_EXCEPTIONS |
| 1265 | } |
| 1266 | |
| 1267 | auto fut = DecodeRowGroups(/*self=*/nullptr, row_groups, column_indices, |
| 1268 | /*cpu_executor=*/nullptr); |
| 1269 | return fut.MoveResult(); |
| 1270 | } |
| 1271 | |
| 1272 | Future<std::shared_ptr<Table>> FileReaderImpl::DecodeRowGroups( |
| 1273 | std::shared_ptr<FileReaderImpl> self, const std::vector<int>& row_groups, |
| 1274 | const std::vector<int>& column_indices, ::arrow::internal::Executor* cpu_executor) { |
| 1275 | // `self` is used solely to keep `this` alive in an async context - but we use this |
| 1276 | // in a sync context too so use `this` over `self` |
| 1277 | std::vector<std::shared_ptr<ColumnReaderImpl>> readers; |
| 1278 | std::shared_ptr<::arrow::Schema> result_schema; |
| 1279 | RETURN_NOT_OK(GetFieldReaders(column_indices, row_groups, &readers, &result_schema)); |
| 1280 | // OptionalParallelForAsync requires an executor |
| 1281 | if (!cpu_executor) cpu_executor = ::arrow::internal::GetCpuThreadPool(); |
| 1282 | |
| 1283 | auto read_column = [row_groups, self, this](size_t i, |
| 1284 | std::shared_ptr<ColumnReaderImpl> reader) |
| 1285 | -> ::arrow::Result<std::shared_ptr<::arrow::ChunkedArray>> { |
| 1286 | std::shared_ptr<::arrow::ChunkedArray> column; |
| 1287 | RETURN_NOT_OK(ReadColumn(static_cast<int>(i), row_groups, reader.get(), &column)); |
| 1288 | return column; |
| 1289 | }; |
| 1290 | auto make_table = [result_schema, row_groups, self, |
| 1291 | this](const ::arrow::ChunkedArrayVector& columns) |
| 1292 | -> ::arrow::Result<std::shared_ptr<Table>> { |
| 1293 | int64_t num_rows = 0; |
| 1294 | if (!columns.empty()) { |
| 1295 | num_rows = columns[0]->length(); |
| 1296 | } else { |
| 1297 | for (int i : row_groups) { |
| 1298 | num_rows += parquet_reader()->metadata()->RowGroup(i)->num_rows(); |
| 1299 | } |
| 1300 | } |
| 1301 | auto table = Table::Make(std::move(result_schema), columns, num_rows); |
| 1302 | RETURN_NOT_OK(table->Validate()); |
| 1303 | return table; |
| 1304 | }; |
| 1305 | return ::arrow::internal::OptionalParallelForAsync(reader_properties_.use_threads(), |
| 1306 | std::move(readers), read_column, |
| 1307 | cpu_executor) |
| 1308 | .Then(std::move(make_table)); |
| 1309 | } |
| 1310 | |
| 1311 | std::shared_ptr<RowGroupReader> FileReaderImpl::RowGroup(int row_group_index) { |
no test coverage detected