ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset method#8301

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
jorisvandenbossche wants to merge 6 commits into from
Closed

ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset method #8301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9222a92
ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset m…
jorisvandenbossche Sep 29, 2020
a4faba0
add docstring + tests
jorisvandenbossche Oct 9, 2020
a5064ee
Make ParquetFileFragment::row_groups() explicit empty/no selection
bkietz Oct 9, 2020
01c5307
correct test
bkietz Oct 9, 2020
529af73
update tests
jorisvandenbossche Oct 10, 2020
13cdd3c
remove superfluous row_groups_ setter
jorisvandenbossche Oct 10, 2020
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
104 changes: 73 additions & 31 deletions cpp/src/arrow/dataset/file_parquet.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -395,14 +395,10 @@ Result<ScanTaskIterator> ParquetFileFormat::ScanFile(std::shared_ptr<ScanOptions
} else {
// since we are not scanning this fragment with a filter, don't bother loading
// statistics
row_groups = parquet_fragment->row_groups();
if (row_groups.empty()) {
// empty vector represents all row groups
std::shared_ptr<parquet::FileMetaData> metadata =
reader->parquet_reader()->metadata();
int num_row_groups = metadata->num_row_groups();
row_groups = RowGroupInfo::FromCount(num_row_groups);
if (parquet_fragment->row_groups() == nullptr) {
parquet_fragment->SetNumRowGroups(reader->num_row_groups());
}
row_groups = *parquet_fragment->row_groups();
}
}

Expand All@@ -424,7 +420,7 @@ Result<std::shared_ptr<FileFragment>> ParquetFileFormat::MakeFragment(
std::shared_ptr<Schema> physical_schema) {
return std::shared_ptr<FileFragment>(new ParquetFileFragment(
std::move(source), shared_from_this(), std::move(partition_expression),
std::move(physical_schema), {}));
std::move(physical_schema)));
}

