读取和写入数据集

本节包含一些用于读取和写入数据集的食谱。数据集是包含表格数据的多个文件集合。

读取分区数据集

构成数据集的单个数据文件通常会根据某种分区方案分布在多个不同的目录中。

这简化了数据的管理,并且还允许通过检查文件路径并利用分区方案提供的保证来进行数据集的部分读取。

此食谱演示了读取分区数据集的基本知识。首先让我们检查一下我们的数据

我们数据集中文件的列表
const std::string& directory_base = airquality_basedir;

// Create a filesystem
std::shared_ptr<arrow::fs::LocalFileSystem> fs =
    std::make_shared<arrow::fs::LocalFileSystem>();

// Create a file selector which describes which files are part of
// the dataset.  This selector performs a recursive search of a base
// directory which is typical with partitioned datasets.  You can also
// create a dataset from a list of one or more paths.
arrow::fs::FileSelector selector;
selector.base_dir = directory_base;
selector.recursive = true;

// List out the files so we can see how our data is partitioned.
// This step is not necessary for reading a dataset
ARROW_ASSIGN_OR_RAISE(std::vector<arrow::fs::FileInfo> file_infos,
                      fs->GetFileInfo(selector));
int num_printed = 0;
for (const auto& path : file_infos) {
  if (path.IsFile()) {
    rout << path.path().substr(directory_base.size()) << std::endl;
    if (++num_printed == 10) {
      rout << "..." << std::endl;
      break;
    }
  }
}
代码输出
/Month=8/Day=15/chunk-0.parquet
/Month=8/Day=20/chunk-0.parquet
/Month=8/Day=24/chunk-0.parquet
/Month=8/Day=23/chunk-0.parquet
/Month=8/Day=16/chunk-0.parquet
/Month=8/Day=13/chunk-0.parquet
/Month=8/Day=25/chunk-0.parquet
/Month=8/Day=18/chunk-0.parquet
/Month=8/Day=1/chunk-0.parquet
/Month=8/Day=17/chunk-0.parquet
...

注意

这种 key=value 的分区方案在 Arrow 中称为“hive”分区。

现在我们有了文件系统和选择器,我们可以继续创建数据集。为此,我们需要选择一种格式和一种分区方案。一旦我们拥有了所有必要的组件,我们就可以创建一个 arrow::dataset::Dataset 实例。

创建 arrow::dataset::Dataset 实例
// Create a file format which describes the format of the files.
// Here we specify we are reading parquet files.  We could pick a different format
// such as Arrow-IPC files or CSV files or we could customize the parquet format with
// additional reading & parsing options.
std::shared_ptr<arrow::dataset::ParquetFileFormat> format =
    std::make_shared<arrow::dataset::ParquetFileFormat>();

// Create a partitioning factory.  A partitioning factory will be used by a dataset
// factory to infer the partitioning schema from the filenames.  All we need to
// specify is the flavor of partitioning which, in our case, is "hive".
//
// Alternatively, we could manually create a partitioning scheme from a schema.  This
// is typically not necessary for hive partitioning as inference works well.
std::shared_ptr<arrow::dataset::PartitioningFactory> partitioning_factory =
    arrow::dataset::HivePartitioning::MakeFactory();

arrow::dataset::FileSystemFactoryOptions options;
options.partitioning = partitioning_factory;

// Create a dataset factory
ARROW_ASSIGN_OR_RAISE(
    std::shared_ptr<arrow::dataset::DatasetFactory> dataset_factory,
    arrow::dataset::FileSystemDatasetFactory::Make(fs, selector, format, options));

// Create the dataset, this will scan the dataset directory to find all the files
// and may scan some file metadata in order to determine the dataset schema.
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::dataset::Dataset> dataset,
                      dataset_factory->Finish());

rout << "We discovered the following schema for the dataset:" << std::endl
     << std::endl
     << dataset->schema()->ToString() << std::endl;
代码输出
We discovered the following schema for the dataset:

Ozone: int32
Solar.R: int32
Wind: double
Temp: int32
Month: int32
Day: int32

一旦我们有了数据集对象,我们就可以读取数据。从数据集读取数据有时称为“扫描”数据集,我们用来执行此操作的对象是 arrow::dataset::Scanner。以下代码片段展示了如何将整个数据集扫描到内存中的表格中

将数据集扫描到 arrow::Table 中
// Create a scanner
arrow::dataset::ScannerBuilder scanner_builder(dataset);
ARROW_RETURN_NOT_OK(scanner_builder.UseThreads(true));
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::dataset::Scanner> scanner,
                      scanner_builder.Finish());

// Scan the dataset.  There are a variety of other methods available on the scanner as
// well
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Table> table, scanner->ToTable());
rout << "Read in a table with " << table->num_rows() << " rows and "
     << table->num_columns() << " columns";
代码输出
Read in a table with 153 rows and 6 columns