LazyParquetChunkLoader.h

Go to the documentation of this file.

/*
 * Copyright 2022 HEAVY.AI, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <arrow/filesystem/filesystem.h>
#include <parquet/schema.h>

#include "DataMgr/Chunk/Chunk.h"
#include "DataPreview.h"
#include "ForeignTableSchema.h"
#include "Interval.h"
#include "ParquetEncoder.h"
#include "ParquetShared.h"
#include "StringDictionary/StringDictionary.h"

extern size_t g_max_import_threads;

namespace foreign_storage {

class LazyParquetChunkLoader {
 public:
  // The number of elements in a batch that are read from the Parquet file;
  // this number is subject to change with performance tuning.
  // Most filesystems use a default block size of 4096 bytes.
  const static int batch_reader_num_elements = 4096;

  LazyParquetChunkLoader(std::shared_ptr<arrow::fs::FileSystem> file_system,
                         FileReaderMap* file_reader_cache,
                         const ForeignTable* foreign_table);

  std::list<std::unique_ptr<ChunkMetadata>> loadChunk(
      const std::vector<RowGroupInterval>& row_group_intervals,
      const int parquet_column_index,
      std::list<Chunk_NS::Chunk>& chunks,
      StringDictionary* string_dictionary = nullptr,
      RejectedRowIndices* rejected_row_indices = nullptr);

  std::list<RowGroupMetadata> metadataScan(
      const std::vector<std::string>& file_paths,
      const ForeignTableSchema& schema,
      const bool do_metadata_stats_validation = true);

  static bool isColumnMappingSupported(const ColumnDescriptor* omnisci_column,
                                       const parquet::ColumnDescriptor* parquet_column);

  std::pair<size_t, size_t> loadRowGroups(
      const RowGroupInterval& row_group_interval,
      const std::map<int, Chunk_NS::Chunk>& chunks,
      const ForeignTableSchema& schema,
      const std::map<int, StringDictionary*>& column_dictionaries,
      const int num_threads = 1);

  DataPreview previewFiles(const std::vector<std::string>& files,
                           const size_t max_num_rows,
                           const ForeignTable& table);

 private:
  static SQLTypeInfo suggestColumnMapping(
      const parquet::ColumnDescriptor* parquet_column);

  std::list<std::unique_ptr<ChunkMetadata>> appendRowGroups(
      const std::vector<RowGroupInterval>& row_group_intervals,
      const int parquet_column_index,
      const ColumnDescriptor* column_descriptor,
      std::list<Chunk_NS::Chunk>& chunks,
      StringDictionary* string_dictionary,
      RejectedRowIndices* rejected_row_indices,
      const bool is_for_detect = false,
      const std::optional<int64_t> max_levels_read = std::nullopt);

  std::shared_ptr<arrow::fs::FileSystem> file_system_;
  FileReaderMap* file_reader_cache_;

  const ForeignTable* foreign_table_;
};
}  // namespace foreign_storage