Merge pull request ClickHouse#25902 from Avogar/arrow-nested

Refactor ArrowColumnToCHColumn, support inserting Nested as Array(Struct) in Arrow/ORC/Parquet

Author: nikitamikhaylov

src/Columns/ColumnLowCardinality.h

@@ -194,6 +194,7 @@ class ColumnLowCardinality final : public COWHelper<IColumn, ColumnLowCardinalit
     const IColumnUnique & getDictionary() const { return dictionary.getColumnUnique(); }
     IColumnUnique & getDictionary() { return dictionary.getColumnUnique(); }
     const ColumnPtr & getDictionaryPtr() const { return dictionary.getColumnUniquePtr(); }
+    ColumnPtr & getDictionaryPtr() { return dictionary.getColumnUniquePtr(); }
     /// IColumnUnique & getUnique() { return static_cast<IColumnUnique &>(*column_unique); }
     /// ColumnPtr getUniquePtr() const { return column_unique; }
 

src/Core/Settings.h

@@ -528,6 +528,9 @@ class IColumn;
     M(Bool, input_format_tsv_empty_as_default, false, "Treat empty fields in TSV input as default values.", 0) \
     M(Bool, input_format_tsv_enum_as_number, false, "Treat inserted enum values in TSV formats as enum indices \\N", 0) \
     M(Bool, input_format_null_as_default, true, "For text input formats initialize null fields with default values if data type of this field is not nullable", 0) \
+    M(Bool, input_format_arrow_import_nested, false, "Allow to insert array of structs into Nested table in Arrow input format.", 0) \
+    M(Bool, input_format_orc_import_nested, false, "Allow to insert array of structs into Nested table in ORC input format.", 0) \
+    M(Bool, input_format_parquet_import_nested, false, "Allow to insert array of structs into Nested table in Parquet input format.", 0) \
     \
     M(DateTimeInputFormat, date_time_input_format, FormatSettings::DateTimeInputFormat::Basic, "Method to read DateTime from text input formats. Possible values: 'basic' and 'best_effort'.", 0) \
     M(DateTimeOutputFormat, date_time_output_format, FormatSettings::DateTimeOutputFormat::Simple, "Method to write DateTime to text output. Possible values: 'simple', 'iso', 'unix_timestamp'.", 0) \

src/DataTypes/NestedUtils.cpp

@@ -208,6 +208,18 @@ void validateArraySizes(const Block & block)
     }
 }
 
+std::unordered_set<String> getAllTableNames(const Block & block)
+{
+    std::unordered_set<String> nested_table_names;
+    for (auto & name : block.getNames())
+    {
+        auto nested_table_name = Nested::extractTableName(name);
+        if (!nested_table_name.empty())
+            nested_table_names.insert(nested_table_name);
+    }
+    return nested_table_names;
+}
+
 }
 
 }

src/DataTypes/NestedUtils.h

@@ -28,6 +28,9 @@ namespace Nested
 
     /// Check that sizes of arrays - elements of nested data structures - are equal.
     void validateArraySizes(const Block & block);
+
+    /// Get all nested tables names from a block.
+    std::unordered_set<String> getAllTableNames(const Block & block);
 }
 
 }

src/Formats/FormatFactory.cpp

@@ -88,6 +88,7 @@ FormatSettings getFormatSettings(ContextPtr context, const Settings & settings)
     format_settings.json.quote_denormals = settings.output_format_json_quote_denormals;
     format_settings.null_as_default = settings.input_format_null_as_default;
     format_settings.parquet.row_group_size = settings.output_format_parquet_row_group_size;
+    format_settings.parquet.import_nested = settings.input_format_parquet_import_nested;
     format_settings.pretty.charset = settings.output_format_pretty_grid_charset.toString() == "ASCII" ? FormatSettings::Pretty::Charset::ASCII : FormatSettings::Pretty::Charset::UTF8;
     format_settings.pretty.color = settings.output_format_pretty_color;
     format_settings.pretty.max_column_pad_width = settings.output_format_pretty_max_column_pad_width;
@@ -114,6 +115,8 @@ FormatSettings getFormatSettings(ContextPtr context, const Settings & settings)
     format_settings.with_names_use_header = settings.input_format_with_names_use_header;
     format_settings.write_statistics = settings.output_format_write_statistics;
     format_settings.arrow.low_cardinality_as_dictionary = settings.output_format_arrow_low_cardinality_as_dictionary;
