static ColumnPtr truncateNestedDataIfNull()

in cpp-ch/local-engine/Storages/Output/NormalFileWriter.cpp [43:154]


static ColumnPtr truncateNestedDataIfNull(const ColumnPtr & column)
{
    if (const auto * col_const = checkAndGetColumn<ColumnConst>(column.get()))
    {
        size_t s = col_const->size();
        auto new_data = truncateNestedDataIfNull(col_const->getDataColumnPtr());
        return ColumnConst::create(std::move(new_data), s);
    }
    else if (const auto * col_array = checkAndGetColumn<ColumnArray>(column.get()))
    {
        auto new_data = truncateNestedDataIfNull(col_array->getDataPtr());
        return ColumnArray::create(std::move(new_data), col_array->getOffsetsPtr());
    }
    else if (const auto * col_map = checkAndGetColumn<ColumnMap>(column.get()))
    {
        auto new_nested = truncateNestedDataIfNull(col_map->getNestedColumnPtr());
        return ColumnMap::create(std::move(new_nested));
    }
    else if (const auto * col_tuple = checkAndGetColumn<ColumnTuple>(column.get()))
    {
        Columns new_columns;
        for (size_t i = 0; i < col_tuple->tupleSize(); ++i)
            new_columns.emplace_back(truncateNestedDataIfNull(col_tuple->getColumnPtr(i)));
        return ColumnTuple::create(std::move(new_columns));
    }
    else if (const auto * col_nullable = checkAndGetColumn<ColumnNullable>(column.get()))
    {
        const auto & null_map = col_nullable->getNullMapData();
        auto nested = truncateNestedDataIfNull(col_nullable->getNestedColumnPtr());
        const auto * nested_array = checkAndGetColumn<ColumnArray>(nested.get());
        const auto * nested_map = checkAndGetColumn<ColumnMap>(nested.get());
        const auto * nested_tuple = checkAndGetColumn<ColumnTuple>(nested.get());

        if (!memoryIsZero(null_map.data(), 0, null_map.size()) && (nested_array || nested_map || nested_tuple))
        {
            /// Process Nullable(Array) or Nullable(Map)
            if (nested_array || nested_map)
            {
                if (!nested_array)
                    nested_array = checkAndGetColumn<ColumnArray>(&nested_map->getNestedColumn());

                const auto & offsets = nested_array->getOffsets();
                size_t total_data_size = 0;
                for (size_t i = 0; i < null_map.size(); ++i)
                    total_data_size += (offsets[i] - offsets[i - 1]) * (!null_map[i]);

                auto new_nested_array = nested_array->cloneEmpty();
                new_nested_array->reserve(nested_array->size());
                auto & new_nested_array_data = assert_cast<ColumnArray &>(*new_nested_array).getData();
                new_nested_array_data.reserve(total_data_size);

                for (size_t i = 0; i < null_map.size(); ++i)
                    if (null_map[i])
                        new_nested_array->insertDefault();
                    else
                        new_nested_array->insertFrom(*nested_array, i);

                if (nested_map)
                {
                    auto new_nested_map = ColumnMap::create(std::move(new_nested_array));
                    return ColumnNullable::create(std::move(new_nested_map), col_nullable->getNullMapColumnPtr());
                }
                else
                {
                    return ColumnNullable::create(std::move(new_nested_array), col_nullable->getNullMapColumnPtr());
                }
            }
            else
            {
                /// Process Nullable(Tuple)
                const auto & nested_columns = nested_tuple->getColumns();
                Columns new_nested_columns(nested_columns.size());
                for (size_t i = 0; i < nested_columns.size(); ++i)
                {
                    const auto & nested_column = nested_columns[i];
                    TypeIndex type_index = nested_column->getDataType();
                    if (const auto * nullable_nested_column = checkAndGetColumn<ColumnNullable>(nested_column.get()))
                        type_index = nullable_nested_column->getNestedColumnPtr()->getDataType();

                    bool should_truncate = type_index == TypeIndex::Array || type_index == TypeIndex::Map || type_index == TypeIndex::Tuple;
                    if (should_truncate)
                    {
                        auto new_nested_column = nested_column->cloneEmpty();
                        new_nested_column->reserve(nested_column->size());
                        for (size_t j = 0; j < null_map.size(); ++j)
                        {
                            if (null_map[j])
                                new_nested_column->insertDefault();
                            else
                                new_nested_column->insertFrom(*nested_column, j);
                        }
                        new_nested_columns[i] = std::move(new_nested_column);
                    }
                    else
                    {
                        new_nested_columns[i] = nested_column;
                    }
                }

                auto new_nested_tuple = ColumnTuple::create(std::move(new_nested_columns));
                return ColumnNullable::create(std::move(new_nested_tuple), col_nullable->getNullMapColumnPtr());
            }
        }
        else
        {
            auto new_nested = truncateNestedDataIfNull(nested);
            return ColumnNullable::create(std::move(new_nested), col_nullable->getNullMapColumnPtr());
        }
    }
    else
        return column;
}