Skip to content

Commit 721e6f9

Browse files
fsaintjacquesbkietz
authored and
wesm
committed
ARROW-6161: [C++][Dataset] Implements ParquetFragment
This implements the ParquetFragment whose main goal is to expose a streaming iterator of RecordBatch. Closes apache#5083 from fsaintjacques/ARROW-6161-parquet-scan and squashes the following commits: 3fc7ce6 <Benjamin Kietzman> ignore lint for awful doxygen section f28ef32 <Benjamin Kietzman> hide problematic inheritance from doxygen 2d86954 <François Saint-Jacques> Remove test_data.h.in a9b7b91 <François Saint-Jacques> Follow file convention 022527f <Benjamin Kietzman> explicit move return for Result return f7a28d1 <Benjamin Kietzman> clang-format 21a2be9 <Benjamin Kietzman> iwyu 870cfb6 <Benjamin Kietzman> move WriteRecordBatch to dataset/test_utils.h 4e480b3 <Benjamin Kietzman> replace RepeatedRecordBatch with GeneratedRecordBatch f0ee598 <Benjamin Kietzman> simplify repetetive array building 82f408d <Benjamin Kietzman> move dataset test utilities under arrow/dataset/ b2ead08 <Benjamin Kietzman> include iterator 02c2c8c <Benjamin Kietzman> add method impls to make MSVC happy cccb62d <Benjamin Kietzman> remove AllColumnIndices and AllRowGroups, add arrow::internal::Iota c88761f <Benjamin Kietzman> use RecordBatchIterator typedef 88adb7d <Benjamin Kietzman> first half of review comments 88f7118 <Benjamin Kietzman> include dataset.h fa3d5cc <Benjamin Kietzman> move destructor to source file for MSVC 6c7d18b <François Saint-Jacques> Formating c43a3a9 <François Saint-Jacques> Fix ASAN reported issue 016fa86 <François Saint-Jacques> Remove superfluous destructor 3c893c0 <François Saint-Jacques> Fix linter ddbedd8 <François Saint-Jacques> Make gcc-5 happy 1bee33c <François Saint-Jacques> cmake-format c3d840a <Benjamin Kietzman> refactoring, recommenting 189d11e <François Saint-Jacques> Refactor FileSource::Open to switch 6e19abe <François Saint-Jacques> Improve ScanOptions 0ce5684 <François Saint-Jacques> Fix unused variables ff48877 <François Saint-Jacques> Format and lint e01dec8 <François Saint-Jacques> Refactor test to use in-memory RecordBatchReader 1c960be <François Saint-Jacques> Fix test 054bfd7 <Benjamin Kietzman> add repeatedrecordbatch 9fdded8 <François Saint-Jacques> Add ParquetFragment integration test d951aba <François Saint-Jacques> Add integration data fixture class a80503a <François Saint-Jacques> Move back code into arrow/dataset and add proper cmake configuration 8614626 <François Saint-Jacques> Refactor parquet::arrow::FileReader aa24c1f <Benjamin Kietzman> merge conflict/lint fix d469ce0 <Benjamin Kietzman> pass Iterators by unique_ptr fa94688 <Benjamin Kietzman> make RecordBatchReader a refinement of the Iterator interface 5fbc38b <François Saint-Jacques> Initial commit Lead-authored-by: François Saint-Jacques <fsaintjacques@gmail.com> Co-authored-by: Benjamin Kietzman <bengilgit@gmail.com> Signed-off-by: Wes McKinney <wesm+git@apache.org>
1 parent f8b7425 commit 721e6f9

33 files changed

Lines changed: 826 additions & 103 deletions

‎cpp/src/arrow/adapters/orc/adapter.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,8 @@
4141
#include "arrow/util/checked_cast.h"
4242
#include "arrow/util/decimal.h"
4343
#include "arrow/util/key_value_metadata.h"
44-
#include "arrow/util/lazy.h"
4544
#include "arrow/util/macros.h"
45+
#include "arrow/util/range.h"
4646
#include "arrow/util/visibility.h"
4747

4848
#include "orc/Exceptions.hh"

‎cpp/src/arrow/adapters/orc/adapter_util.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
#include "arrow/status.h"
2525
#include "arrow/util/checked_cast.h"
2626
#include "arrow/util/decimal.h"
27-
#include "arrow/util/lazy.h"
27+
#include "arrow/util/range.h"
2828

