| 326 | // (Doc section: Reading and writing partitioned data #3) |
| 327 | |
| 328 | arrow::Status RunDatasetDocumentation(const std::string& format_name, |
| 329 | const std::string& uri, const std::string& mode) { |
| 330 | ARROW_RETURN_NOT_OK(arrow::compute::Initialize()); |
| 331 | |
| 332 | std::string base_path; |
| 333 | std::shared_ptr<ds::FileFormat> format; |
| 334 | std::string root_path; |
| 335 | ARROW_ASSIGN_OR_RAISE(auto fs, fs::FileSystemFromUri(uri, &root_path)); |
| 336 | |
| 337 | if (format_name == "feather") { |
| 338 | format = std::make_shared<ds::IpcFileFormat>(); |
| 339 | ARROW_ASSIGN_OR_RAISE(base_path, CreateExampleFeatherDataset(fs, root_path)); |
| 340 | } else if (format_name == "parquet") { |
| 341 | format = std::make_shared<ds::ParquetFileFormat>(); |
| 342 | ARROW_ASSIGN_OR_RAISE(base_path, CreateExampleParquetDataset(fs, root_path)); |
| 343 | } else if (format_name == "parquet_hive") { |
| 344 | format = std::make_shared<ds::ParquetFileFormat>(); |
| 345 | ARROW_ASSIGN_OR_RAISE(base_path, |
| 346 | CreateExampleParquetHivePartitionedDataset(fs, root_path)); |
| 347 | } else { |
| 348 | std::cerr << "Unknown format: " << format_name << std::endl; |
| 349 | std::cerr << "Supported formats: feather, parquet, parquet_hive" << std::endl; |
| 350 | return arrow::Status::ExecutionError("Dataset creating failed."); |
| 351 | } |
| 352 | |
| 353 | std::shared_ptr<arrow::Table> table; |
| 354 | if (mode == "no_filter") { |
| 355 | ARROW_ASSIGN_OR_RAISE(table, ScanWholeDataset(fs, format, base_path)); |
| 356 | } else if (mode == "filter") { |
| 357 | ARROW_ASSIGN_OR_RAISE(table, FilterAndSelectDataset(fs, format, base_path)); |
| 358 | } else if (mode == "project") { |
| 359 | ARROW_ASSIGN_OR_RAISE(table, ProjectDataset(fs, format, base_path)); |
| 360 | } else if (mode == "select_project") { |
| 361 | ARROW_ASSIGN_OR_RAISE(table, SelectAndProjectDataset(fs, format, base_path)); |
| 362 | } else if (mode == "partitioned") { |
| 363 | ARROW_ASSIGN_OR_RAISE(table, ScanPartitionedDataset(fs, format, base_path)); |
| 364 | } else if (mode == "filter_partitioned") { |
| 365 | ARROW_ASSIGN_OR_RAISE(table, FilterPartitionedDataset(fs, format, base_path)); |
| 366 | } else { |
| 367 | std::cerr << "Unknown mode: " << mode << std::endl; |
| 368 | std::cerr |
| 369 | << "Supported modes: no_filter, filter, project, select_project, partitioned" |
| 370 | << std::endl; |
| 371 | return arrow::Status::ExecutionError("Dataset reading failed."); |
| 372 | } |
| 373 | std::cout << "Read " << table->num_rows() << " rows" << std::endl; |
| 374 | std::cout << table->ToString() << std::endl; |
| 375 | return arrow::Status::OK(); |
| 376 | } |
| 377 | |
| 378 | int main(int argc, char** argv) { |
| 379 | if (argc < 3) { |
no test coverage detected