+    format_settings.arrow.import_nested = settings.input_format_arrow_import_nested;
+    format_settings.orc.import_nested = settings.input_format_orc_import_nested;
 
     /// Validate avro_schema_registry_url with RemoteHostFilter when non-empty and in Server context
     if (format_settings.schema.is_server)

src/Formats/FormatSettings.h

@@ -53,6 +53,7 @@ struct FormatSettings
     {
         UInt64 row_group_size = 1000000;
         bool low_cardinality_as_dictionary = false;
+        bool import_nested = false;
     } arrow;
 
     struct
@@ -100,6 +101,7 @@ struct FormatSettings
     struct
     {
         UInt64 row_group_size = 1000000;
+        bool import_nested = false;
     } parquet;
 
     struct Pretty
@@ -174,6 +176,11 @@ struct FormatSettings
         bool deduce_templates_of_expressions = true;
         bool accurate_types_of_literals = true;
     } values;
+
+    struct
+    {
+        bool import_nested = false;
+    } orc;
 };
 
 }

src/Processors/Formats/Impl/ArrowBlockInputFormat.cpp

@@ -22,8 +22,8 @@ namespace ErrorCodes
     extern const int CANNOT_READ_ALL_DATA;
 }
 
-ArrowBlockInputFormat::ArrowBlockInputFormat(ReadBuffer & in_, const Block & header_, bool stream_)
-    : IInputFormat(header_, in_), stream{stream_}
+ArrowBlockInputFormat::ArrowBlockInputFormat(ReadBuffer & in_, const Block & header_, bool stream_, const FormatSettings & format_settings_)
+    : IInputFormat(header_, in_), stream{stream_}, format_settings(format_settings_)
 {
 }
 
@@ -102,7 +102,7 @@ void ArrowBlockInputFormat::prepareReader()
         schema = file_reader->schema();
     }
 
-    arrow_column_to_ch_column = std::make_unique<ArrowColumnToCHColumn>(getPort().getHeader(), std::move(schema), "Arrow");
+    arrow_column_to_ch_column = std::make_unique<ArrowColumnToCHColumn>(getPort().getHeader(), "Arrow", format_settings.arrow.import_nested);
 
     if (stream)
         record_batch_total = -1;
@@ -119,19 +119,19 @@ void registerInputFormatProcessorArrow(FormatFactory & factory)
         [](ReadBuffer & buf,
            const Block & sample,
            const RowInputFormatParams & /* params */,
-           const FormatSettings & /* format_settings */)
+           const FormatSettings & format_settings)
         {
-            return std::make_shared<ArrowBlockInputFormat>(buf, sample, false);
+            return std::make_shared<ArrowBlockInputFormat>(buf, sample, false, format_settings);
         });
     factory.markFormatAsColumnOriented("Arrow");
     factory.registerInputFormatProcessor(
         "ArrowStream",
         [](ReadBuffer & buf,
            const Block & sample,
            const RowInputFormatParams & /* params */,
-           const FormatSettings & /* format_settings */)
+           const FormatSettings & format_settings)
         {
-            return std::make_shared<ArrowBlockInputFormat>(buf, sample, true);
+            return std::make_shared<ArrowBlockInputFormat>(buf, sample, true, format_settings);
         });
 }
 

src/Processors/Formats/Impl/ArrowBlockInputFormat.h

@@ -6,6 +6,7 @@
 #if USE_ARROW
 
 #include <Processors/Formats/IInputFormat.h>
+#include <Formats/FormatSettings.h>
 
 namespace arrow { class RecordBatchReader; }
 namespace arrow::ipc { class RecordBatchFileReader; }
@@ -19,7 +20,7 @@ class ArrowColumnToCHColumn;
 class ArrowBlockInputFormat : public IInputFormat
 {
 public:
-    ArrowBlockInputFormat(ReadBuffer & in_, const Block & header_, bool stream_);
+    ArrowBlockInputFormat(ReadBuffer & in_, const Block & header_, bool stream_, const FormatSettings & format_settings_);
 
     void resetParser() override;
 
@@ -41,6 +42,8 @@ class ArrowBlockInputFormat : public IInputFormat
     int record_batch_total = 0;
     int record_batch_current = 0;
 
+    const FormatSettings format_settings;
+
     void prepareReader();
 };