2929
#include "orc/Exceptions.hh"
3030
#include "orc/OrcFile.hh"

‎cpp/src/arrow/array/diff.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@
3232
#include "arrow/pretty_print.h"
3333
#include "arrow/status.h"
3434
#include "arrow/type_traits.h"
35-
#include "arrow/util/lazy.h"
3635
#include "arrow/util/logging.h"
36+
#include "arrow/util/range.h"
3737
#include "arrow/util/stl.h"
3838
#include "arrow/util/string.h"
3939
#include "arrow/util/variant.h"

‎cpp/src/arrow/array_test.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
#include "arrow/util/bit_util.h"
4646
#include "arrow/util/checked_cast.h"
4747
#include "arrow/util/decimal.h"
48-
#include "arrow/util/lazy.h"
48+
#include "arrow/util/range.h"
4949

5050
// This file is compiled together with array-*-test.cc into a single
5151
// executable array-test.

‎cpp/src/arrow/dataset/CMakeLists.txt‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,28 @@ arrow_install_all_headers("arrow/dataset")
2323
# pkg-config support
2424
arrow_add_pkg_config("arrow-dataset")
2525

26-
set(ARROW_DATASET_SRCS scanner.cc)
26+
set(ARROW_DATASET_SRCS scanner.cc file_base.cc)
27+
set(ARROW_DATASET_LINK_STATIC arrow_static)
28+
set(ARROW_DATASET_LINK_SHARED arrow_shared)
29+
30+
if(ARROW_PARQUET)
31+
set(ARROW_DATASET_LINK_STATIC ${ARROW_DATASET_LINK_STATIC} parquet_static)
32+
set(ARROW_DATASET_LINK_SHARED ${ARROW_DATASET_LINK_SHARED} parquet_shared)
33+
set(ARROW_DATASET_SRCS ${ARROW_DATASET_SRCS} file_parquet.cc)
34+
set(ARROW_DATASET_PRIVATE_INCLUDES ${PROJECT_SOURCE_DIR}/src/parquet)
35+
endif()
2736

2837
add_arrow_lib(arrow_dataset
2938
OUTPUTS
3039
ARROW_DATASET_LIBRARIES
3140
SOURCES
3241
${ARROW_DATASET_SRCS}
42+
PRIVATE_INCLUDES
43+
${ARROW_DATASET_PRIVATE_INCLUDES}
3344
SHARED_LINK_LIBS
34-
arrow_shared
45+
${ARROW_DATASET_LINK_SHARED}
3546
STATIC_LINK_LIBS
36-
arrow_static)
47+
${ARROW_DATASET_LINK_STATIC})
3748

3849
if(ARROW_DATASET_TEST_LINKAGE STREQUAL "static")
3950
set(ARROW_DATASET_TEST_LINK_LIBS arrow_dataset_static ${ARROW_TEST_STATIC_LINK_LIBS})
@@ -53,4 +64,15 @@ if(NOT WIN32)
5364
"arrow-dataset"
5465
LABELS
5566
"arrow_dataset")
67+
68+
if(ARROW_PARQUET)
69+
add_arrow_test(file_parquet_test
70+
EXTRA_LINK_LIBS
71+
${ARROW_DATASET_TEST_LINK_LIBS}
72+
PREFIX
73+
"arrow-dataset"
74+
LABELS
75+
"arrow_dataset")
76+
77+
endif()
5678
endif()

‎cpp/src/arrow/dataset/api.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,4 @@
2222
#include "arrow/dataset/file_base.h"
2323
#include "arrow/dataset/file_csv.h"
2424
#include "arrow/dataset/file_feather.h"
25-
#include "arrow/dataset/file_parquet.h"
2625
#include "arrow/dataset/scanner.h"

‎cpp/src/arrow/dataset/dataset.h‎

Lines changed: 34 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -23,58 +23,65 @@
2323

2424
#include "arrow/dataset/type_fwd.h"
2525
#include "arrow/dataset/visibility.h"
26+
#include "arrow/util/iterator.h"
2627

