Merge pull request ClickHouse#83653 from scanhex12/iceberg_field_id

Read iceberg data files by field ids

Author: scanhex12

src/Formats/FormatParserGroup.cpp

@@ -17,6 +17,32 @@ namespace Setting
     extern const SettingsMaxThreads max_parsing_threads;
 }
 
+void ColumnMapper::setStorageColumnEncoding(std::unordered_map<String, Int64> && storage_encoding_)
+{
+    storage_encoding = std::move(storage_encoding_);
+}
+
+std::pair<std::unordered_map<String, String>, std::unordered_map<String, String>> ColumnMapper::makeMapping(
+    const Block & header,
+    const std::unordered_map<Int64, String> & format_encoding)
+{
+    std::unordered_map<String, String> clickhouse_to_parquet_names;
+    std::unordered_map<String, String> parquet_names_to_clickhouse;
+
+    for (size_t i = 0; i < header.columns(); ++i)
+    {
+        auto column_name = header.getNames()[i];
+        int64_t field_id;
+        if (auto it = storage_encoding.find(column_name); it != storage_encoding.end())
+            field_id = it->second;
+        else
+            continue;
+        clickhouse_to_parquet_names[column_name] = format_encoding.at(field_id);
+        parquet_names_to_clickhouse.emplace(format_encoding.at(field_id), column_name);
+    }
+    return {clickhouse_to_parquet_names, parquet_names_to_clickhouse};
+}
+
 FormatParserGroup::FormatParserGroup(const Settings & settings, size_t num_streams_, std::shared_ptr<const ActionsDAG> filter_actions_dag_, const ContextPtr & context_)
     : max_parsing_threads(settings[Setting::max_parsing_threads])
     , max_io_threads(settings[Setting::max_download_threads])

src/Formats/FormatParserGroup.h

@@ -13,6 +13,24 @@ struct FormatParserGroup;
 
 using FormatParserGroupPtr = std::shared_ptr<FormatParserGroup>;
 
+/// Some formats needs to custom mapping between columns in file and clickhouse columns.
+class ColumnMapper
+{
+public:
+    /// clickhouse_column_name -> field_id
+    void setStorageColumnEncoding(std::unordered_map<String, Int64> && storage_encoding_);
+
+    /// clickhouse_column_name -> format_column_name (just join the maps above by field_id).
+    std::pair<std::unordered_map<String, String>, std::unordered_map<String, String>> makeMapping(
+        const Block & header,
+        const std::unordered_map<Int64, String> & format_encoding);
+
+private:
+    std::unordered_map<String, Int64> storage_encoding;
+};
+
+using ColumnMapperPtr = std::shared_ptr<ColumnMapper>;
+
 /// When reading many files in one query, e.g. `SELECT ... FROM file('part{00..99}.parquet')`,
 /// we want the file readers to share some resource limits, e.g. number of threads.
 /// They may also want to share some data structures to avoid initializing multiple copies,
@@ -41,6 +59,8 @@ struct FormatParserGroup
     /// IInputFormat implementation may put arbitrary state here.
     std::shared_ptr<void> opaque;
 
+    ColumnMapperPtr column_mapper;
+
 private:
     /// For lazily initializing the fields above.
     std::once_flag init_flag;

src/Processors/Formats/Impl/ArrowBlockInputFormat.cpp

@@ -174,6 +174,7 @@ void ArrowBlockInputFormat::prepareReader()
         getPort().getHeader(),
         "Arrow",
         format_settings,
+        std::nullopt,
         format_settings.arrow.allow_missing_columns,
         format_settings.null_as_default,
         format_settings.date_time_overflow_behavior,

src/Processors/Formats/Impl/ArrowColumnToCHColumn.cpp

