[C++][Parquet] Fast-path fixed-length LIST reads

cpp/src/parquet/arrow/reader.cc

@@ -124,6 +124,16 @@ class ColumnReaderImpl : public ColumnReader {
 
   virtual ::arrow::Status LoadBatch(int64_t num_records) = 0;
 
+  virtual ::arrow::Status LoadBatchFixedLength(int64_t num_values,
+                                               int64_t expected_list_size) {
+    return LoadBatch(num_values);
+  }
+
+  virtual ::arrow::Status LoadBatchDetectFixedLength(int64_t num_records,
+                                                     int64_t* list_size) {
+    return Status::NotImplemented("fixed-length LIST detection is unavailable");
+  }
+
   virtual ::arrow::Status BuildArray(int64_t length_upper_bound,
                                      std::shared_ptr<::arrow::ChunkedArray>* out) = 0;
   virtual bool IsOrHasRepeatedChild() const = 0;
@@ -495,7 +505,7 @@ class LeafReader : public ColumnReaderImpl {
 
   bool IsOrHasRepeatedChild() const final { return false; }
 
-  Status LoadBatch(int64_t records_to_read) final {
+  Status LoadBatch(int64_t records_to_read) override {
     BEGIN_PARQUET_CATCH_EXCEPTIONS
     out_ = nullptr;
     record_reader_->Reset();
@@ -530,6 +540,57 @@ class LeafReader : public ColumnReaderImpl {
     END_PARQUET_CATCH_EXCEPTIONS
   }
 
+  Status LoadBatchFixedLength(int64_t values_to_read, int64_t expected_list_size) final {
+    BEGIN_PARQUET_CATCH_EXCEPTIONS
+    out_ = nullptr;
+    record_reader_->Reset();
+    record_reader_->Reserve(values_to_read);
+    int64_t remaining = values_to_read;
+    while (remaining > 0 && record_reader_->HasMoreData()) {
+      const int64_t values_read =
+          record_reader_->ReadRecordsFixedLength(remaining, expected_list_size);
+      if (values_read < 0) {
+        if (remaining != values_to_read) {
+          return Status::Invalid(
+              "fixed-length LIST detector rejected a later page after decoding values; "
+              "disable ArrowReaderProperties::set_fixed_list_fast_path() for "
+              "variable-length "
+              "LIST columns");
+        }
+        return Status::NotImplemented(
+            "fixed-size list fast path is unavailable for this leaf");
+      }
+      remaining -= values_read;
+      if (values_read == 0) {
+        NextRowGroup();
+      }
+    }
+    if (remaining != 0) {
+      return Status::Invalid(
+          "Fixed-size list fast path reached end of column before reading ",
+          values_to_read,
+          " values; disable ArrowReaderProperties::set_fixed_list_fast_path() "
+          "for variable-length LIST columns");
+    }
+    RETURN_NOT_OK(TransferColumnData(record_reader_.get(), nullptr, field_, descr_,
+                                     ctx_.get(), &out_));
+    return Status::OK();
+    END_PARQUET_CATCH_EXCEPTIONS
+  }
+
+  Status LoadBatchDetectFixedLength(int64_t records_to_read, int64_t* list_size) final {
+    BEGIN_PARQUET_CATCH_EXCEPTIONS
+    out_ = nullptr;
+    record_reader_->Reset();
+    const int64_t detected = record_reader_->DetectFixedLength();
+    if (detected <= 0) {
+      return Status::NotImplemented("page is not effectively fixed-length");
+    }
+    *list_size = detected;
+    return LoadBatchFixedLength(records_to_read * detected, detected);
+    END_PARQUET_CATCH_EXCEPTIONS
+  }
+
   ::arrow::Status BuildArray(int64_t length_upper_bound,
                              std::shared_ptr<::arrow::ChunkedArray>* out) final {
     *out = out_;
@@ -567,7 +628,7 @@ class ExtensionReader : public ColumnReaderImpl {
     return storage_reader_->GetRepLevels(data, length);
   }
 
-  Status LoadBatch(int64_t number_of_records) final {
+  Status LoadBatch(int64_t number_of_records) override {
     return storage_reader_->LoadBatch(number_of_records);
   }
 
@@ -611,10 +672,46 @@ class ListReader : public ColumnReaderImpl {
 
   bool IsOrHasRepeatedChild() const final { return true; }
 
-  Status LoadBatch(int64_t number_of_records) final {
+  Status LoadBatch(int64_t number_of_records) override {
+    if (fixed_list_fast_path_) {
+      auto status = item_reader_->LoadBatchFixedLength(
+          number_of_records * fixed_list_size_, fixed_list_size_);
+      if (status.ok()) {
+        return status;
+      }
+      if (!status.IsNotImplemented()) {
+        return status;
+      }
+      fixed_list_fast_path_ = false;
+    }
+    if (ctx_->reader_properties->fixed_list_fast_path() &&
+        field_->type()->id() == ::arrow::Type::LIST && !field_->nullable() &&
+        !item_reader_->IsOrHasRepeatedChild()) {
+      int64_t detected_size = 0;
+      auto status = LoadBatchDetectFixedLength(number_of_records, &detected_size);
+      if (status.ok()) {
+        return status;
+      }
+    }
     return item_reader_->LoadBatch(number_of_records);
   }
 
+  Status LoadBatchDetectFixedLength(int64_t number_of_records,
+                                    int64_t* list_size) override {
+    if (field_->type()->id() != ::arrow::Type::LIST || field_->nullable() ||
+        item_reader_->IsOrHasRepeatedChild()) {
+      return Status::NotImplemented(
+          "LIST shape is not eligible for fixed-length detection");
+    }
+    auto status = item_reader_->LoadBatchDetectFixedLength(number_of_records, list_size);
+    if (!status.ok()) {
+      return status;
+    }
+    fixed_list_size_ = *list_size;
+    fixed_list_fast_path_ = true;
+    return Status::OK();
+  }
+
   virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
       std::shared_ptr<ArrayData> data) {
     if (field_->type()->id() == ::arrow::Type::MAP) {
@@ -627,6 +724,32 @@ class ListReader : public ColumnReaderImpl {
 
   Status BuildArray(int64_t length_upper_bound,
                     std::shared_ptr<ChunkedArray>* out) override {
+    if (fixed_list_fast_path_ && fixed_list_size_ > 0) {
+      std::shared_ptr<ChunkedArray> child;
+      RETURN_NOT_OK(
+          item_reader_->BuildArray(length_upper_bound * fixed_list_size_, &child));
+      ARROW_ASSIGN_OR_RAISE(auto child_data, ChunksToSingle(*child));
+      const int64_t length = child_data->length / fixed_list_size_;
+      if (length * fixed_list_size_ != child_data->length) {
+        return Status::Invalid(
+            "Detected fixed-length LIST child count is not divisible by ",
+            fixed_list_size_);
+      }
+      std::shared_ptr<ResizableBuffer> offsets_buffer;
+      ARROW_ASSIGN_OR_RAISE(
+          offsets_buffer,
+          AllocateResizableBuffer(sizeof(IndexType) * (length + 1), ctx_->pool));
+      auto* offsets = reinterpret_cast<IndexType*>(offsets_buffer->mutable_data());
+      for (int64_t i = 0; i <= length; ++i) {
+        offsets[i] = static_cast<IndexType>(i * fixed_list_size_);
+      }
+      auto data = std::make_shared<ArrayData>(
+          field_->type(), length,
+          std::vector<std::shared_ptr<Buffer>>{nullptr, offsets_buffer},
+          std::vector<std::shared_ptr<ArrayData>>{std::move(child_data)});
+      *out = std::make_shared<ChunkedArray>(::arrow::MakeArray(std::move(data)));
+      return Status::OK();
+    }
     const int16_t* def_levels;
     const int16_t* rep_levels;
     int64_t num_levels;
@@ -683,11 +806,14 @@ class ListReader : public ColumnReaderImpl {
 
  protected:
   std::shared_ptr<ReaderContext> ctx_;
+  ColumnReaderImpl* item_reader() const { return item_reader_.get(); }
 
  private:
   std::shared_ptr<Field> field_;
   ::parquet::internal::LevelInfo level_info_;
   std::unique_ptr<ColumnReaderImpl> item_reader_;
+  int64_t fixed_list_size_ = 0;
+  bool fixed_list_fast_path_ = false;
 };
 
 template <typename IndexType>
@@ -727,6 +853,50 @@ class PARQUET_NO_EXPORT FixedSizeListReader : public ListReader<int32_t> {
                       std::unique_ptr<ColumnReaderImpl> child_reader)
       : ListReader(std::move(ctx), std::move(field), level_info,
                    std::move(child_reader)) {}
+
+  Status LoadBatch(int64_t records_to_read) final {
+    const auto& type = checked_cast<const ::arrow::FixedSizeListType&>(*field()->type());
+    if (ctx_->reader_properties->fixed_list_fast_path() && !field()->nullable() &&
+        !item_reader()->field()->nullable() && type.list_size() > 0) {
+      auto status = item_reader()->LoadBatchFixedLength(
+          records_to_read * type.list_size(), type.list_size());
+      if (status.IsNotImplemented()) {
+        return ListReader<int32_t>::LoadBatch(records_to_read);
+      }
+      return status;
+    }
+    return ListReader<int32_t>::LoadBatch(records_to_read);
+  }
+
+  // A FixedSizeList schema already carries the invariant that every non-null
+  // slot contributes exactly list_size child values.  When the list itself is
+  // required, there is no need to decode or materialize Dremel levels: the
+  // child reader has already produced a flat value array and its length gives
+  // us the number of records.  This is the same fast path used by readers
+  // which detect effectively fixed-size LIST pages from their encoded levels.
+  Status BuildArray(int64_t length_upper_bound,
+                    std::shared_ptr<::arrow::ChunkedArray>* out) final {
+    const auto& type = checked_cast<const ::arrow::FixedSizeListType&>(*field()->type());
+    const int32_t list_size = type.list_size();
+    if (ctx_->reader_properties->fixed_list_fast_path() && !field()->nullable() &&
+        list_size > 0) {
+      std::shared_ptr<ChunkedArray> child;
+      RETURN_NOT_OK(item_reader()->BuildArray(length_upper_bound * list_size, &child));
+      ARROW_ASSIGN_OR_RAISE(auto child_data, ChunksToSingle(*child));
+      if (child_data->length % list_size != 0) {
+        return Status::Invalid("Fixed-size list child value count ", child_data->length,
+                               " is not divisible by list size ", list_size);
+      }
+      auto data = std::make_shared<ArrayData>(
+          field()->type(), child_data->length / list_size,
+          std::vector<std::shared_ptr<Buffer>>{nullptr},
+          std::vector<std::shared_ptr<ArrayData>>{std::move(child_data)});
+      *out = std::make_shared<ChunkedArray>(::arrow::MakeArray(std::move(data)));
+      return Status::OK();
+    }
+    return ListReader<int32_t>::BuildArray(length_upper_bound, out);
+  }
+
   ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
       std::shared_ptr<ArrayData> data) final {
     DCHECK_EQ(data->buffers.size(), 2);

cpp/src/parquet/arrow/reader_writer_benchmark.cc

@@ -285,7 +285,8 @@ struct Examples<bool> {
 
 static void BenchmarkReadTable(::benchmark::State& state, const Table& table,
                                std::shared_ptr<WriterProperties> properties,
-                               int64_t num_values = -1, int64_t total_bytes = -1) {
+                               int64_t num_values = -1, int64_t total_bytes = -1,
+                               bool fixed_list_fast_path = false) {
   // Make sure we roundtrip Arrow types by storing the schema
   auto arrow_properties = ArrowWriterProperties::Builder().store_schema()->build();
   auto output = CreateOutputStream();
@@ -296,8 +297,10 @@ static void BenchmarkReadTable(::benchmark::State& state, const Table& table,
   for (auto _ : state) {
     auto reader =
         ParquetFileReader::Open(std::make_shared<::arrow::io::BufferReader>(buffer));
-    auto arrow_reader_result =
-        FileReader::Make(::arrow::default_memory_pool(), std::move(reader));
+    ArrowReaderProperties reader_properties;
+    reader_properties.set_fixed_list_fast_path(fixed_list_fast_path);
+    auto arrow_reader_result = FileReader::Make(::arrow::default_memory_pool(),
+                                                std::move(reader), reader_properties);
     EXIT_NOT_OK(arrow_reader_result.status());
     auto arrow_reader = std::move(*arrow_reader_result);
 
@@ -680,6 +683,103 @@ static void BM_ReadListColumn(::benchmark::State& state) {
 
 BENCHMARK(BM_ReadListColumn)->Apply(NestedReadArguments);
 
+// Fixed-size-list reader benchmark.  This exercises the same physical Parquet
+// LIST encoding as a vector column while exposing the fixed-size invariant to
+// the Arrow reader.  The reader fast path can therefore avoid decoding the
+// definition/repetition levels and rebuilding offsets.
+static std::shared_ptr<Table> MakeVectorFloatTable(int32_t list_size, bool fixed_size) {
+  constexpr int64_t kTotalValues = 4 * 1024 * 1024;
+  const int64_t num_records = kTotalValues / list_size;
+  NumericBuilder<::arrow::FloatType> builder;
+  std::vector<float> values(static_cast<size_t>(num_records * list_size));
+  for (size_t i = 0; i < values.size(); ++i) {
+    values[i] = static_cast<float>((i * 17) % 100003) / 101.0f;
+  }
+  EXIT_NOT_OK(builder.AppendValues(values.data(), static_cast<int64_t>(values.size())));
+  std::shared_ptr<Array> child;
+  EXIT_NOT_OK(builder.Finish(&child));
+  auto child_field = ::arrow::field("item", ::arrow::float32(), /*nullable=*/false);
+  std::shared_ptr<::arrow::DataType> type;
+  std::shared_ptr<Array> array;
+  if (fixed_size) {
+    type = ::arrow::fixed_size_list(child_field, list_size);
+    auto array_result = ::arrow::FixedSizeListArray::FromArrays(child, type);
+    EXIT_NOT_OK(array_result.status());
+    array = *array_result;
+  } else {
+    type = ::arrow::list(child_field);
+    std::shared_ptr<::arrow::Buffer> offsets =
+        ::arrow::AllocateBuffer(sizeof(int32_t) * (num_records + 1)).ValueOrDie();
+    auto* offset_data = offsets->mutable_data_as<int32_t>();
+    for (int64_t i = 0; i <= num_records; ++i) {
+      offset_data[i] = static_cast<int32_t>(i * list_size);
+    }
+    auto data = ::arrow::ArrayData::Make(
+        type, num_records,
+        std::vector<std::shared_ptr<::arrow::Buffer>>{nullptr, std::move(offsets)},
+        std::vector<std::shared_ptr<::arrow::ArrayData>>{child->data()}, nullptr, 0);
+    array = ::arrow::MakeArray(std::move(data));
+  }
+  auto schema = ::arrow::schema({field("vector", type, /*nullable=*/false)});
+  return Table::Make(schema, {std::move(array)}, num_records);
+}
+
+template <Encoding::type encoding>
+static void BM_ReadFixedSizeFloatList(::benchmark::State& state) {
+  const int32_t list_size = static_cast<int32_t>(state.range(0));
+  auto table = MakeVectorFloatTable(list_size, /*fixed_size=*/true);
+  auto properties = WriterProperties::Builder()
+                        .disable_dictionary()
+                        ->compression(Compression::ZSTD)
+                        ->encoding(encoding)
+                        ->build();
+  BenchmarkReadTable(state, *table, properties, table->num_rows(),
+                     table->num_rows() * list_size * sizeof(float),
+                     /*fixed_list_fast_path=*/true);
+}
+
+static void BM_ReadFixedSizeFloatListPlain(::benchmark::State& state) {
+  BM_ReadFixedSizeFloatList<Encoding::PLAIN>(state);
+}
+static void BM_ReadFixedSizeFloatListByteStreamSplit(::benchmark::State& state) {
+  BM_ReadFixedSizeFloatList<Encoding::BYTE_STREAM_SPLIT>(state);
+}
+static void BM_ReadFixedSizeFloatListAlp(::benchmark::State& state) {
+  BM_ReadFixedSizeFloatList<Encoding::ALP>(state);
+}
+
+BENCHMARK(BM_ReadFixedSizeFloatListPlain)->Arg(3)->Arg(16)->Arg(768);
+BENCHMARK(BM_ReadFixedSizeFloatListByteStreamSplit)->Arg(3)->Arg(16)->Arg(768);
+BENCHMARK(BM_ReadFixedSizeFloatListAlp)->Arg(3)->Arg(16)->Arg(768);
+
+template <Encoding::type encoding>
+static void BM_ReadFloatList(::benchmark::State& state) {
+  const int32_t list_size = static_cast<int32_t>(state.range(0));
+  auto table = MakeVectorFloatTable(list_size, /*fixed_size=*/false);
+  auto properties = WriterProperties::Builder()
+                        .disable_dictionary()
+                        ->compression(Compression::ZSTD)
+                        ->encoding(encoding)
+                        ->build();
+  BenchmarkReadTable(state, *table, properties, table->num_rows(),
+                     table->num_rows() * list_size * sizeof(float),
+                     /*fixed_list_fast_path=*/true);
+}
+
+static void BM_ReadFloatListPlain(::benchmark::State& state) {
+  BM_ReadFloatList<Encoding::PLAIN>(state);
+}
+static void BM_ReadFloatListByteStreamSplit(::benchmark::State& state) {
+  BM_ReadFloatList<Encoding::BYTE_STREAM_SPLIT>(state);
+}
+static void BM_ReadFloatListAlp(::benchmark::State& state) {
+  BM_ReadFloatList<Encoding::ALP>(state);
+}
+
+BENCHMARK(BM_ReadFloatListPlain)->Arg(3)->Arg(16)->Arg(768);
+BENCHMARK(BM_ReadFloatListByteStreamSplit)->Arg(3)->Arg(16)->Arg(768);
+BENCHMARK(BM_ReadFloatListAlp)->Arg(3)->Arg(16)->Arg(768);
+
 static void BM_ReadListOfStructColumn(::benchmark::State& state) {
   constexpr int64_t kNumValues = BENCHMARK_SIZE / 10;
   const double null_probability = static_cast<double>(state.range(0)) / 100.0;