2728
namespace arrow {
2829
namespace dataset {
2930

3031
/// \brief A granular piece of a Dataset, such as an individual file,
31-
/// which can be read/scanned separately from other fragments
32+
/// which can be read/scanned separately from other fragments.
33+
///
34+
/// A DataFragment yields a collection of RecordBatch, encapsulated in one or
35+
/// more ScanTasks.
3236
class ARROW_DS_EXPORT DataFragment {
3337
public:
34-
virtual ~DataFragment() = default;
38+
/// \brief Scan returns an iterator of ScanTasks, each of which yields
39+
/// RecordBatches from this DataFragment.
40+
virtual Status Scan(std::shared_ptr<ScanContext> scan_context,
41+
std::unique_ptr<ScanTaskIterator>* out) = 0;
3542

3643
/// \brief Return true if the fragment can benefit from parallel
3744
/// scanning
3845
virtual bool splittable() const = 0;
3946

40-
/// \brief Partition options to use when scanning this fragment. May be
41-
/// nullptr
47+
/// \brief Filtering, schema reconciliation, and partition options to use when
48+
/// scanning this fragment. May be nullptr, which indicates that no filtering
49+
/// or schema reconciliation will be performed and all partitions will be
50+
/// scanned.
4251
virtual std::shared_ptr<ScanOptions> scan_options() const = 0;
43-
};
44-
45-
/// \brief Conditions to apply to a dataset when reading to include or
46-
/// exclude fragments, filter out rows, etc.
47-
struct DataSelector {
48-
std::vector<std::shared_ptr<Filter>> filters;
4952

50-
// TODO(wesm): Select specific partition keys, file path globs, or
51-
// other common desirable selections
53+
virtual ~DataFragment() = default;
5254
};
5355

5456
/// \brief A basic component of a Dataset which yields zero or more
55-
/// DataFragments
57+
/// DataFragments. A DataSource acts as a discovery mechanism of DataFragments
58+
/// and partitions, e.g. files deeply nested in a directory.
5659
class ARROW_DS_EXPORT DataSource {
5760
public:
58-
virtual ~DataSource() = default;
61+
/// \brief GetFragments returns an iterator of DataFragments. The ScanOptions
62+
/// controls filtering and schema inference.
63+
virtual std::unique_ptr<DataFragmentIterator> GetFragments(
64+
std::shared_ptr<ScanOptions> options) = 0;
5965

6066
virtual std::string type() const = 0;
6167

62-
virtual std::unique_ptr<DataFragmentIterator> GetFragments(
63-
const DataSelector& selector) = 0;
68+
virtual ~DataSource() = default;
6469
};
6570

6671
/// \brief A DataSource consisting of a flat sequence of DataFragments
6772
class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
6873
public:
6974
std::unique_ptr<DataFragmentIterator> GetFragments(
70-
const DataSelector& selector) override;
75+
std::shared_ptr<ScanOptions> options) override {
76+
return MakeIterator(fragments_);
77+
}
7178

7279
private:
7380
DataFragmentVector fragments_;
7481
};
7582

7683
/// \brief Top-level interface for a Dataset with fragments coming
77-
/// from possibly multiple sources
84+
/// from possibly multiple sources.
7885
class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
7986
public:
8087
/// \param[in] source a single input data source
@@ -110,5 +117,14 @@ class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
110117
std::vector<std::shared_ptr<DataSource>> sources_;
111118
};
112119

120+
/// \brief Conditions to apply to a dataset when reading to include or
121+
/// exclude fragments, filter out rows, etc.
122+
struct DataSelector {
123+
std::vector<std::shared_ptr<Filter>> filters;
124+
125+
// TODO(wesm): Select specific partition keys, file path globs, or
126+
// other common desirable selections
127+
};
128+
113129
} // namespace dataset
114130
} // namespace arrow

‎cpp/src/arrow/dataset/file_base.cc‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include "arrow/dataset/file_base.h"
19+
20+
#include "arrow/filesystem/filesystem.h"
21+
#include "arrow/io/interfaces.h"
22+
#include "arrow/io/memory.h"
23+
24+
namespace arrow {
25+
namespace dataset {
26+
27+
Status FileSource::Open(std::shared_ptr<arrow::io::RandomAccessFile>* out) const {
28+
switch (type_) {
29+
case PATH:
30+
return filesystem_->OpenInputFile(path_, out);
31+
case BUFFER:
32+
*out = std::make_shared<::arrow::io::BufferReader>(buffer_);
33+
return Status::OK();
34+
}
35+
36+
return Status::OK();
37+
}
38+
39+
Status FileBasedDataFragment::Scan(std::shared_ptr<ScanContext> scan_context,
40+
std::unique_ptr<ScanTaskIterator>* out) {
41+
return format_->ScanFile(source_, scan_options_, scan_context, out);
42+
}
43+
44+
} // namespace dataset
45+
} // namespace arrow