@@ -1468,6 +1468,7 @@ ArrowColumnToCHColumn::ArrowColumnToCHColumn(
     const Block & header_,
     const std::string & format_name_,
     const FormatSettings & format_settings_,
+    const std::optional<std::unordered_map<String, String>> & parquet_columns_to_clickhouse_,
     bool allow_missing_columns_,
     bool null_as_default_,
     FormatSettings::DateTimeOverflowBehavior date_time_overflow_behavior_,
@@ -1485,6 +1486,7 @@ ArrowColumnToCHColumn::ArrowColumnToCHColumn(
     , case_insensitive_matching(case_insensitive_matching_)
     , is_stream(is_stream_)
     , enable_json_parsing(enable_json_parsing_)
+    , parquet_columns_to_clickhouse(parquet_columns_to_clickhouse_)
 {
 }
 
@@ -1504,12 +1506,14 @@ Chunk ArrowColumnToCHColumn::arrowTableToCHChunk(
 
         auto arrow_field = table->schema()->GetFieldByName(column_name);
 
+        if (parquet_columns_to_clickhouse)
+            column_name = parquet_columns_to_clickhouse->at(column_name);
+
         if (case_insensitive_matching)
             boost::to_lower(column_name);
 
         name_to_arrow_column[std::move(column_name)] = {std::move(arrow_column), std::move(arrow_field)};
     }
-
     return arrowColumnsToCHChunk(name_to_arrow_column, num_rows, metadata, block_missing_values);
 }
 

src/Processors/Formats/Impl/ArrowColumnToCHColumn.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstdint>
+#include <unordered_map>
 
 #include "config.h"
 
@@ -25,6 +26,7 @@ class ArrowColumnToCHColumn
         const Block & header_,
         const std::string & format_name_,
         const FormatSettings & format_settings_,
+        const std::optional<std::unordered_map<String, String>> & parquet_columns_to_clickhouse_,
         bool allow_missing_columns_,
         bool null_as_default_,
         FormatSettings::DateTimeOverflowBehavior date_time_overflow_behavior_,
@@ -90,6 +92,8 @@ class ArrowColumnToCHColumn
     /// To avoid converting dictionary from Arrow Dictionary
     /// to LowCardinality every chunk we save it and reuse.
     std::unordered_map<std::string, DictionaryInfo> dictionary_infos;
+
+    std::optional<std::unordered_map<String, String>> parquet_columns_to_clickhouse;
 };
 
 }

src/Processors/Formats/Impl/ArrowFieldIndexUtil.h

