| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| |
| #include <iostream> |
| #include <memory> |
| #include "graphar/writer_util.h" |
| #ifdef ARROW_ORC |
| #include "arrow/adapters/orc/adapter.h" |
| #endif |
| #include <arrow/compute/api.h> |
| #include "arrow/api.h" |
| #include "arrow/csv/api.h" |
| #include "arrow/dataset/api.h" |
| #include "parquet/arrow/reader.h" |
| #if defined(ARROW_VERSION) && ARROW_VERSION <= 12000000 |
| #include "arrow/dataset/file_json.h" |
| #endif |
| #include "arrow/filesystem/api.h" |
| #include "arrow/filesystem/s3fs.h" |
| #include "arrow/ipc/writer.h" |
| #include "parquet/arrow/writer.h" |
| #include "simple-uri-parser/uri_parser.h" |
| |
| #include "graphar/expression.h" |
| #include "graphar/filesystem.h" |
| #include "graphar/fwd.h" |
| #include "graphar/general_params.h" |
| |
| namespace graphar::detail { |
| template <typename U, typename T> |
| static Status CastToLargeOffsetArray( |
| const std::shared_ptr<arrow::Array>& in, |
| const std::shared_ptr<arrow::DataType>& to_type, |
| std::shared_ptr<arrow::Array>& out) { // NOLINT(runtime/references) |
| auto array_data = in->data()->Copy(); |
| auto offset = array_data->buffers[1]; |
| using from_offset_type = typename U::offset_type; |
| using to_string_offset_type = typename T::offset_type; |
| auto raw_value_offsets_ = |
| offset == NULLPTR |
| ? NULLPTR |
| : reinterpret_cast<const from_offset_type*>(offset->data()); |
| std::vector<to_string_offset_type> to_offset(offset->size() / |
| sizeof(from_offset_type)); |
| for (size_t i = 0; i < to_offset.size(); ++i) { |
| to_offset[i] = raw_value_offsets_[i]; |
| } |
| std::shared_ptr<arrow::Buffer> buffer; |
| arrow::TypedBufferBuilder<to_string_offset_type> buffer_builder; |
| RETURN_NOT_ARROW_OK( |
| buffer_builder.Append(to_offset.data(), to_offset.size())); |
| RETURN_NOT_ARROW_OK(buffer_builder.Finish(&buffer)); |
| array_data->type = to_type; |
| array_data->buffers[1] = buffer; |
| out = arrow::MakeArray(array_data); |
| RETURN_NOT_ARROW_OK(out->ValidateFull()); |
| return Status::OK(); |
| } |
| |
| template <typename U, typename T> |
| static Status CastToLargeOffsetArray( |
| const std::shared_ptr<arrow::ChunkedArray>& in, |
| const std::shared_ptr<arrow::DataType>& to_type, |
| std::shared_ptr<arrow::ChunkedArray>& out) { // NOLINT(runtime/references) |
| std::vector<std::shared_ptr<arrow::Array>> chunks; |
| for (auto const& chunk : in->chunks()) { |
| std::shared_ptr<arrow::Array> array; |
| auto status = CastToLargeOffsetArray<U, T>(chunk, to_type, array); |
| GAR_RETURN_NOT_OK(status); |
| chunks.emplace_back(array); |
| } |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(out, arrow::ChunkedArray::Make(chunks)); |
| return Status::OK(); |
| } |
| } // namespace graphar::detail |
| |
| namespace graphar { |
| namespace ds = arrow::dataset; |
| |
| std::shared_ptr<ds::FileFormat> FileSystem::GetFileFormat( |
| const FileType type) const { |
| switch (type) { |
| case CSV: |
| return std::make_shared<ds::CsvFileFormat>(); |
| case PARQUET: |
| return std::make_shared<ds::ParquetFileFormat>(); |
| case JSON: |
| return std::make_shared<ds::JsonFileFormat>(); |
| #ifdef ARROW_ORC |
| case ORC: |
| return std::make_shared<ds::OrcFileFormat>(); |
| #endif |
| default: |
| return nullptr; |
| } |
| } |
| |
| Result<std::shared_ptr<arrow::Table>> FileSystem::ReadFileToTable( |
| const std::string& path, FileType file_type, |
| const std::vector<int>& column_indices) const noexcept { |
| parquet::arrow::FileReaderBuilder builder; |
| auto open_file_status = builder.OpenFile(path); |
| if (!open_file_status.ok()) { |
| return Status::Invalid("Failed to open file: ", path, " - ", |
| open_file_status.ToString()); |
| } |
| builder.memory_pool(arrow::default_memory_pool()); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto reader, builder.Build()); |
| std::shared_ptr<arrow::Table> table; |
| if (column_indices.empty()) { |
| arrow::Status read_status = reader->ReadTable(&table); |
| if (!read_status.ok()) { |
| return Status::Invalid("Failed to read table from file: ", path, " - ", |
| read_status.ToString()); |
| } |
| } else { |
| arrow::Status read_status = reader->ReadTable(column_indices, &table); |
| if (!read_status.ok()) { |
| return Status::Invalid("Failed to read table from file: ", path, " - ", |
| read_status.ToString()); |
| } |
| } |
| return table; |
| } |
| |
| Result<std::shared_ptr<arrow::Table>> FileSystem::ReadFileToTable( |
| const std::string& path, FileType file_type, |
| const util::FilterOptions& options) const noexcept { |
| std::shared_ptr<ds::FileFormat> format = GetFileFormat(file_type); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| auto factory, arrow::dataset::FileSystemDatasetFactory::Make( |
| arrow_fs_, {path}, format, |
| arrow::dataset::FileSystemFactoryOptions())); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto dataset, factory->Finish()); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto scan_builder, dataset->NewScan()); |
| #if ARROW_VERSION >= 21000000 |
| RETURN_NOT_ARROW_OK(arrow::compute::Initialize()); |
| #endif |
| // Apply the row filter and select the specified columns |
| if (options.filter) { |
| GAR_ASSIGN_OR_RAISE(auto filter, options.filter->Evaluate()); |
| RETURN_NOT_ARROW_OK(scan_builder->Filter(filter)); |
| } |
| if (options.columns) { |
| RETURN_NOT_ARROW_OK(scan_builder->Project(*options.columns)); |
| } |
| |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto scanner, scan_builder->Finish()); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto table, scanner->ToTable()); |
| // cast string array to large string array as we need concatenate chunks in |
| // some places, e.g., in vineyard |
| for (int i = 0; i < table->num_columns(); ++i) { |
| std::shared_ptr<arrow::DataType> type = table->column(i)->type(); |
| if (type->id() == arrow::Type::STRING) { |
| type = arrow::large_utf8(); |
| } else if (type->id() == arrow::Type::BINARY) { |
| type = arrow::large_binary(); |
| } |
| if (type->Equals(table->column(i)->type())) { |
| continue; |
| } |
| // do casting |
| auto field = table->field(i)->WithType(type); |
| std::shared_ptr<arrow::ChunkedArray> chunked_array; |
| |
| if (table->num_rows() == 0) { |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| chunked_array, arrow::ChunkedArray::MakeEmpty(type)); |
| } else if (type->Equals(arrow::large_utf8())) { |
| auto status = detail::CastToLargeOffsetArray<arrow::StringArray, |
| arrow::LargeStringArray>( |
| table->column(i), type, chunked_array); |
| GAR_RETURN_NOT_OK(status); |
| } else if (type->Equals(arrow::large_binary())) { |
| auto status = detail::CastToLargeOffsetArray<arrow::BinaryArray, |
| arrow::LargeBinaryArray>( |
| table->column(i), type, chunked_array); |
| GAR_RETURN_NOT_OK(status); |
| } else { |
| // noop |
| chunked_array = table->column(i); |
| } |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(table, table->RemoveColumn(i)); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| table, table->AddColumn(i, field, chunked_array)); |
| } |
| return table; |
| } |
| |
| template <typename T> |
| Result<T> FileSystem::ReadFileToValue(const std::string& path) const noexcept { |
| T ret; |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto input, |
| arrow_fs_->OpenInputStream(path)); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto bytes, |
| input->Read(sizeof(T), &ret)); |
| ARROW_UNUSED(bytes); |
| return ret; |
| } |
| |
| template <> |
| Result<std::string> FileSystem::ReadFileToValue(const std::string& path) const |
| noexcept { |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto access_file, |
| arrow_fs_->OpenInputFile(path)); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto bytes, access_file->GetSize()); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto buffer, |
| access_file->ReadAt(0, bytes)); |
| return buffer->ToString(); |
| } |
| |
| template <typename T> |
| Status FileSystem::WriteValueToFile(const T& value, |
| const std::string& path) const noexcept { |
| // try to create the directory, oss filesystem may not support this, ignore |
| ARROW_UNUSED(arrow_fs_->CreateDir(path.substr(0, path.find_last_of("/")))); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto ofstream, |
| arrow_fs_->OpenOutputStream(path)); |
| RETURN_NOT_ARROW_OK(ofstream->Write(&value, sizeof(T))); |
| RETURN_NOT_ARROW_OK(ofstream->Close()); |
| return Status::OK(); |
| } |
| |
| template <> |
| Status FileSystem::WriteValueToFile(const std::string& value, |
| const std::string& path) const noexcept { |
| // try to create the directory, oss filesystem may not support this, ignore |
| ARROW_UNUSED(arrow_fs_->CreateDir(path.substr(0, path.find_last_of("/")))); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto ofstream, |
| arrow_fs_->OpenOutputStream(path)); |
| RETURN_NOT_ARROW_OK(ofstream->Write(value.c_str(), value.size())); |
| RETURN_NOT_ARROW_OK(ofstream->Close()); |
| return Status::OK(); |
| } |
| |
| Status FileSystem::WriteTableToFile( |
| const std::shared_ptr<arrow::Table>& table, FileType file_type, |
| const std::string& path, |
| const std::shared_ptr<WriterOptions>& options) const noexcept { |
| // try to create the directory, oss filesystem may not support this, ignore |
| ARROW_UNUSED(arrow_fs_->CreateDir(path.substr(0, path.find_last_of("/")))); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto output_stream, |
| arrow_fs_->OpenOutputStream(path)); |
| switch (file_type) { |
| case FileType::CSV: { |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| auto writer, |
| arrow::csv::MakeCSVWriter(output_stream.get(), table->schema(), |
| options->getCsvOption())); |
| RETURN_NOT_ARROW_OK(writer->WriteTable(*table)); |
| RETURN_NOT_ARROW_OK(writer->Close()); |
| break; |
| } |
| case FileType::PARQUET: { |
| auto schema = table->schema(); |
| RETURN_NOT_ARROW_OK(parquet::arrow::WriteTable( |
| *table, arrow::default_memory_pool(), output_stream, 64 * 1024 * 1024, |
| options->getParquetWriterProperties(), |
| options->getArrowWriterProperties())); |
| break; |
| } |
| #ifdef ARROW_ORC |
| case FileType::ORC: { |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| auto writer, arrow::adapters::orc::ORCFileWriter::Open( |
| output_stream.get(), options->getOrcOption())); |
| RETURN_NOT_ARROW_OK(writer->Write(*table)); |
| RETURN_NOT_ARROW_OK(writer->Close()); |
| break; |
| } |
| #endif |
| default: |
| return Status::Invalid( |
| "Unsupported file type: ", FileTypeToString(file_type), " for wrting."); |
| } |
| return Status::OK(); |
| } |
| |
| Status FileSystem::WriteLabelTableToFile( |
| const std::shared_ptr<arrow::Table>& table, const std::string& path) const |
| noexcept { |
| // try to create the directory, oss filesystem may not support this, ignore |
| ARROW_UNUSED(arrow_fs_->CreateDir(path.substr(0, path.find_last_of("/")))); |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(auto output_stream, |
| arrow_fs_->OpenOutputStream(path)); |
| auto schema = table->schema(); |
| auto column_num = schema->num_fields(); |
| parquet::WriterProperties::Builder builder; |
| builder.compression(arrow::Compression::type::ZSTD); // enable compression |
| builder.encoding(parquet::Encoding::RLE); |
| RETURN_NOT_ARROW_OK(parquet::arrow::WriteTable( |
| *table, arrow::default_memory_pool(), output_stream, 64 * 1024 * 1024, |
| builder.build(), parquet::default_arrow_writer_properties())); |
| return Status::OK(); |
| } |
| |
| Status FileSystem::CopyFile(const std::string& src_path, |
| const std::string& dst_path) const noexcept { |
| // try to create the directory, oss filesystem may not support this, ignore |
| ARROW_UNUSED( |
| arrow_fs_->CreateDir(dst_path.substr(0, dst_path.find_last_of("/")))); |
| RETURN_NOT_ARROW_OK(arrow_fs_->CopyFile(src_path, dst_path)); |
| return Status::OK(); |
| } |
| |
| Result<IdType> FileSystem::GetFileNumOfDir(const std::string& dir_path, |
| bool recursive) const noexcept { |
| arrow::fs::FileSelector file_selector; |
| file_selector.base_dir = dir_path; |
| file_selector.allow_not_found = false; // if dir_path not exist, return error |
| file_selector.recursive = recursive; |
| arrow::fs::FileInfoVector file_infos; |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN(file_infos, |
| arrow_fs_->GetFileInfo(file_selector)); |
| return static_cast<IdType>(file_infos.size()); |
| } |
| |
| FileSystem::~FileSystem() {} |
| |
| Result<std::shared_ptr<FileSystem>> FileSystemFromUriOrPath( |
| const std::string& uri_string, std::string* out_path) { |
| if (uri_string.length() >= 1 && uri_string[0] == '/') { |
| // if the uri_string is an absolute path, we need to create a local file |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| auto arrow_fs, |
| arrow::fs::FileSystemFromUriOrPath(uri_string, out_path)); |
| // arrow would delete the last slash, so use uri string |
| if (out_path != nullptr) { |
| *out_path = uri_string; |
| } |
| return std::make_shared<FileSystem>(arrow_fs); |
| } |
| |
| GAR_RETURN_ON_ARROW_ERROR_AND_ASSIGN( |
| auto arrow_fs, arrow::fs::FileSystemFromUriOrPath(uri_string)); |
| auto uri = uri::parse_uri(uri_string); |
| if (uri.error != uri::Error::None) { |
| return Status::Invalid("Failed to parse URI: ", uri_string); |
| } |
| if (out_path != nullptr) { |
| if (uri.scheme == "file" || uri.scheme == "hdfs" || uri.scheme.empty()) { |
| *out_path = uri.path; |
| } else if (uri.scheme == "s3" || uri.scheme == "gs") { |
| // bucket name is the host, path is the path |
| *out_path = uri.authority.host + uri.path; |
| } else { |
| return Status::Invalid("Unrecognized filesystem type in URI: ", |
| uri_string); |
| } |
| } |
| return std::make_shared<FileSystem>(arrow_fs); |
| } |
| |
| // arrow::fs::InitializeS3 and arrow::fs::FinalizeS3 need arrow_version >= 15 |
| Status InitializeS3() { |
| #if defined(ARROW_VERSION) && ARROW_VERSION > 14000000 |
| auto options = arrow::fs::S3GlobalOptions::Defaults(); |
| #else |
| arrow::fs::S3GlobalOptions options; |
| options.log_level = arrow::fs::S3LogLevel::Fatal; |
| #endif |
| #if defined(ARROW_VERSION) && ARROW_VERSION >= 15000000 |
| RETURN_NOT_ARROW_OK(arrow::fs::InitializeS3(options)); |
| #endif |
| return Status::OK(); |
| } |
| |
| Status FinalizeS3() { |
| #if defined(ARROW_VERSION) && ARROW_VERSION >= 15000000 |
| RETURN_NOT_ARROW_OK(arrow::fs::FinalizeS3()); |
| #endif |
| return Status::OK(); |
| } |
| |
| /// template specialization for std::string |
| template Result<IdType> FileSystem::ReadFileToValue<IdType>( |
| const std::string&) const noexcept; |
| /// template specialization for std::string |
| template Status FileSystem::WriteValueToFile<IdType>(const IdType&, |
| const std::string&) const |
| noexcept; |
| } // namespace graphar |