//
Expand DownExpand Up@@ -530,6 +526,14 @@ bool RowGroupInfo::Satisfy(const Expression& predicate) const {
// ParquetFileFragment
//

ParquetFileFragment::ParquetFileFragment(FileSource source,
std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema)
: FileFragment(std::move(source), std::move(format), std::move(partition_expression),
std::move(physical_schema)),
parquet_format_(checked_cast<ParquetFileFormat&>(*format_)) {}

ParquetFileFragment::ParquetFileFragment(FileSource source,
std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
Expand All@@ -540,22 +544,28 @@ ParquetFileFragment::ParquetFileFragment(FileSource source,
row_groups_(std::move(row_groups)),
parquet_format_(checked_cast<ParquetFileFormat&>(*format_)),
has_complete_metadata_(RowGroupInfosAreComplete(row_groups_) &&
physical_schema_ != nullptr) {
if (!row_groups_.empty()) {
// Empty row_groups_ indicates selection of all row groups in the file, so we must
// open a reader to determine the real count.
num_row_groups_ = static_cast<int>(row_groups_.size());
physical_schema_ != nullptr),
num_row_groups_(static_cast<int>(row_groups_.size())) {}

void ParquetFileFragment::SetNumRowGroups(int num_row_groups) {
DCHECK_EQ(num_row_groups_, -1);
num_row_groups_ = num_row_groups;
row_groups_ = RowGroupInfo::FromCount(num_row_groups);
}

const std::vector<RowGroupInfo>* ParquetFileFragment::row_groups() {
auto lock = physical_schema_mutex_.Lock();
if (num_row_groups_ == -1) {
return nullptr;
}
return &row_groups_;
}

Result<int> ParquetFileFragment::GetNumRowGroups() {
auto lock = physical_schema_mutex_.Lock();
if (num_row_groups_ == -1) {
ARROW_ASSIGN_OR_RAISE(auto reader, parquet_format_.GetReader(source_));
num_row_groups_ = reader->num_row_groups();
if (row_groups_.empty()) {
row_groups_ = RowGroupInfo::FromCount(num_row_groups_);
}
SetNumRowGroups(reader->num_row_groups());
}
return num_row_groups_;
}
Expand DownExpand Up@@ -584,23 +594,22 @@ Status ParquetFileFragment::EnsureCompleteMetadata(parquet::arrow::FileReader* r
}
physical_schema_ = std::move(schema);

std::shared_ptr<parquet::FileMetaData> metadata = reader->parquet_reader()->metadata();
int num_row_groups = metadata->num_row_groups();

if (row_groups_.empty()) {
num_row_groups_ = num_row_groups;
row_groups_ = RowGroupInfo::FromCount(num_row_groups);
}

for (const RowGroupInfo& info : row_groups_) {
// Ensure RowGroups are indexing valid RowGroups before augmenting.
if (info.id() >= num_row_groups) {
return Status::IndexError("Trying to scan row group ", info.id(), " but ",
source_.path(), " only has ", num_row_groups,
" row groups");
if (num_row_groups_ == -1) {
SetNumRowGroups(reader->num_row_groups());
} else {
// a subselection was specified; validate that now
int num_row_groups = reader->num_row_groups();
for (const RowGroupInfo& info : row_groups_) {
// Ensure RowGroups are indexing valid RowGroups before augmenting.
if (info.id() >= num_row_groups) {
return Status::IndexError("Trying to scan row group ", info.id(), " but ",
source_.path(), " only has ", num_row_groups,
" row groups");
}
}
}

parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get();
for (RowGroupInfo& info : row_groups_) {
// Augment a RowGroup with statistics if missing.
if (info.HasStatistics()) continue;
Expand DownExpand Up@@ -631,6 +640,39 @@ Result<FragmentVector> ParquetFileFragment::SplitByRowGroup(
return fragments;
}

Result<std::shared_ptr<Fragment>> ParquetFileFragment::Subset(
const std::shared_ptr<Expression>& predicate) {
RETURN_NOT_OK(EnsureCompleteMetadata());
ARROW_ASSIGN_OR_RAISE(auto row_groups, FilterRowGroups(*predicate));

ARROW_ASSIGN_OR_RAISE(auto new_fragment,
parquet_format_.MakeFragment(source_, partition_expression(),
row_groups, physical_schema_));
return new_fragment;
Comment thread
bkietz marked this conversation as resolved.
}

Result<std::shared_ptr<Fragment>> ParquetFileFragment::Subset(
const std::vector<int> row_group_ids) {
std::vector<RowGroupInfo> row_groups;

@jorisvandenbosschejorisvandenbosscheSep 29, 2020

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here I would still need to check that all row_group_ids exist in the fragment

if (num_row_groups_ == -1) {
row_groups = RowGroupInfo::FromIdentifiers(row_group_ids);
} else {
row_groups = row_groups_;
auto end = std::remove_if(
row_groups.begin(), row_groups.end(), [&row_group_ids](const RowGroupInfo& info) {
return std::find(row_group_ids.begin(), row_group_ids.end(), info.id()) ==
row_group_ids.end();
});
row_groups.erase(end, row_groups.end());
}

ARROW_ASSIGN_OR_RAISE(auto new_fragment,
parquet_format_.MakeFragment(source_, partition_expression(),
row_groups, physical_schema_));
return new_fragment;
}

Result<std::vector<RowGroupInfo>> ParquetFileFragment::FilterRowGroups(
const Expression& predicate) {
DCHECK(has_complete_metadata_);
Expand Down
18 changes: 14 additions & 4 deletions cpp/src/arrow/dataset/file_parquet.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -215,9 +215,9 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
public:
Result<FragmentVector> SplitByRowGroup(const std::shared_ptr<Expression>& predicate);

/// \brief Return the RowGroups selected by this fragment. An empty list
/// represents all RowGroups in the parquet file.
const std::vector<RowGroupInfo>& row_groups() const { return row_groups_; }
/// \brief Return the RowGroups selected by this fragment, or nullptr
/// if all RowGroups in the parquet file are selected.
const std::vector<RowGroupInfo>* row_groups();

/// \brief Return the number of row groups selected by this fragment.
Result<int> GetNumRowGroups();
Expand All@@ -232,12 +232,20 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
/// \brief Ensure attached statistics are complete and the physical schema is cached.
Status EnsureCompleteMetadata(parquet::arrow::FileReader* reader = NULLPTR);

/// \brief Return a filtered subset of the ParquetFileFragment.
Result<std::shared_ptr<Fragment>> Subset(const std::shared_ptr<Expression>& predicate);
Result<std::shared_ptr<Fragment>> Subset(const std::vector<int> row_group_ids);
Comment thread
bkietz marked this conversation as resolved.

private:
ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema,
std::vector<RowGroupInfo> row_groups);

ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema);

// Overridden to opportunistically set metadata since a reader must be opened anyway.
Result<std::shared_ptr<Schema>> ReadPhysicalSchemaImpl() override {
ARROW_RETURN_NOT_OK(EnsureCompleteMetadata());
Expand All@@ -247,9 +255,11 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
// Return a filtered subset of RowGroupInfos.
Result<std::vector<RowGroupInfo>> FilterRowGroups(const Expression& predicate);

void SetNumRowGroups(int);

std::vector<RowGroupInfo> row_groups_;
ParquetFileFormat& parquet_format_;
bool has_complete_metadata_;
bool has_complete_metadata_ = false;
int num_row_groups_ = -1;

friend class ParquetFileFormat;
Expand Down
16 changes: 11 additions & 5 deletions cpp/src/arrow/dataset/file_parquet_test.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -197,7 +197,7 @@ class TestParquetFileFormat : public ArrowParquetWriterMixin {
auto expected = expected_row_groups[i];
auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragments[i]);

EXPECT_EQ(parquet_fragment->row_groups(),
EXPECT_EQ(*parquet_fragment->row_groups(),
RowGroupInfo::FromIdentifiers({expected}));
EXPECT_EQ(SingleBatch(parquet_fragment.get())->num_rows(), expected + 1);
}
Expand DownExpand Up@@ -514,14 +514,20 @@ TEST_F(TestParquetFileFormat, ExplicitRowGroupSelection) {
return internal::checked_pointer_cast<ParquetFileFragment>(fragment);
};

// empty selection is identical to selecting all row groups
EXPECT_TRUE(row_groups_fragment({})->row_groups().empty());
CountRowsAndBatchesInScan(row_groups_fragment({}), kTotalNumRows, kNumRowGroups);
EXPECT_OK_AND_ASSIGN(auto all_row_groups_fragment,
format_->MakeFragment(*source, scalar(true)));

// null selection is identical to selecting all row groups
EXPECT_EQ(internal::checked_pointer_cast<ParquetFileFragment>(all_row_groups_fragment)
->row_groups(),
nullptr);
CountRowsAndBatchesInScan(all_row_groups_fragment, kTotalNumRows, kNumRowGroups);

// individual selection selects a single row group
for (int i = 0; i < kNumRowGroups; ++i) {
CountRowsAndBatchesInScan(row_groups_fragment({i}), i + 1, 1);
EXPECT_EQ(row_groups_fragment({i})->row_groups(), RowGroupInfo::FromIdentifiers({i}));
EXPECT_EQ(*row_groups_fragment({i})->row_groups(),
RowGroupInfo::FromIdentifiers({i}));
}

for (int i = 0; i < kNumRowGroups; ++i) {
Expand Down
66 changes: 62 additions & 4 deletions python/pyarrow/_dataset.pyx
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1001,11 +1001,12 @@ cdef class ParquetFileFragment(FileFragment):
@property
def row_groups(self):
cdef:
vector[CRowGroupInfo] c_row_groups
const vector[CRowGroupInfo]* c_row_groups
c_row_groups = self.parquet_file_fragment.row_groups()
if c_row_groups.empty():
if c_row_groups == nullptr:
return None
return [RowGroupInfo.wrap(row_group) for row_group in c_row_groups]
return [RowGroupInfo.wrap(c_row_groups.at(i))
for i in range(c_row_groups.size())]

@property
def num_row_groups(self):
Expand DownExpand Up@@ -1035,7 +1036,7 @@ cdef class ParquetFileFragment(FileFragment):

Returns
-------
A list of Fragment.
A list of Fragments
"""
cdef:
vector[shared_ptr[CFragment]] c_fragments
Expand All@@ -1050,6 +1051,63 @@ cdef class ParquetFileFragment(FileFragment):

return [Fragment.wrap(c_fragment) for c_fragment in c_fragments]

def subset(self, Expression filter=None, Schema schema=None,
object row_group_ids=None):
"""
Create a subset of the fragment (viewing a subset of the row groups).

Subset can be specified by either a filter predicate (with optional
schema) or by a list of row group IDs. Note that when using a filter,
the resulting fragment can be empty (viewing no row groups).

Parameters
----------
filter : Expression, default None
Only include the row groups which satisfy this predicate (using
the Parquet RowGroup statistics).
schema : Schema, default None
Schema to use when filtering row groups. Defaults to the
Fragment's phsyical schema
row_group_ids : list of ints
The row group IDs to include in the subset. Can only be specified
if `filter` is None.

Returns
-------
ParquetFileFragment
"""
cdef:
shared_ptr[CExpression] c_filter
vector[int] c_row_group_ids
shared_ptr[CFragment] c_fragment

if filter is not None and row_group_ids is not None:
raise ValueError(
"Cannot specify both 'filter' and 'row_group_ids'."
)

if filter is not None:
schema = schema or self.physical_schema
c_filter = _insert_implicit_casts(filter, schema)
with nogil:
c_fragment = move(GetResultValue(
self.parquet_file_fragment.SubsetWithFilter(
move(c_filter))))
elif row_group_ids is not None:
c_row_group_ids = [
<int> row_group for row_group in sorted(set(row_group_ids))
]
with nogil:
c_fragment = move(GetResultValue(
self.parquet_file_fragment.SubsetWithIds(
move(c_row_group_ids))))
else:
raise ValueError(
"Need to specify one of 'filter' or 'row_group_ids'"
)

return Fragment.wrap(c_fragment)


cdef class ParquetReadOptions(_Weakrefable):
"""
Expand Down
6 changes: 5 additions & 1 deletion python/pyarrow/includes/libarrow_dataset.pxd
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,10 +250,14 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:

cdef cppclass CParquetFileFragment "arrow::dataset::ParquetFileFragment"(
CFileFragment):
const vector[CRowGroupInfo]& row_groups() const
const vector[CRowGroupInfo]* row_groups() const
CResult[int] GetNumRowGroups()
CResult[vector[shared_ptr[CFragment]]] SplitByRowGroup(
shared_ptr[CExpression] predicate)
CResult[shared_ptr[CFragment]] SubsetWithFilter "Subset"(
shared_ptr[CExpression] predicate)
CResult[shared_ptr[CFragment]] SubsetWithIds "Subset"(
vector[int] row_group_ids)
CStatus EnsureCompleteMetadata()

cdef cppclass CFileSystemDatasetWriteOptions \
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset method#8301

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
jorisvandenbossche wants to merge 6 commits into from
Closed

ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset method #8301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9222a92
ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset m…
jorisvandenbossche Sep 29, 2020
a4faba0
add docstring + tests
jorisvandenbossche Oct 9, 2020
a5064ee
Make ParquetFileFragment::row_groups() explicit empty/no selection
bkietz Oct 9, 2020
01c5307
correct test
bkietz Oct 9, 2020
529af73
update tests
jorisvandenbossche Oct 10, 2020
13cdd3c
remove superfluous row_groups_ setter
jorisvandenbossche Oct 10, 2020
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
104 changes: 73 additions & 31 deletions cpp/src/arrow/dataset/file_parquet.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -395,14 +395,10 @@ Result<ScanTaskIterator> ParquetFileFormat::ScanFile(std::shared_ptr<ScanOptions
} else {
// since we are not scanning this fragment with a filter, don't bother loading
// statistics
row_groups = parquet_fragment->row_groups();
if (row_groups.empty()) {
// empty vector represents all row groups
std::shared_ptr<parquet::FileMetaData> metadata =
reader->parquet_reader()->metadata();
int num_row_groups = metadata->num_row_groups();
row_groups = RowGroupInfo::FromCount(num_row_groups);
if (parquet_fragment->row_groups() == nullptr) {
parquet_fragment->SetNumRowGroups(reader->num_row_groups());
}
row_groups = *parquet_fragment->row_groups();
}
}

Expand All@@ -424,7 +420,7 @@ Result<std::shared_ptr<FileFragment>> ParquetFileFormat::MakeFragment(
std::shared_ptr<Schema> physical_schema) {
return std::shared_ptr<FileFragment>(new ParquetFileFragment(
std::move(source), shared_from_this(), std::move(partition_expression),
std::move(physical_schema), {}));
std::move(physical_schema)));
}

//
Expand DownExpand Up@@ -530,6 +526,14 @@ bool RowGroupInfo::Satisfy(const Expression& predicate) const {
// ParquetFileFragment
//

ParquetFileFragment::ParquetFileFragment(FileSource source,
std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema)
: FileFragment(std::move(source), std::move(format), std::move(partition_expression),
std::move(physical_schema)),
parquet_format_(checked_cast<ParquetFileFormat&>(*format_)) {}

ParquetFileFragment::ParquetFileFragment(FileSource source,
std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
Expand All@@ -540,22 +544,28 @@ ParquetFileFragment::ParquetFileFragment(FileSource source,
row_groups_(std::move(row_groups)),
parquet_format_(checked_cast<ParquetFileFormat&>(*format_)),
has_complete_metadata_(RowGroupInfosAreComplete(row_groups_) &&
physical_schema_ != nullptr) {
if (!row_groups_.empty()) {
// Empty row_groups_ indicates selection of all row groups in the file, so we must
// open a reader to determine the real count.
num_row_groups_ = static_cast<int>(row_groups_.size());
physical_schema_ != nullptr),
num_row_groups_(static_cast<int>(row_groups_.size())) {}

void ParquetFileFragment::SetNumRowGroups(int num_row_groups) {
DCHECK_EQ(num_row_groups_, -1);
num_row_groups_ = num_row_groups;
row_groups_ = RowGroupInfo::FromCount(num_row_groups);
}

const std::vector<RowGroupInfo>* ParquetFileFragment::row_groups() {
auto lock = physical_schema_mutex_.Lock();
if (num_row_groups_ == -1) {
return nullptr;
}
return &row_groups_;
}

Result<int> ParquetFileFragment::GetNumRowGroups() {
auto lock = physical_schema_mutex_.Lock();
if (num_row_groups_ == -1) {
ARROW_ASSIGN_OR_RAISE(auto reader, parquet_format_.GetReader(source_));
num_row_groups_ = reader->num_row_groups();
if (row_groups_.empty()) {
row_groups_ = RowGroupInfo::FromCount(num_row_groups_);
}
SetNumRowGroups(reader->num_row_groups());
}
return num_row_groups_;
}
Expand DownExpand Up@@ -584,23 +594,22 @@ Status ParquetFileFragment::EnsureCompleteMetadata(parquet::arrow::FileReader* r
}
physical_schema_ = std::move(schema);

std::shared_ptr<parquet::FileMetaData> metadata = reader->parquet_reader()->metadata();
int num_row_groups = metadata->num_row_groups();

if (row_groups_.empty()) {
num_row_groups_ = num_row_groups;
row_groups_ = RowGroupInfo::FromCount(num_row_groups);
}

for (const RowGroupInfo& info : row_groups_) {
// Ensure RowGroups are indexing valid RowGroups before augmenting.
if (info.id() >= num_row_groups) {
return Status::IndexError("Trying to scan row group ", info.id(), " but ",
source_.path(), " only has ", num_row_groups,
" row groups");
if (num_row_groups_ == -1) {
SetNumRowGroups(reader->num_row_groups());
} else {
// a subselection was specified; validate that now
int num_row_groups = reader->num_row_groups();
for (const RowGroupInfo& info : row_groups_) {
// Ensure RowGroups are indexing valid RowGroups before augmenting.
if (info.id() >= num_row_groups) {
return Status::IndexError("Trying to scan row group ", info.id(), " but ",
source_.path(), " only has ", num_row_groups,
" row groups");
}
}
}

parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get();
for (RowGroupInfo& info : row_groups_) {
// Augment a RowGroup with statistics if missing.
if (info.HasStatistics()) continue;
Expand DownExpand Up@@ -631,6 +640,39 @@ Result<FragmentVector> ParquetFileFragment::SplitByRowGroup(
return fragments;
}

Result<std::shared_ptr<Fragment>> ParquetFileFragment::Subset(
const std::shared_ptr<Expression>& predicate) {
RETURN_NOT_OK(EnsureCompleteMetadata());
ARROW_ASSIGN_OR_RAISE(auto row_groups, FilterRowGroups(*predicate));

ARROW_ASSIGN_OR_RAISE(auto new_fragment,
parquet_format_.MakeFragment(source_, partition_expression(),
row_groups, physical_schema_));
return new_fragment;
Comment thread
bkietz marked this conversation as resolved.
}

Result<std::shared_ptr<Fragment>> ParquetFileFragment::Subset(
const std::vector<int> row_group_ids) {
std::vector<RowGroupInfo> row_groups;

@jorisvandenbosschejorisvandenbosscheSep 29, 2020

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here I would still need to check that all row_group_ids exist in the fragment

if (num_row_groups_ == -1) {
row_groups = RowGroupInfo::FromIdentifiers(row_group_ids);
} else {
row_groups = row_groups_;
auto end = std::remove_if(
row_groups.begin(), row_groups.end(), [&row_group_ids](const RowGroupInfo& info) {
return std::find(row_group_ids.begin(), row_group_ids.end(), info.id()) ==
row_group_ids.end();
});
row_groups.erase(end, row_groups.end());
}

ARROW_ASSIGN_OR_RAISE(auto new_fragment,
parquet_format_.MakeFragment(source_, partition_expression(),
row_groups, physical_schema_));
return new_fragment;
}

Result<std::vector<RowGroupInfo>> ParquetFileFragment::FilterRowGroups(
const Expression& predicate) {
DCHECK(has_complete_metadata_);
Expand Down
18 changes: 14 additions & 4 deletions cpp/src/arrow/dataset/file_parquet.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -215,9 +215,9 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
public:
Result<FragmentVector> SplitByRowGroup(const std::shared_ptr<Expression>& predicate);

/// \brief Return the RowGroups selected by this fragment. An empty list
/// represents all RowGroups in the parquet file.
const std::vector<RowGroupInfo>& row_groups() const { return row_groups_; }
/// \brief Return the RowGroups selected by this fragment, or nullptr
/// if all RowGroups in the parquet file are selected.
const std::vector<RowGroupInfo>* row_groups();

/// \brief Return the number of row groups selected by this fragment.
Result<int> GetNumRowGroups();
Expand All@@ -232,12 +232,20 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
/// \brief Ensure attached statistics are complete and the physical schema is cached.
Status EnsureCompleteMetadata(parquet::arrow::FileReader* reader = NULLPTR);

/// \brief Return a filtered subset of the ParquetFileFragment.
Result<std::shared_ptr<Fragment>> Subset(const std::shared_ptr<Expression>& predicate);
Result<std::shared_ptr<Fragment>> Subset(const std::vector<int> row_group_ids);
Comment thread
bkietz marked this conversation as resolved.

private:
ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema,
std::vector<RowGroupInfo> row_groups);

ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema);

// Overridden to opportunistically set metadata since a reader must be opened anyway.
Result<std::shared_ptr<Schema>> ReadPhysicalSchemaImpl() override {
ARROW_RETURN_NOT_OK(EnsureCompleteMetadata());
Expand All@@ -247,9 +255,11 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
// Return a filtered subset of RowGroupInfos.
Result<std::vector<RowGroupInfo>> FilterRowGroups(const Expression& predicate);

void SetNumRowGroups(int);

std::vector<RowGroupInfo> row_groups_;
ParquetFileFormat& parquet_format_;
bool has_complete_metadata_;
bool has_complete_metadata_ = false;
int num_row_groups_ = -1;

friend class ParquetFileFormat;
Expand Down
16 changes: 11 additions & 5 deletions cpp/src/arrow/dataset/file_parquet_test.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -197,7 +197,7 @@ class TestParquetFileFormat : public ArrowParquetWriterMixin {
auto expected = expected_row_groups[i];
auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragments[i]);

EXPECT_EQ(parquet_fragment->row_groups(),
EXPECT_EQ(*parquet_fragment->row_groups(),
RowGroupInfo::FromIdentifiers({expected}));
EXPECT_EQ(SingleBatch(parquet_fragment.get())->num_rows(), expected + 1);
}
Expand DownExpand Up@@ -514,14 +514,20 @@ TEST_F(TestParquetFileFormat, ExplicitRowGroupSelection) {
return internal::checked_pointer_cast<ParquetFileFragment>(fragment);
};

// empty selection is identical to selecting all row groups
EXPECT_TRUE(row_groups_fragment({})->row_groups().empty());
CountRowsAndBatchesInScan(row_groups_fragment({}), kTotalNumRows, kNumRowGroups);
EXPECT_OK_AND_ASSIGN(auto all_row_groups_fragment,
format_->MakeFragment(*source, scalar(true)));

// null selection is identical to selecting all row groups
EXPECT_EQ(internal::checked_pointer_cast<ParquetFileFragment>(all_row_groups_fragment)
->row_groups(),
nullptr);
CountRowsAndBatchesInScan(all_row_groups_fragment, kTotalNumRows, kNumRowGroups);

// individual selection selects a single row group
for (int i = 0; i < kNumRowGroups; ++i) {
CountRowsAndBatchesInScan(row_groups_fragment({i}), i + 1, 1);
EXPECT_EQ(row_groups_fragment({i})->row_groups(), RowGroupInfo::FromIdentifiers({i}));
EXPECT_EQ(*row_groups_fragment({i})->row_groups(),
RowGroupInfo::FromIdentifiers({i}));
}

for (int i = 0; i < kNumRowGroups; ++i) {
Expand Down
66 changes: 62 additions & 4 deletions python/pyarrow/_dataset.pyx
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1001,11 +1001,12 @@ cdef class ParquetFileFragment(FileFragment):
@property
def row_groups(self):
cdef:
vector[CRowGroupInfo] c_row_groups
const vector[CRowGroupInfo]* c_row_groups
c_row_groups = self.parquet_file_fragment.row_groups()
if c_row_groups.empty():
if c_row_groups == nullptr:
return None
return [RowGroupInfo.wrap(row_group) for row_group in c_row_groups]
return [RowGroupInfo.wrap(c_row_groups.at(i))
for i in range(c_row_groups.size())]

@property
def num_row_groups(self):
Expand DownExpand Up@@ -1035,7 +1036,7 @@ cdef class ParquetFileFragment(FileFragment):

Returns
-------
A list of Fragment.
A list of Fragments
"""
cdef:
vector[shared_ptr[CFragment]] c_fragments
Expand All@@ -1050,6 +1051,63 @@ cdef class ParquetFileFragment(FileFragment):

return [Fragment.wrap(c_fragment) for c_fragment in c_fragments]

def subset(self, Expression filter=None, Schema schema=None,
object row_group_ids=None):
"""
Create a subset of the fragment (viewing a subset of the row groups).

Subset can be specified by either a filter predicate (with optional
schema) or by a list of row group IDs. Note that when using a filter,
the resulting fragment can be empty (viewing no row groups).

Parameters
----------
filter : Expression, default None
Only include the row groups which satisfy this predicate (using
the Parquet RowGroup statistics).
schema : Schema, default None
Schema to use when filtering row groups. Defaults to the
Fragment's phsyical schema
row_group_ids : list of ints
The row group IDs to include in the subset. Can only be specified
if `filter` is None.

Returns
-------
ParquetFileFragment
"""
cdef:
shared_ptr[CExpression] c_filter
vector[int] c_row_group_ids
shared_ptr[CFragment] c_fragment

if filter is not None and row_group_ids is not None:
raise ValueError(
"Cannot specify both 'filter' and 'row_group_ids'."
)

if filter is not None:
schema = schema or self.physical_schema
c_filter = _insert_implicit_casts(filter, schema)
with nogil:
c_fragment = move(GetResultValue(
self.parquet_file_fragment.SubsetWithFilter(
move(c_filter))))
elif row_group_ids is not None:
c_row_group_ids = [
<int> row_group for row_group in sorted(set(row_group_ids))
]
with nogil:
c_fragment = move(GetResultValue(
self.parquet_file_fragment.SubsetWithIds(
move(c_row_group_ids))))
else:
raise ValueError(
"Need to specify one of 'filter' or 'row_group_ids'"
)

return Fragment.wrap(c_fragment)


cdef class ParquetReadOptions(_Weakrefable):
"""
Expand Down
6 changes: 5 additions & 1 deletion python/pyarrow/includes/libarrow_dataset.pxd
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,10 +250,14 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:

cdef cppclass CParquetFileFragment "arrow::dataset::ParquetFileFragment"(
CFileFragment):
const vector[CRowGroupInfo]& row_groups() const
const vector[CRowGroupInfo]* row_groups() const
CResult[int] GetNumRowGroups()
CResult[vector[shared_ptr[CFragment]]] SplitByRowGroup(
shared_ptr[CExpression] predicate)
CResult[shared_ptr[CFragment]] SubsetWithFilter "Subset"(
shared_ptr[CExpression] predicate)
CResult[shared_ptr[CFragment]] SubsetWithIds "Subset"(
vector[int] row_group_ids)
CStatus EnsureCompleteMetadata()

cdef cppclass CFileSystemDatasetWriteOptions \
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset method#8301

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
jorisvandenbossche wants to merge 6 commits into from
Closed

ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset method #8301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9222a92
ARROW-10100: [C++][Python][Dataset] Add ParquetFileFragment::Subset m…
jorisvandenbossche Sep 29, 2020
a4faba0
add docstring + tests
jorisvandenbossche Oct 9, 2020
a5064ee
Make ParquetFileFragment::row_groups() explicit empty/no selection
bkietz Oct 9, 2020
01c5307
correct test
bkietz Oct 9, 2020
529af73
update tests
jorisvandenbossche Oct 10, 2020
13cdd3c
remove superfluous row_groups_ setter
jorisvandenbossche Oct 10, 2020
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
104 changes: 73 additions & 31 deletions cpp/src/arrow/dataset/file_parquet.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -395,14 +395,10 @@ Result<ScanTaskIterator> ParquetFileFormat::ScanFile(std::shared_ptr<ScanOptions
} else {
// since we are not scanning this fragment with a filter, don't bother loading
// statistics
row_groups = parquet_fragment->row_groups();
if (row_groups.empty()) {
// empty vector represents all row groups
std::shared_ptr<parquet::FileMetaData> metadata =
reader->parquet_reader()->metadata();
int num_row_groups = metadata->num_row_groups();
row_groups = RowGroupInfo::FromCount(num_row_groups);
if (parquet_fragment->row_groups() == nullptr) {
parquet_fragment->SetNumRowGroups(reader->num_row_groups());
}
row_groups = *parquet_fragment->row_groups();
}
}

Expand All@@ -424,7 +420,7 @@ Result<std::shared_ptr<FileFragment>> ParquetFileFormat::MakeFragment(
std::shared_ptr<Schema> physical_schema) {
return std::shared_ptr<FileFragment>(new ParquetFileFragment(
std::move(source), shared_from_this(), std::move(partition_expression),
std::move(physical_schema), {}));
std::move(physical_schema)));
}

//
Expand DownExpand Up@@ -530,6 +526,14 @@ bool RowGroupInfo::Satisfy(const Expression& predicate) const {
// ParquetFileFragment
//

ParquetFileFragment::ParquetFileFragment(FileSource source,
std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema)
: FileFragment(std::move(source), std::move(format), std::move(partition_expression),
std::move(physical_schema)),
parquet_format_(checked_cast<ParquetFileFormat&>(*format_)) {}

ParquetFileFragment::ParquetFileFragment(FileSource source,
std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
Expand All@@ -540,22 +544,28 @@ ParquetFileFragment::ParquetFileFragment(FileSource source,
row_groups_(std::move(row_groups)),
parquet_format_(checked_cast<ParquetFileFormat&>(*format_)),
has_complete_metadata_(RowGroupInfosAreComplete(row_groups_) &&
physical_schema_ != nullptr) {
if (!row_groups_.empty()) {
// Empty row_groups_ indicates selection of all row groups in the file, so we must
// open a reader to determine the real count.
num_row_groups_ = static_cast<int>(row_groups_.size());
physical_schema_ != nullptr),
num_row_groups_(static_cast<int>(row_groups_.size())) {}

void ParquetFileFragment::SetNumRowGroups(int num_row_groups) {
DCHECK_EQ(num_row_groups_, -1);
num_row_groups_ = num_row_groups;
row_groups_ = RowGroupInfo::FromCount(num_row_groups);
}

const std::vector<RowGroupInfo>* ParquetFileFragment::row_groups() {
auto lock = physical_schema_mutex_.Lock();
if (num_row_groups_ == -1) {
return nullptr;
}
return &row_groups_;
}

Result<int> ParquetFileFragment::GetNumRowGroups() {
auto lock = physical_schema_mutex_.Lock();
if (num_row_groups_ == -1) {
ARROW_ASSIGN_OR_RAISE(auto reader, parquet_format_.GetReader(source_));
num_row_groups_ = reader->num_row_groups();
if (row_groups_.empty()) {
row_groups_ = RowGroupInfo::FromCount(num_row_groups_);
}
SetNumRowGroups(reader->num_row_groups());
}
return num_row_groups_;
}
Expand DownExpand Up@@ -584,23 +594,22 @@ Status ParquetFileFragment::EnsureCompleteMetadata(parquet::arrow::FileReader* r
}
physical_schema_ = std::move(schema);

std::shared_ptr<parquet::FileMetaData> metadata = reader->parquet_reader()->metadata();
int num_row_groups = metadata->num_row_groups();

if (row_groups_.empty()) {
num_row_groups_ = num_row_groups;
row_groups_ = RowGroupInfo::FromCount(num_row_groups);
}

for (const RowGroupInfo& info : row_groups_) {
// Ensure RowGroups are indexing valid RowGroups before augmenting.
if (info.id() >= num_row_groups) {
return Status::IndexError("Trying to scan row group ", info.id(), " but ",
source_.path(), " only has ", num_row_groups,
" row groups");
if (num_row_groups_ == -1) {
SetNumRowGroups(reader->num_row_groups());
} else {
// a subselection was specified; validate that now
int num_row_groups = reader->num_row_groups();
for (const RowGroupInfo& info : row_groups_) {
// Ensure RowGroups are indexing valid RowGroups before augmenting.
if (info.id() >= num_row_groups) {
return Status::IndexError("Trying to scan row group ", info.id(), " but ",
source_.path(), " only has ", num_row_groups,
" row groups");
}
}
}

parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get();
for (RowGroupInfo& info : row_groups_) {
// Augment a RowGroup with statistics if missing.
if (info.HasStatistics()) continue;
Expand DownExpand Up@@ -631,6 +640,39 @@ Result<FragmentVector> ParquetFileFragment::SplitByRowGroup(
return fragments;
}

Result<std::shared_ptr<Fragment>> ParquetFileFragment::Subset(
const std::shared_ptr<Expression>& predicate) {
RETURN_NOT_OK(EnsureCompleteMetadata());
ARROW_ASSIGN_OR_RAISE(auto row_groups, FilterRowGroups(*predicate));

ARROW_ASSIGN_OR_RAISE(auto new_fragment,
parquet_format_.MakeFragment(source_, partition_expression(),
row_groups, physical_schema_));
return new_fragment;
Comment thread
bkietz marked this conversation as resolved.
}

Result<std::shared_ptr<Fragment>> ParquetFileFragment::Subset(
const std::vector<int> row_group_ids) {
std::vector<RowGroupInfo> row_groups;

@jorisvandenbosschejorisvandenbosscheSep 29, 2020

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here I would still need to check that all row_group_ids exist in the fragment

if (num_row_groups_ == -1) {
row_groups = RowGroupInfo::FromIdentifiers(row_group_ids);
} else {
row_groups = row_groups_;
auto end = std::remove_if(
row_groups.begin(), row_groups.end(), [&row_group_ids](const RowGroupInfo& info) {
return std::find(row_group_ids.begin(), row_group_ids.end(), info.id()) ==
row_group_ids.end();
});
row_groups.erase(end, row_groups.end());
}

ARROW_ASSIGN_OR_RAISE(auto new_fragment,
parquet_format_.MakeFragment(source_, partition_expression(),
row_groups, physical_schema_));
return new_fragment;
}

Result<std::vector<RowGroupInfo>> ParquetFileFragment::FilterRowGroups(
const Expression& predicate) {
DCHECK(has_complete_metadata_);
Expand Down
18 changes: 14 additions & 4 deletions cpp/src/arrow/dataset/file_parquet.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -215,9 +215,9 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
public:
Result<FragmentVector> SplitByRowGroup(const std::shared_ptr<Expression>& predicate);

/// \brief Return the RowGroups selected by this fragment. An empty list
/// represents all RowGroups in the parquet file.
const std::vector<RowGroupInfo>& row_groups() const { return row_groups_; }
/// \brief Return the RowGroups selected by this fragment, or nullptr
/// if all RowGroups in the parquet file are selected.
const std::vector<RowGroupInfo>* row_groups();

/// \brief Return the number of row groups selected by this fragment.
Result<int> GetNumRowGroups();
Expand All@@ -232,12 +232,20 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
/// \brief Ensure attached statistics are complete and the physical schema is cached.
Status EnsureCompleteMetadata(parquet::arrow::FileReader* reader = NULLPTR);

/// \brief Return a filtered subset of the ParquetFileFragment.
Result<std::shared_ptr<Fragment>> Subset(const std::shared_ptr<Expression>& predicate);
Result<std::shared_ptr<Fragment>> Subset(const std::vector<int> row_group_ids);
Comment thread
bkietz marked this conversation as resolved.

private:
ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema,
std::vector<RowGroupInfo> row_groups);

ParquetFileFragment(FileSource source, std::shared_ptr<FileFormat> format,
std::shared_ptr<Expression> partition_expression,
std::shared_ptr<Schema> physical_schema);

// Overridden to opportunistically set metadata since a reader must be opened anyway.
Result<std::shared_ptr<Schema>> ReadPhysicalSchemaImpl() override {
ARROW_RETURN_NOT_OK(EnsureCompleteMetadata());
Expand All@@ -247,9 +255,11 @@ class ARROW_DS_EXPORT ParquetFileFragment : public FileFragment {
// Return a filtered subset of RowGroupInfos.
Result<std::vector<RowGroupInfo>> FilterRowGroups(const Expression& predicate);

void SetNumRowGroups(int);

std::vector<RowGroupInfo> row_groups_;
ParquetFileFormat& parquet_format_;
bool has_complete_metadata_;
bool has_complete_metadata_ = false;
int num_row_groups_ = -1;

friend class ParquetFileFormat;
Expand Down
16 changes: 11 additions & 5 deletions cpp/src/arrow/dataset/file_parquet_test.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -197,7 +197,7 @@ class TestParquetFileFormat : public ArrowParquetWriterMixin {
auto expected = expected_row_groups[i];
auto parquet_fragment = checked_pointer_cast<ParquetFileFragment>(fragments[i]);

EXPECT_EQ(parquet_fragment->row_groups(),
EXPECT_EQ(*parquet_fragment->row_groups(),
RowGroupInfo::FromIdentifiers({expected}));
EXPECT_EQ(SingleBatch(parquet_fragment.get())->num_rows(), expected + 1);
}
Expand DownExpand Up@@ -514,14 +514,20 @@ TEST_F(TestParquetFileFormat, ExplicitRowGroupSelection) {
return internal::checked_pointer_cast<ParquetFileFragment>(fragment);
};

// empty selection is identical to selecting all row groups
EXPECT_TRUE(row_groups_fragment({})->row_groups().empty());
CountRowsAndBatchesInScan(row_groups_fragment({}), kTotalNumRows, kNumRowGroups);
EXPECT_OK_AND_ASSIGN(auto all_row_groups_fragment,
format_->MakeFragment(*source, scalar(true)));

// null selection is identical to selecting all row groups
EXPECT_EQ(internal::checked_pointer_cast<ParquetFileFragment>(all_row_groups_fragment)
->row_groups(),
nullptr);
CountRowsAndBatchesInScan(all_row_groups_fragment, kTotalNumRows, kNumRowGroups);

// individual selection selects a single row group
for (int i = 0; i < kNumRowGroups; ++i) {
CountRowsAndBatchesInScan(row_groups_fragment({i}), i + 1, 1);
EXPECT_EQ(row_groups_fragment({i})->row_groups(), RowGroupInfo::FromIdentifiers({i}));
EXPECT_EQ(*row_groups_fragment({i})->row_groups(),
RowGroupInfo::FromIdentifiers({i}));
}

for (int i = 0; i < kNumRowGroups; ++i) {
Expand Down
66 changes: 62 additions & 4 deletions python/pyarrow/_dataset.pyx
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1001,11 +1001,12 @@ cdef class ParquetFileFragment(FileFragment):
@property
def row_groups(self):
cdef:
vector[CRowGroupInfo] c_row_groups
const vector[CRowGroupInfo]* c_row_groups
c_row_groups = self.parquet_file_fragment.row_groups()
if c_row_groups.empty():
if c_row_groups == nullptr:
return None
return [RowGroupInfo.wrap(row_group) for row_group in c_row_groups]
return [RowGroupInfo.wrap(c_row_groups.at(i))
for i in range(c_row_groups.size())]

@property
def num_row_groups(self):
Expand DownExpand Up@@ -1035,7 +1036,7 @@ cdef class ParquetFileFragment(FileFragment):

Returns
-------
A list of Fragment.
A list of Fragments
"""
cdef:
vector[shared_ptr[CFragment]] c_fragments
Expand All@@ -1050,6 +1051,63 @@ cdef class ParquetFileFragment(FileFragment):

return [Fragment.wrap(c_fragment) for c_fragment in c_fragments]

def subset(self, Expression filter=None, Schema schema=None,
object row_group_ids=None):
"""
Create a subset of the fragment (viewing a subset of the row groups).

Subset can be specified by either a filter predicate (with optional
schema) or by a list of row group IDs. Note that when using a filter,
the resulting fragment can be empty (viewing no row groups).

Parameters
----------
filter : Expression, default None
Only include the row groups which satisfy this predicate (using
the Parquet RowGroup statistics).
schema : Schema, default None
Schema to use when filtering row groups. Defaults to the
Fragment's phsyical schema
row_group_ids : list of ints
The row group IDs to include in the subset. Can only be specified
if `filter` is None.

Returns
-------
ParquetFileFragment
"""
cdef:
shared_ptr[CExpression] c_filter
vector[int] c_row_group_ids
shared_ptr[CFragment] c_fragment

if filter is not None and row_group_ids is not None:
raise ValueError(
"Cannot specify both 'filter' and 'row_group_ids'."
)

if filter is not None:
schema = schema or self.physical_schema
c_filter = _insert_implicit_casts(filter, schema)
with nogil:
c_fragment = move(GetResultValue(
self.parquet_file_fragment.SubsetWithFilter(
move(c_filter))))
elif row_group_ids is not None:
c_row_group_ids = [
<int> row_group for row_group in sorted(set(row_group_ids))
]
with nogil:
c_fragment = move(GetResultValue(
self.parquet_file_fragment.SubsetWithIds(
move(c_row_group_ids))))
else:
raise ValueError(
"Need to specify one of 'filter' or 'row_group_ids'"
)

return Fragment.wrap(c_fragment)


cdef class ParquetReadOptions(_Weakrefable):
"""
Expand Down
6 changes: 5 additions & 1 deletion python/pyarrow/includes/libarrow_dataset.pxd
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,10 +250,14 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:

cdef cppclass CParquetFileFragment "arrow::dataset::ParquetFileFragment"(
CFileFragment):
const vector[CRowGroupInfo]& row_groups() const
const vector[CRowGroupInfo]* row_groups() const
CResult[int] GetNumRowGroups()
CResult[vector[shared_ptr[CFragment]]] SplitByRowGroup(
shared_ptr[CExpression] predicate)
CResult[shared_ptr[CFragment]] SubsetWithFilter "Subset"(
shared_ptr[CExpression] predicate)
CResult[shared_ptr[CFragment]] SubsetWithIds "Subset"(
vector[int] row_group_ids)
CStatus EnsureCompleteMetadata()

cdef cppclass CFileSystemDatasetWriteOptions \
Expand Down
Loading