@@ -80,7 +80,8 @@ class ArrowFieldIndexUtil
     std::vector<ClickHouseIndexToParquetIndex> findRequiredIndices(
         const Block & header,
         const arrow::Schema & schema,
-        const parquet::FileMetaData & file)
+        const parquet::FileMetaData & file,
+        const std::optional<std::unordered_map<String, String>> & clickhouse_to_parquet_names)
     {
         std::vector<ClickHouseIndexToParquetIndex> required_indices;
         std::unordered_set<int> added_indices;
@@ -90,6 +91,13 @@ class ArrowFieldIndexUtil
         {
             const auto & named_col = header.getByPosition(i);
             std::string col_name = named_col.name;
+            if (clickhouse_to_parquet_names)
+            {
+                if (auto it = clickhouse_to_parquet_names->find(col_name); it != clickhouse_to_parquet_names->end())
+                    col_name = it->second;
+                else
+                    continue;
+            }
             if (ignore_case)
                 boost::to_lower(col_name);
             findRequiredIndices(col_name, i, named_col.type, fields_indices, added_indices, required_indices, file);

src/Processors/Formats/Impl/ORCBlockInputFormat.cpp

@@ -139,6 +139,7 @@ void ORCBlockInputFormat::prepareReader()
         getPort().getHeader(),
         "ORC",
         format_settings,
+        std::nullopt,
         format_settings.orc.allow_missing_columns,
         format_settings.null_as_default,
         format_settings.date_time_overflow_behavior,

src/Processors/Formats/Impl/Parquet/ParquetRecordReader.cpp

@@ -318,6 +318,7 @@ ParquetRecordReader::ParquetRecordReader(
     std::shared_ptr<::arrow::io::RandomAccessFile> arrow_file,
     const FormatSettings & format_settings,
     std::vector<int> row_groups_indices_,
+    const std::optional<std::vector<Int32>> & column_indices_,
     std::shared_ptr<parquet::FileMetaData> metadata)
     : file_reader(createFileReader(std::move(arrow_file), reader_properties_, std::move(metadata)))
     , arrow_properties(arrow_properties_)
@@ -338,20 +339,26 @@ ParquetRecordReader::ParquetRecordReader(
 
     parquet_col_indice.reserve(header.columns());
     column_readers.reserve(header.columns());
-    for (const auto & col_with_name : header)
+    if (!column_indices_)
     {
-        auto it = parquet_columns.find(col_with_name.name);
-        if (it == parquet_columns.end())
-            throw Exception(ErrorCodes::PARQUET_EXCEPTION, "no column with '{}' in parquet file", col_with_name.name);
+        for (const auto & col_with_name : header)
+        {
+            auto it = parquet_columns.find(col_with_name.name);
+            if (it == parquet_columns.end())
+                throw Exception(ErrorCodes::PARQUET_EXCEPTION, "no column with '{}' in parquet file", col_with_name.name);
 
-        const auto & node = it->second;
-        if (!node->is_primitive())
-            throw Exception(ErrorCodes::NOT_IMPLEMENTED, "arrays and maps are not implemented in native parquet reader");
+            const auto & node = it->second;
+            if (!node->is_primitive())
+                throw Exception(ErrorCodes::NOT_IMPLEMENTED, "arrays and maps are not implemented in native parquet reader");
 
-        auto idx = file_reader->metadata()->schema()->ColumnIndex(*node);
-        chassert(idx >= 0);
-        parquet_col_indice.push_back(idx);
+            auto idx = file_reader->metadata()->schema()->ColumnIndex(*node);
+            chassert(idx >= 0);
+            parquet_col_indice.push_back(idx);
+        }
     }
+    else
+        parquet_col_indice = *column_indices_;
+
     if (arrow_properties.pre_buffer())
     {
         THROW_PARQUET_EXCEPTION(file_reader->PreBuffer(

src/Processors/Formats/Impl/Parquet/ParquetRecordReader.h

@@ -25,6 +25,7 @@ class ParquetRecordReader
         std::shared_ptr<::arrow::io::RandomAccessFile> arrow_file,
         const FormatSettings & format_settings,
         std::vector<int> row_groups_indices_,
+        const std::optional<std::vector<Int32>> & column_indices_,
         std::shared_ptr<parquet::FileMetaData> metadata = nullptr);
 
     Chunk readChunk();

src/Processors/Formats/Impl/ParquetBlockInputFormat.cpp

@@ -1,4 +1,5 @@
-#include "ParquetBlockInputFormat.h"
+#include <memory>
+#include <Processors/Formats/Impl/ParquetBlockInputFormat.h>
 
 #if USE_PARQUET
 
@@ -632,7 +633,22 @@ void ParquetBlockInputFormat::initializeIfNeeded()
         format_settings.parquet.case_insensitive_column_matching,
         format_settings.parquet.allow_missing_columns);
 
-    auto index_mapping = field_util.findRequiredIndices(getPort().getHeader(), *schema, *metadata);
+    std::optional<std::unordered_map<String, String>> clickhouse_to_parquet_names;
+    if (parser_group && parser_group->column_mapper)
+    {
+        auto header = getPort().getHeader();
+        const auto & group_node = metadata->schema()->group_node();
+
+        std::unordered_map<Int64, String> parquet_field_ids;
+        parquet_names_to_clickhouse = std::unordered_map<String, String>{};
+        for (int i = 0; i < group_node->field_count(); ++i)
+            parquet_field_ids[group_node->field(i)->field_id()] = group_node->field(i)->name();
+
+        auto result = parser_group->column_mapper->makeMapping(header, parquet_field_ids);
+        clickhouse_to_parquet_names = std::move(result.first);
+        parquet_names_to_clickhouse = std::move(result.second);
+    }
+    auto index_mapping = field_util.findRequiredIndices(getPort().getHeader(), *schema, *metadata, clickhouse_to_parquet_names);
 
     for (const auto & [clickhouse_header_index, parquet_indexes] : index_mapping)
     {
@@ -848,7 +864,8 @@ void ParquetBlockInputFormat::initializeRowGroupBatchReader(size_t row_group_bat
             reader_properties,
             arrow_file,
             format_settings,
-            row_group_batch.row_groups_idxs);
+            row_group_batch.row_groups_idxs,
+            column_indices);
     }
     else
     {
@@ -875,6 +892,7 @@ void ParquetBlockInputFormat::initializeRowGroupBatchReader(size_t row_group_bat
             getPort().getHeader(),
             "Parquet",
             format_settings,
+            parquet_names_to_clickhouse,
             format_settings.parquet.allow_missing_columns,
             format_settings.null_as_default,
             format_settings.date_time_overflow_behavior,