‎cpp/src/arrow/dataset/file_base.h‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,20 @@
2121
#include <string>
2222
#include <utility>
2323

24+
#include "arrow/buffer.h"
25+
#include "arrow/dataset/dataset.h"
2426
#include "arrow/dataset/scanner.h"
2527
#include "arrow/dataset/type_fwd.h"
2628
#include "arrow/dataset/visibility.h"
2729
#include "arrow/dataset/writer.h"
30+
#include "arrow/io/file.h"
2831
#include "arrow/util/compression.h"
2932

3033
namespace arrow {
3134
namespace dataset {
3235

33-
/// \brief Contains the location of a file to be read
36+
/// \brief The path and filesystem where an actual file is located or a buffer which can
37+
/// be read like a file
3438
class ARROW_DS_EXPORT FileSource {
3539
public:
3640
enum SourceType { PATH, BUFFER };
@@ -77,6 +81,9 @@ class ARROW_DS_EXPORT FileSource {
7781
/// when file source type is BUFFER
7882
std::shared_ptr<Buffer> buffer() const { return buffer_; }
7983

84+
/// \brief Get a RandomAccessFile which views this file source
85+
Status Open(std::shared_ptr<arrow::io::RandomAccessFile>* out) const;
86+
8087
private:
8188
explicit FileSource(SourceType type,
8289
Compression::type compression = Compression::UNCOMPRESSED)
@@ -86,7 +93,7 @@ class ARROW_DS_EXPORT FileSource {
8693

8794
// PATH-based source
8895
std::string path_;
89-
fs::FileSystem* filesystem_;
96+
fs::FileSystem* filesystem_ = NULLPTR;
9097

9198
// BUFFER-based source
9299
std::shared_ptr<Buffer> buffer_;
@@ -119,7 +126,7 @@ class ARROW_DS_EXPORT FileFormat {
119126
virtual bool IsKnownExtension(const std::string& ext) const = 0;
120127

121128
/// \brief Open a file for scanning
122-
virtual Status ScanFile(const FileSource& location,
129+
virtual Status ScanFile(const FileSource& source,
123130
std::shared_ptr<ScanOptions> scan_options,
124131
std::shared_ptr<ScanContext> scan_context,
125132
std::unique_ptr<ScanTaskIterator>* out) const = 0;
@@ -128,14 +135,20 @@ class ARROW_DS_EXPORT FileFormat {
128135
/// \brief A DataFragment that is stored in a file with a known format
129136
class ARROW_DS_EXPORT FileBasedDataFragment : public DataFragment {
130137
public:
131-
FileBasedDataFragment(const FileSource& location, std::shared_ptr<FileFormat> format,
132-
std::shared_ptr<ScanOptions>);
138+
FileBasedDataFragment(const FileSource& source, std::shared_ptr<FileFormat> format,
139+
std::shared_ptr<ScanOptions> scan_options)
140+
: source_(source),
141+
format_(std::move(format)),
142+
scan_options_(std::move(scan_options)) {}
143+
144+
Status Scan(std::shared_ptr<ScanContext> scan_context,
145+
std::unique_ptr<ScanTaskIterator>* out) override;
133146

134-
const FileSource& location() const { return location_; }
147+
const FileSource& source() const { return source_; }
135148
std::shared_ptr<FileFormat> format() const { return format_; }
136149

137150
protected:
138-
FileSource location_;
151+
FileSource source_;
139152
std::shared_ptr<FileFormat> format_;
140153
std::shared_ptr<ScanOptions> scan_options_;
141154
};

‎cpp/src/arrow/dataset/file_csv.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class ARROW_DS_EXPORT CsvFileFormat : public FileFormat {
6060
bool IsKnownExtension(const std::string& ext) const override;
6161

6262
/// \brief Open a file for scanning
63-
Status ScanFile(const FileSource& location, std::shared_ptr<ScanOptions> scan_options,
63+
Status ScanFile(const FileSource& source, std::shared_ptr<ScanOptions> scan_options,
6464
std::shared_ptr<ScanContext> scan_context,
6565
std::unique_ptr<ScanTaskIterator>* out) const override;
6666
};

0 commit comments

Comments
 (0)