Skip to content

ARROW-7058: [C++] FileSystemDataSourceDiscovery should apply partition schemes relative to its base dir #5772

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bkietz wants to merge 3 commits into from
Closed

ARROW-7058: [C++] FileSystemDataSourceDiscovery should apply partition schemes relative to its base dir #5772

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1bae5ff
ARROW-7058: [C++] FilSystemDataSourceDiscovery should apply partition…
bkietz Nov 4, 2019
f2f9689
add DCHECK for path containing selector
bkietz Nov 4, 2019
c0edfa5
make base_dir of partition schemes explicitly optional
bkietz Nov 5, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 9 additions & 12 deletions cpp/src/arrow/dataset/discovery.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,25 +34,21 @@ namespace arrow {
namespace dataset {

FileSystemDataSourceDiscovery::FileSystemDataSourceDiscovery(
fs::FileSystem* filesystem, std::vector<fs::FileStats> files,
fs::FileSystem* filesystem, fs::Selector selector, std::vector<fs::FileStats> files,
std::shared_ptr<FileFormat> format)
: fs_(filesystem), files_(std::move(files)), format_(std::move(format)) {}

Status FileSystemDataSourceDiscovery::Make(fs::FileSystem* filesystem,
std::vector<fs::FileStats> files,
std::shared_ptr<FileFormat> format,
std::shared_ptr<DataSourceDiscovery>* out) {
out->reset(new FileSystemDataSourceDiscovery(filesystem, files, format));
return Status::OK();
}
: fs_(filesystem),
selector_(std::move(selector)),
files_(std::move(files)),
format_(std::move(format)) {}

Status FileSystemDataSourceDiscovery::Make(fs::FileSystem* filesystem,
fs::Selector selector,
std::shared_ptr<FileFormat> format,
std::shared_ptr<DataSourceDiscovery>* out) {
std::vector<fs::FileStats> files;
RETURN_NOT_OK(filesystem->GetTargetStats(selector, &files));
return Make(filesystem, files, format, out);
out->reset(new FileSystemDataSourceDiscovery(filesystem, selector, files, format));
return Status::OK();
}

static inline Status InspectSchema(fs::FileSystem* fs,
Expand Down Expand Up @@ -85,7 +81,8 @@ Status FileSystemDataSourceDiscovery::Finish(std::shared_ptr<DataSource>* out) {
PathPartitions partitions;

if (partition_scheme_ != nullptr) {
RETURN_NOT_OK(ApplyPartitionScheme(*partition_scheme_, files_, &partitions));
RETURN_NOT_OK(
ApplyPartitionScheme(*partition_scheme_, selector_, files_, &partitions));
}

return FileSystemBasedDataSource::Make(fs_, files_, root_partition(),
Expand Down
7 changes: 2 additions & 5 deletions cpp/src/arrow/dataset/discovery.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,10 +89,6 @@ class ARROW_DS_EXPORT DataSourceDiscovery {
/// of fs::FileStats or a fs::Selector.
class ARROW_DS_EXPORT FileSystemDataSourceDiscovery : public DataSourceDiscovery {
public:
static Status Make(fs::FileSystem* filesystem, std::vector<fs::FileStats> files,
std::shared_ptr<FileFormat> format,
std::shared_ptr<DataSourceDiscovery>* out);

static Status Make(fs::FileSystem* filesystem, fs::Selector selector,
std::shared_ptr<FileFormat> format,
std::shared_ptr<DataSourceDiscovery>* out);
Expand All @@ -102,11 +98,12 @@ class ARROW_DS_EXPORT FileSystemDataSourceDiscovery : public DataSourceDiscovery
Status Finish(std::shared_ptr<DataSource>* out) override;

protected:
FileSystemDataSourceDiscovery(fs::FileSystem* filesystem,
FileSystemDataSourceDiscovery(fs::FileSystem* filesystem, fs::Selector selector,
std::vector<fs::FileStats> files,
std::shared_ptr<FileFormat> format);

fs::FileSystem* fs_;
fs::Selector selector_;
std::vector<fs::FileStats> files_;
std::shared_ptr<FileFormat> format_;
};
Expand Down
32 changes: 21 additions & 11 deletions cpp/src/arrow/dataset/discovery_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "arrow/dataset/partition.h"
#include "arrow/dataset/test_util.h"
#include "arrow/filesystem/test_util.h"

Expand All @@ -31,16 +32,11 @@ class FileSystemDataSourceDiscoveryTest : public TestFileSystemBasedDataSource {
void MakeDiscovery(const std::vector<fs::FileStats>& files) {
MakeFileSystem(files);
ASSERT_OK(
FileSystemDataSourceDiscovery::Make(fs_.get(), files, format_, &discovery_));
}

void MakeDiscovery(const std::vector<fs::FileStats>& files, fs::Selector selector) {
MakeFileSystem(files);
ASSERT_OK(
FileSystemDataSourceDiscovery::Make(fs_.get(), selector, format_, &discovery_));
FileSystemDataSourceDiscovery::Make(fs_.get(), selector_, format_, &discovery_));
}

protected:
fs::Selector selector_;
std::shared_ptr<DataSourceDiscovery> discovery_;
std::shared_ptr<FileFormat> format_ = std::make_shared<DummyFileFormat>();
};
Expand All @@ -53,15 +49,29 @@ TEST_F(FileSystemDataSourceDiscoveryTest, Basic) {
}

TEST_F(FileSystemDataSourceDiscoveryTest, Selector) {
// This test ensure that the Selector is enforced.
fs::Selector selector;
selector.base_dir = "A";
MakeDiscovery({fs::File("0"), fs::File("A/a")}, selector);
selector_.base_dir = "A";
MakeDiscovery({fs::File("0"), fs::File("A/a")});

ASSERT_OK(discovery_->Finish(&source_));
// "0" doesn't match selector, so it has been dropped:
AssertFragmentsAreFromPath(source_->GetFragments(options_), {"A/a"});
}

TEST_F(FileSystemDataSourceDiscoveryTest, Partition) {
selector_.base_dir = "a=ignored/base";
MakeDiscovery(
{fs::File(selector_.base_dir + "/a=1"), fs::File(selector_.base_dir + "/a=2")});

auto partition_scheme =
std::make_shared<HivePartitionScheme>(schema({field("a", int32())}));

ASSERT_OK(discovery_->SetPartitionScheme(partition_scheme));
ASSERT_OK(discovery_->Finish(&source_));

AssertFragmentsAreFromPath(source_->GetFragments(options_),
{selector_.base_dir + "/a=1", selector_.base_dir + "/a=2"});
}

TEST_F(FileSystemDataSourceDiscoveryTest, Inspect) {
auto s = schema({field("f64", float64())});
format_ = std::make_shared<DummyFileFormat>(s);
Expand Down
7 changes: 5 additions & 2 deletions cpp/src/arrow/dataset/partition.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,10 +88,13 @@ Result<std::shared_ptr<Expression>> HivePartitionScheme::Parse(
return ConvertPartitionKeys(GetUnconvertedKeys(path), *schema_);
}

Status ApplyPartitionScheme(const PartitionScheme& scheme,
Status ApplyPartitionScheme(const PartitionScheme& scheme, const fs::Selector& selector,
std::vector<fs::FileStats> files, PathPartitions* out) {
for (const auto& file : files) {
const auto& path = file.path();
// XXX is this the right way to drop the base dir?
DCHECK(std::equal(selector.base_dir.begin(), selector.base_dir.end(),
Comment thread
bkietz marked this conversation as resolved.
Outdated
file.path().begin()));
const auto& path = file.path().substr(selector.base_dir.size());
Comment thread
bkietz marked this conversation as resolved.
Outdated
std::shared_ptr<Expression> partition;
RETURN_NOT_OK(scheme.Parse(path, &partition));
out->emplace(path, std::move(partition));
Expand Down
5 changes: 3 additions & 2 deletions cpp/src/arrow/dataset/partition.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@ namespace arrow {

namespace fs {
struct FileStats;
}
struct Selector;
} // namespace fs

namespace dataset {

Expand Down Expand Up @@ -168,7 +169,7 @@ class ARROW_DS_EXPORT FunctionPartitionScheme : public PartitionScheme {
/// \brief Mapping from path to partition expressions.
using PathPartitions = std::unordered_map<std::string, std::shared_ptr<Expression>>;

Status ApplyPartitionScheme(const PartitionScheme& scheme,
Status ApplyPartitionScheme(const PartitionScheme& scheme, const fs::Selector& selector,
std::vector<fs::FileStats> files, PathPartitions* out);

// TODO(bkietz) use RE2 and named groups to provide RegexpPartitionScheme
Expand Down