| /* |
| * Licensed to the Apache Software Foundation (ASF) under one or more |
| * contributor license agreements. See the NOTICE file distributed with |
| * this work for additional information regarding copyright ownership. |
| * The ASF licenses this file to You 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. |
| */ |
| #include "CHColumnToSparkRow.h" |
| #include <Columns/ColumnArray.h> |
| #include <Columns/ColumnConst.h> |
| #include <Columns/ColumnNullable.h> |
| #include <Columns/ColumnString.h> |
| #include <Columns/IColumn.h> |
| #include <DataTypes/DataTypeArray.h> |
| #include <DataTypes/DataTypeLowCardinality.h> |
| #include <DataTypes/DataTypeMap.h> |
| #include <DataTypes/DataTypeNullable.h> |
| #include <DataTypes/DataTypeTuple.h> |
| #include <DataTypes/DataTypesDecimal.h> |
| #include <DataTypes/ObjectUtils.h> |
| #include <jni/jni_common.h> |
| #include <Common/Exception.h> |
| |
| namespace DB |
| { |
| namespace ErrorCodes |
| { |
| extern const int LOGICAL_ERROR; |
| extern const int UNKNOWN_TYPE; |
| } |
| } |
| |
| |
| namespace local_engine |
| { |
| using namespace DB; |
| |
| int64_t calculateBitSetWidthInBytes(int64_t num_fields) |
| { |
| return ((num_fields + 63) / 64) * 8; |
| } |
| |
| static int64_t calculatedFixeSizePerRow(int64_t num_cols) |
| { |
| return calculateBitSetWidthInBytes(num_cols) + num_cols * 8; |
| } |
| |
| int64_t roundNumberOfBytesToNearestWord(int64_t num_bytes) |
| { |
| auto remainder = num_bytes & 0x07; // This is equivalent to `numBytes % 8` |
| return num_bytes + ((8 - remainder) & 0x7); |
| } |
| |
| |
| void bitSet(char * bitmap, size_t index) |
| { |
| int64_t mask = 1L << (index & 0x3f); // mod 64 and shift |
| int64_t word_offset = (index >> 6) * 8; |
| int64_t word; |
| memcpy(&word, bitmap + word_offset, sizeof(int64_t)); |
| int64_t value = word | mask; |
| memcpy(bitmap + word_offset, &value, sizeof(int64_t)); |
| } |
| |
| ALWAYS_INLINE bool isBitSet(const char * bitmap, size_t index) |
| { |
| assert(index >= 0); |
| int64_t mask = 1L << (index & 63); |
| int64_t word_offset = static_cast<int64_t>(index >> 6) * 8L; |
| int64_t word = *reinterpret_cast<const int64_t *>(bitmap + word_offset); |
| return word & mask; |
| } |
| |
| static void writeFixedLengthNonNullableValue( |
| char * buffer_address, |
| int64_t field_offset, |
| const ColumnWithTypeAndName & col, |
| size_t num_rows, |
| const std::vector<int64_t> & offsets, |
| const MaskVector & masks = nullptr) |
| { |
| FixedLengthDataWriter writer(col.type); |
| |
| if (writer.getWhichDataType().isDecimal32()) |
| { |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| auto field = (*col.column)[row_idx]; |
| writer.write(field, buffer_address + offsets[i] + field_offset); |
| } |
| } |
| else |
| { |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| writer.unsafeWrite(col.column->getDataAt(row_idx), buffer_address + offsets[i] + field_offset); |
| } |
| } |
| } |
| |
| static void writeFixedLengthNullableValue( |
| char * buffer_address, |
| int64_t field_offset, |
| const ColumnWithTypeAndName & col, |
| int32_t col_index, |
| size_t num_rows, |
| const std::vector<int64_t> & offsets, |
| const MaskVector & masks = nullptr) |
| { |
| const auto & nullable_column = checkAndGetColumn<ColumnNullable>(*col.column); |
| const auto & null_map = nullable_column.getNullMapData(); |
| const auto & nested_column = nullable_column.getNestedColumn(); |
| FixedLengthDataWriter writer(col.type); |
| |
| if (writer.getWhichDataType().isDecimal32()) |
| { |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| if (null_map[row_idx]) |
| bitSet(buffer_address + offsets[i], col_index); |
| else |
| { |
| auto field = (*col.column)[row_idx]; |
| writer.write(field, buffer_address + offsets[i] + field_offset); |
| } |
| } |
| } |
| else |
| { |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| if (null_map[row_idx]) |
| bitSet(buffer_address + offsets[i], col_index); |
| else |
| writer.unsafeWrite(nested_column.getDataAt(row_idx), buffer_address + offsets[i] + field_offset); |
| } |
| } |
| } |
| |
| static void writeVariableLengthNonNullableValue( |
| char * buffer_address, |
| int64_t field_offset, |
| const ColumnWithTypeAndName & col, |
| size_t num_rows, |
| const std::vector<int64_t> & offsets, |
| std::vector<int64_t> & buffer_cursor, |
| const MaskVector & masks = nullptr) |
| { |
| const auto type_without_nullable{removeNullable(col.type)}; |
| const bool use_raw_data = BackingDataLengthCalculator::isDataTypeSupportRawData(type_without_nullable); |
| const bool big_endian = BackingDataLengthCalculator::isBigEndianInSparkRow(type_without_nullable); |
| VariableLengthDataWriter writer(col.type, buffer_address, offsets, buffer_cursor); |
| if (use_raw_data) |
| { |
| if (!big_endian) |
| { |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| StringRef str = col.column->getDataAt(row_idx); |
| int64_t offset_and_size = writer.writeUnalignedBytes(i, str.data, str.size, 0); |
| memcpy(buffer_address + offsets[i] + field_offset, &offset_and_size, 8); |
| } |
| } |
| else |
| { |
| Field field; |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| StringRef str_view = col.column->getDataAt(row_idx); |
| String buf(str_view.data, str_view.size); |
| BackingDataLengthCalculator::swapDecimalEndianBytes(buf); |
| int64_t offset_and_size = writer.writeUnalignedBytes(i, buf.data(), buf.size(), 0); |
| memcpy(buffer_address + offsets[i] + field_offset, &offset_and_size, 8); |
| } |
| } |
| } |
| else |
| { |
| Field field; |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| field = (*col.column)[row_idx]; |
| int64_t offset_and_size = writer.write(i, field, 0); |
| memcpy(buffer_address + offsets[i] + field_offset, &offset_and_size, 8); |
| } |
| } |
| } |
| |
| static void writeVariableLengthNullableValue( |
| char * buffer_address, |
| int64_t field_offset, |
| const ColumnWithTypeAndName & col, |
| int32_t col_index, |
| size_t num_rows, |
| const std::vector<int64_t> & offsets, |
| std::vector<int64_t> & buffer_cursor, |
| const MaskVector & masks = nullptr) |
| { |
| const auto & nullable_column = checkAndGetColumn<ColumnNullable>(*col.column); |
| const auto & null_map = nullable_column.getNullMapData(); |
| const auto & nested_column = nullable_column.getNestedColumn(); |
| const auto type_without_nullable{removeNullable(col.type)}; |
| const bool use_raw_data = BackingDataLengthCalculator::isDataTypeSupportRawData(type_without_nullable); |
| const bool big_endian = BackingDataLengthCalculator::isBigEndianInSparkRow(type_without_nullable); |
| VariableLengthDataWriter writer(col.type, buffer_address, offsets, buffer_cursor); |
| if (use_raw_data) |
| { |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| if (null_map[row_idx]) |
| bitSet(buffer_address + offsets[i], col_index); |
| else if (!big_endian) |
| { |
| StringRef str = nested_column.getDataAt(row_idx); |
| int64_t offset_and_size = writer.writeUnalignedBytes(i, str.data, str.size, 0); |
| memcpy(buffer_address + offsets[i] + field_offset, &offset_and_size, 8); |
| } |
| else |
| { |
| Field field; |
| nested_column.get(row_idx, field); |
| StringRef str_view = nested_column.getDataAt(row_idx); |
| String buf(str_view.data, str_view.size); |
| BackingDataLengthCalculator::swapDecimalEndianBytes(buf); |
| int64_t offset_and_size = writer.writeUnalignedBytes(i, buf.data(), buf.size(), 0); |
| memcpy(buffer_address + offsets[i] + field_offset, &offset_and_size, 8); |
| } |
| } |
| } |
| else |
| { |
| Field field; |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| if (null_map[row_idx]) |
| bitSet(buffer_address + offsets[i], col_index); |
| else |
| { |
| field = nested_column[row_idx]; |
| int64_t offset_and_size = writer.write(i, field, 0); |
| memcpy(buffer_address + offsets[i] + field_offset, &offset_and_size, 8); |
| } |
| } |
| } |
| } |
| |
| |
| static void writeValue( |
| char * buffer_address, |
| int64_t field_offset, |
| const ColumnWithTypeAndName & col, |
| int32_t col_index, |
| int64_t num_rows, |
| const std::vector<int64_t> & offsets, |
| std::vector<int64_t> & buffer_cursor, |
| const MaskVector & masks = nullptr) |
| { |
| const auto type_without_nullable{removeNullable(col.type)}; |
| const auto is_nullable = isColumnNullable(*col.column); |
| if (BackingDataLengthCalculator::isFixedLengthDataType(type_without_nullable)) |
| { |
| if (is_nullable) |
| writeFixedLengthNullableValue(buffer_address, field_offset, col, col_index, num_rows, offsets, masks); |
| else |
| writeFixedLengthNonNullableValue(buffer_address, field_offset, col, num_rows, offsets, masks); |
| } |
| else if (BackingDataLengthCalculator::isVariableLengthDataType(type_without_nullable)) |
| { |
| if (is_nullable) |
| writeVariableLengthNullableValue(buffer_address, field_offset, col, col_index, num_rows, offsets, buffer_cursor, masks); |
| else |
| writeVariableLengthNonNullableValue(buffer_address, field_offset, col, num_rows, offsets, buffer_cursor, masks); |
| } |
| else |
| throw Exception(ErrorCodes::UNKNOWN_TYPE, "Doesn't support type {} for writeValue", col.type->getName()); |
| } |
| |
| SparkRowInfo::SparkRowInfo( |
| const DB::ColumnsWithTypeAndName & cols, const DB::DataTypes & dataTypes, size_t col_size, size_t row_size, const MaskVector & masks) |
| : types(dataTypes) |
| , num_rows(masks == nullptr ? row_size : masks->size()) |
| , num_cols(col_size) |
| , null_bitset_width_in_bytes(calculateBitSetWidthInBytes(num_cols)) |
| , total_bytes(0) |
| , offsets(num_rows, 0) |
| , lengths(num_rows, 0) |
| , buffer_cursor(num_rows, 0) |
| , buffer_address(nullptr) |
| { |
| int64_t fixed_size_per_row = calculatedFixeSizePerRow(num_cols); |
| |
| /// Initialize lengths and buffer_cursor |
| for (size_t i = 0; i < num_rows; i++) |
| { |
| lengths[i] = fixed_size_per_row; |
| buffer_cursor[i] = fixed_size_per_row; |
| } |
| |
| for (int64_t col_idx = 0; col_idx < num_cols; ++col_idx) |
| { |
| const auto & col = cols[col_idx]; |
| /// No need to calculate backing data length for fixed length types |
| const auto type_without_nullable = removeLowCardinalityAndNullable(col.type); |
| if (BackingDataLengthCalculator::isVariableLengthDataType(type_without_nullable)) |
| { |
| if (BackingDataLengthCalculator::isDataTypeSupportRawData(type_without_nullable)) |
| { |
| auto column = col.column->convertToFullIfNeeded(); |
| if (const auto * nullable_column = checkAndGetColumn<ColumnNullable>(&*column)) |
| { |
| const auto & nested_column = nullable_column->getNestedColumn(); |
| const auto & null_map = nullable_column->getNullMapData(); |
| for (size_t i = 0; i < num_rows; ++i) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| if (!null_map[row_idx]) |
| lengths[i] += roundNumberOfBytesToNearestWord(nested_column.getDataAt(row_idx).size); |
| } |
| } |
| else |
| { |
| for (size_t i = 0; i < num_rows; ++i) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| lengths[i] += roundNumberOfBytesToNearestWord(column->getDataAt(row_idx).size); |
| } |
| } |
| } |
| else |
| { |
| BackingDataLengthCalculator calculator(type_without_nullable); |
| for (size_t i = 0; i < num_rows; ++i) |
| { |
| size_t row_idx = masks == nullptr ? i : masks->at(i); |
| const auto field = (*col.column)[row_idx]; |
| lengths[i] += calculator.calculate(field); |
| } |
| } |
| } |
| } |
| |
| /// Initialize offsets |
| for (size_t i = 1; i < num_rows; ++i) |
| offsets[i] = offsets[i - 1] + lengths[i - 1]; |
| |
| /// Initialize total_bytes |
| for (size_t i = 0; i < num_rows; ++i) |
| total_bytes += lengths[i]; |
| } |
| |
| SparkRowInfo::SparkRowInfo(const Block & block, const MaskVector & masks) |
| : SparkRowInfo(block.getColumnsWithTypeAndName(), block.getDataTypes(), block.columns(), block.rows(), masks) |
| { |
| } |
| |
| const DB::DataTypes & SparkRowInfo::getDataTypes() const |
| { |
| return types; |
| } |
| |
| int64_t SparkRowInfo::getFieldOffset(int32_t col_idx) const |
| { |
| return null_bitset_width_in_bytes + 8L * col_idx; |
| } |
| |
| int64_t SparkRowInfo::getNullBitsetWidthInBytes() const |
| { |
| return null_bitset_width_in_bytes; |
| } |
| |
| void SparkRowInfo::setNullBitsetWidthInBytes(int64_t null_bitset_width_in_bytes_) |
| { |
| null_bitset_width_in_bytes = null_bitset_width_in_bytes_; |
| } |
| |
| int64_t SparkRowInfo::getNumCols() const |
| { |
| return num_cols; |
| } |
| |
| void SparkRowInfo::setNumCols(int64_t num_cols_) |
| { |
| num_cols = num_cols_; |
| } |
| |
| int64_t SparkRowInfo::getNumRows() const |
| { |
| return num_rows; |
| } |
| |
| void SparkRowInfo::setNumRows(int64_t num_rows_) |
| { |
| num_rows = num_rows_; |
| } |
| |
| char * SparkRowInfo::getBufferAddress() const |
| { |
| return buffer_address; |
| } |
| |
| void SparkRowInfo::setBufferAddress(char * buffer_address_) |
| { |
| buffer_address = buffer_address_; |
| } |
| |
| const std::vector<int64_t> & SparkRowInfo::getOffsets() const |
| { |
| return offsets; |
| } |
| |
| const std::vector<int64_t> & SparkRowInfo::getLengths() const |
| { |
| return lengths; |
| } |
| |
| std::vector<int64_t> & SparkRowInfo::getBufferCursor() |
| { |
| return buffer_cursor; |
| } |
| |
| int64_t SparkRowInfo::getTotalBytes() const |
| { |
| return total_bytes; |
| } |
| |
| std::unique_ptr<SparkRowInfo> CHColumnToSparkRow::convertCHColumnToSparkRow(const Block & block, const MaskVector & masks) |
| { |
| if (!block.columns()) |
| throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "A block with empty columns"); |
| std::unique_ptr<SparkRowInfo> spark_row_info = std::make_unique<SparkRowInfo>(block, masks); |
| spark_row_info->setBufferAddress(static_cast<char *>(alloc(spark_row_info->getTotalBytes(), 64))); |
| // spark_row_info->setBufferAddress(alignedAlloc(spark_row_info->getTotalBytes(), 64)); |
| memset(spark_row_info->getBufferAddress(), 0, spark_row_info->getTotalBytes()); |
| for (auto col_idx = 0; col_idx < spark_row_info->getNumCols(); col_idx++) |
| { |
| const auto & col = block.getByPosition(col_idx); |
| int64_t field_offset = spark_row_info->getFieldOffset(col_idx); |
| |
| ColumnWithTypeAndName col_full{col.column->convertToFullIfNeeded(), removeLowCardinality(col.type), col.name}; |
| writeValue( |
| spark_row_info->getBufferAddress(), |
| field_offset, |
| col_full, |
| col_idx, |
| spark_row_info->getNumRows(), |
| spark_row_info->getOffsets(), |
| spark_row_info->getBufferCursor(), |
| masks); |
| } |
| return spark_row_info; |
| } |
| |
| void CHColumnToSparkRow::freeMem(char * address, size_t size) |
| { |
| free(address, size); |
| // rollback(size); |
| } |
| |
| BackingDataLengthCalculator::BackingDataLengthCalculator(const DataTypePtr & type_) |
| : type_without_nullable(removeNullable(type_)), which(type_without_nullable) |
| { |
| if (!isFixedLengthDataType(type_without_nullable) && !isVariableLengthDataType(type_without_nullable)) |
| throw Exception( |
| ErrorCodes::UNKNOWN_TYPE, "Doesn't support type {} for BackingDataLengthCalculator", type_without_nullable->getName()); |
| } |
| |
| int64_t BackingDataLengthCalculator::calculate(const Field & field) const |
| { |
| if (field.isNull()) |
| return 0; |
| |
| if (which.isNativeInt() || which.isNativeUInt() || which.isFloat() || which.isDateOrDate32() || which.isDateTime64() |
| || which.isDecimal32() || which.isDecimal64()) |
| return 0; |
| |
| if (which.isStringOrFixedString()) |
| { |
| const auto & str = field.safeGet<String>(); |
| return roundNumberOfBytesToNearestWord(str.size()); |
| } |
| |
| if (which.isDecimal128()) |
| return 16; |
| |
| if (which.isArray()) |
| { |
| /// 内存布局：numElements(8B) | null_bitmap(与numElements成正比) | values(每个值长度与类型有关) | backing buffer |
| const auto & array = field.safeGet<Array>(); /// Array can not be wrapped with Nullable |
| const auto num_elems = array.size(); |
| int64_t res = 8 + calculateBitSetWidthInBytes(num_elems); |
| |
| const auto * array_type = typeid_cast<const DataTypeArray *>(type_without_nullable.get()); |
| const auto & nested_type = array_type->getNestedType(); |
| res += roundNumberOfBytesToNearestWord(getArrayElementSize(nested_type) * num_elems); |
| |
| BackingDataLengthCalculator calculator(nested_type); |
| for (size_t i = 0; i < array.size(); ++i) |
| res += calculator.calculate(array[i]); |
| return res; |
| } |
| |
| if (which.isMap()) |
| { |
| /// 内存布局：Length of UnsafeArrayData of key(8B) | UnsafeArrayData of key | UnsafeArrayData of value |
| int64_t res = 8; |
| |
| /// Construct Array of keys and values from Map |
| const auto & map = field.safeGet<Map>(); /// Map can not be wrapped with Nullable |
| const auto num_keys = map.size(); |
| auto array_key = Array(); |
| auto array_val = Array(); |
| array_key.reserve(num_keys); |
| array_val.reserve(num_keys); |
| for (size_t i = 0; i < num_keys; ++i) |
| { |
| const auto & pair = map[i].safeGet<DB::Tuple>(); |
| array_key.push_back(pair[0]); |
| array_val.push_back(pair[1]); |
| } |
| |
| const auto * map_type = typeid_cast<const DB::DataTypeMap *>(type_without_nullable.get()); |
| |
| const auto & key_type = map_type->getKeyType(); |
| const auto key_array_type = std::make_shared<DataTypeArray>(key_type); |
| BackingDataLengthCalculator calculator_key(key_array_type); |
| res += calculator_key.calculate(array_key); |
| |
| const auto & val_type = map_type->getValueType(); |
| const auto type_array_val = std::make_shared<DataTypeArray>(val_type); |
| BackingDataLengthCalculator calculator_val(type_array_val); |
| res += calculator_val.calculate(array_val); |
| return res; |
| } |
| |
| if (which.isTuple()) |
| { |
| /// 内存布局：null_bitmap(字节数与字段数成正比) | field1 value(8B) | field2 value(8B) | ... | fieldn value(8B) | backing buffer |
| const auto & tuple = field.safeGet<Tuple>(); /// Tuple can not be wrapped with Nullable |
| const auto * type_tuple = typeid_cast<const DataTypeTuple *>(type_without_nullable.get()); |
| const auto & type_fields = type_tuple->getElements(); |
| const auto num_fields = type_fields.size(); |
| int64_t res = calculateBitSetWidthInBytes(num_fields) + 8 * num_fields; |
| for (size_t i = 0; i < num_fields; ++i) |
| { |
| BackingDataLengthCalculator calculator(type_fields[i]); |
| res += calculator.calculate(tuple[i]); |
| } |
| return res; |
| } |
| |
| throw Exception( |
| ErrorCodes::UNKNOWN_TYPE, "Doesn't support type {} for BackingBufferLengthCalculator", type_without_nullable->getName()); |
| } |
| |
| int64_t BackingDataLengthCalculator::getArrayElementSize(const DataTypePtr & nested_type) |
| { |
| const WhichDataType nested_which(removeNullable(nested_type)); |
| if (nested_which.isUInt8() || nested_which.isInt8()) |
| return 1; |
| else if (nested_which.isUInt16() || nested_which.isInt16() || nested_which.isDate()) |
| return 2; |
| else if (nested_which.isUInt32() || nested_which.isInt32() || nested_which.isFloat32() || nested_which.isDate32()) |
| return 4; |
| else if ( |
| nested_which.isUInt64() || nested_which.isInt64() || nested_which.isFloat64() || nested_which.isDateTime64() |
| || nested_which.isDecimal32() || nested_which.isDecimal64()) |
| return 8; |
| else |
| return 8; |
| } |
| |
| bool BackingDataLengthCalculator::isFixedLengthDataType(const DataTypePtr & type_without_nullable) |
| { |
| const WhichDataType which(type_without_nullable); |
| return which.isUInt8() || which.isInt8() || which.isUInt16() || which.isInt16() || which.isDate() || which.isUInt32() || which.isInt32() |
| || which.isFloat32() || which.isDate32() || which.isDecimal32() || which.isUInt64() || which.isInt64() || which.isFloat64() |
| || which.isDateTime64() || which.isDecimal64() || which.isNothing(); |
| } |
| |
| bool BackingDataLengthCalculator::isVariableLengthDataType(const DataTypePtr & type_without_nullable) |
| { |
| const WhichDataType which(type_without_nullable); |
| return which.isStringOrFixedString() || which.isDecimal128() || which.isArray() || which.isMap() || which.isTuple(); |
| } |
| |
| bool BackingDataLengthCalculator::isDataTypeSupportRawData(const DB::DataTypePtr & type_without_nullable) |
| { |
| const WhichDataType which(type_without_nullable); |
| return isFixedLengthDataType(type_without_nullable) || which.isStringOrFixedString() || which.isDecimal128(); |
| } |
| |
| bool BackingDataLengthCalculator::isBigEndianInSparkRow(const DB::DataTypePtr & type_without_nullable) |
| { |
| const WhichDataType which(type_without_nullable); |
| return which.isDecimal128(); |
| } |
| |
| void BackingDataLengthCalculator::swapDecimalEndianBytes(String & buf) |
| { |
| assert(buf.size() == 16); |
| |
| using base_type = Decimal128::NativeType::base_type; |
| auto * decimal128 = reinterpret_cast<Decimal128 *>(buf.data()); |
| for (size_t i = 0; i != std::size(decimal128->value.items); ++i) |
| decimal128->value.items[i] = __builtin_bswap64(decimal128->value.items[i]); |
| |
| base_type * high = reinterpret_cast<base_type *>(buf.data() + 8); |
| base_type * low = reinterpret_cast<base_type *>(buf.data()); |
| std::swap(*high, *low); |
| } |
| |
| VariableLengthDataWriter::VariableLengthDataWriter( |
| const DataTypePtr & type_, char * buffer_address_, const std::vector<int64_t> & offsets_, std::vector<int64_t> & buffer_cursor_) |
| : type_without_nullable(removeNullable(type_)) |
| , which(type_without_nullable) |
| , buffer_address(buffer_address_) |
| , offsets(offsets_) |
| , buffer_cursor(buffer_cursor_) |
| { |
| assert(buffer_address); |
| assert(offsets.size() == buffer_cursor.size()); |
| |
| if (!BackingDataLengthCalculator::isVariableLengthDataType(type_without_nullable)) |
| throw Exception(ErrorCodes::UNKNOWN_TYPE, "VariableLengthDataWriter doesn't support type {}", type_without_nullable->getName()); |
| } |
| |
| int64_t VariableLengthDataWriter::writeArray(size_t row_idx, const DB::Array & array, int64_t parent_offset) |
| { |
| /// 内存布局：numElements(8B) | null_bitmap(与numElements成正比) | values(每个值长度与类型有关) | backing data |
| const auto & offset = offsets[row_idx]; |
| auto & cursor = buffer_cursor[row_idx]; |
| const auto num_elems = array.size(); |
| const auto * array_type = typeid_cast<const DataTypeArray *>(type_without_nullable.get()); |
| const auto & nested_type = array_type->getNestedType(); |
| |
| /// Write numElements(8B) |
| const auto start = cursor; |
| memcpy(buffer_address + offset + cursor, &num_elems, 8); |
| cursor += 8; |
| if (num_elems == 0) |
| return BackingDataLengthCalculator::getOffsetAndSize(start - parent_offset, 8); |
| |
| /// Skip null_bitmap(already reset to zero) |
| const auto len_null_bitmap = calculateBitSetWidthInBytes(num_elems); |
| cursor += len_null_bitmap; |
| |
| /// Skip values(already reset to zero) |
| const auto elem_size = BackingDataLengthCalculator::getArrayElementSize(nested_type); |
| const auto len_values = roundNumberOfBytesToNearestWord(elem_size * num_elems); |
| cursor += len_values; |
| |
| if (BackingDataLengthCalculator::isFixedLengthDataType(removeNullable(nested_type))) |
| { |
| /// If nested type is fixed-length data type, update null_bitmap and values in place |
| FixedLengthDataWriter writer(nested_type); |
| for (size_t i = 0; i < num_elems; ++i) |
| { |
| const auto & elem = array[i]; |
| if (elem.isNull()) |
| bitSet(buffer_address + offset + start + 8, i); |
| else |
| writer.write(elem, buffer_address + offset + start + 8 + len_null_bitmap + i * elem_size); |
| } |
| } |
| else |
| { |
| /// If nested type is not fixed-length data type, update null_bitmap in place |
| /// And append values in backing data recursively |
| VariableLengthDataWriter writer(nested_type, buffer_address, offsets, buffer_cursor); |
| for (size_t i = 0; i < num_elems; ++i) |
| { |
| const auto & elem = array[i]; |
| if (elem.isNull()) |
| bitSet(buffer_address + offset + start + 8, i); |
| else |
| { |
| const auto offset_and_size = writer.write(row_idx, elem, start); |
| memcpy(buffer_address + offset + start + 8 + len_null_bitmap + i * elem_size, &offset_and_size, 8); |
| } |
| } |
| } |
| return BackingDataLengthCalculator::getOffsetAndSize(start - parent_offset, cursor - start); |
| } |
| |
| int64_t VariableLengthDataWriter::writeMap(size_t row_idx, const DB::Map & map, int64_t parent_offset) |
| { |
| /// 内存布局：Length of UnsafeArrayData of key(8B) | UnsafeArrayData of key | UnsafeArrayData of value |
| const auto & offset = offsets[row_idx]; |
| auto & cursor = buffer_cursor[row_idx]; |
| |
| /// Skip length of UnsafeArrayData of key(8B) |
| const auto start = cursor; |
| cursor += 8; |
| |
| /// Even if Map is empty, still write as [unsafe key array numBytes] [unsafe key array] [unsafe value array] |
| const auto num_pairs = map.size(); |
| |
| /// Construct array of keys and array of values from map |
| auto key_array = Array(); |
| auto val_array = Array(); |
| key_array.reserve(num_pairs); |
| val_array.reserve(num_pairs); |
| for (size_t i = 0; i < num_pairs; ++i) |
| { |
| const auto & pair = map[i].safeGet<DB::Tuple>(); |
| key_array.push_back(pair[0]); |
| val_array.push_back(pair[1]); |
| } |
| |
| const auto * map_type = typeid_cast<const DB::DataTypeMap *>(type_without_nullable.get()); |
| |
| /// Append UnsafeArrayData of key |
| const auto & key_type = map_type->getKeyType(); |
| const auto key_array_type = std::make_shared<DataTypeArray>(key_type); |
| VariableLengthDataWriter key_writer(key_array_type, buffer_address, offsets, buffer_cursor); |
| const auto key_array_size = BackingDataLengthCalculator::extractSize(key_writer.write(row_idx, key_array, start + 8)); |
| |
| /// Fill length of UnsafeArrayData of key |
| memcpy(buffer_address + offset + start, &key_array_size, 8); |
| |
| /// Append UnsafeArrayData of value |
| const auto & val_type = map_type->getValueType(); |
| const auto val_array_type = std::make_shared<DataTypeArray>(val_type); |
| VariableLengthDataWriter val_writer(val_array_type, buffer_address, offsets, buffer_cursor); |
| val_writer.write(row_idx, val_array, start + 8 + key_array_size); |
| return BackingDataLengthCalculator::getOffsetAndSize(start - parent_offset, cursor - start); |
| } |
| |
| int64_t VariableLengthDataWriter::writeStruct(size_t row_idx, const DB::Tuple & tuple, int64_t parent_offset) |
| { |
| /// 内存布局：null_bitmap(字节数与字段数成正比) | values(num_fields * 8B) | backing data |
| const auto & offset = offsets[row_idx]; |
| auto & cursor = buffer_cursor[row_idx]; |
| const auto start = cursor; |
| |
| /// Skip null_bitmap |
| const auto * tuple_type = typeid_cast<const DataTypeTuple *>(type_without_nullable.get()); |
| const auto & field_types = tuple_type->getElements(); |
| const auto num_fields = field_types.size(); |
| if (num_fields == 0) |
| return BackingDataLengthCalculator::getOffsetAndSize(start - parent_offset, 0); |
| const auto len_null_bitmap = calculateBitSetWidthInBytes(num_fields); |
| cursor += len_null_bitmap; |
| |
| /// Skip values |
| cursor += num_fields * 8; |
| |
| /// If field type is fixed-length, fill field value in values region |
| /// else append it to backing data region, and update offset_and_size in values region |
| for (size_t i = 0; i < num_fields; ++i) |
| { |
| const auto & field_value = tuple[i]; |
| const auto & field_type = field_types[i]; |
| if (field_value.isNull()) |
| { |
| bitSet(buffer_address + offset + start, i); |
| continue; |
| } |
| |
| if (BackingDataLengthCalculator::isFixedLengthDataType(removeNullable(field_type))) |
| { |
| FixedLengthDataWriter writer(field_type); |
| writer.write(field_value, buffer_address + offset + start + len_null_bitmap + i * 8); |
| } |
| else |
| { |
| VariableLengthDataWriter writer(field_type, buffer_address, offsets, buffer_cursor); |
| const auto offset_and_size = writer.write(row_idx, field_value, start); |
| memcpy(buffer_address + offset + start + len_null_bitmap + 8 * i, &offset_and_size, 8); |
| } |
| } |
| return BackingDataLengthCalculator::getOffsetAndSize(start - parent_offset, cursor - start); |
| } |
| |
| int64_t VariableLengthDataWriter::write(size_t row_idx, const DB::Field & field, int64_t parent_offset) |
| { |
| assert(row_idx < offsets.size()); |
| |
| if (field.isNull()) |
| return 0; |
| |
| if (which.isStringOrFixedString()) |
| { |
| const auto & str = field.safeGet<String>(); |
| return writeUnalignedBytes(row_idx, str.data(), str.size(), parent_offset); |
| } |
| |
| if (which.isDecimal128()) |
| { |
| const auto & decimal_field = field.safeGet<DecimalField<Decimal128>>(); |
| auto decimal128 = decimal_field.getValue(); |
| String buf(reinterpret_cast<const char *>(&decimal128), sizeof(decimal128)); |
| BackingDataLengthCalculator::swapDecimalEndianBytes(buf); |
| return writeUnalignedBytes(row_idx, buf.c_str(), sizeof(Decimal128), parent_offset); |
| } |
| |
| if (which.isArray()) |
| { |
| const auto & array = field.safeGet<Array>(); |
| return writeArray(row_idx, array, parent_offset); |
| } |
| |
| if (which.isMap()) |
| { |
| const auto & map = field.safeGet<Map>(); |
| return writeMap(row_idx, map, parent_offset); |
| } |
| |
| if (which.isTuple()) |
| { |
| const auto & tuple = field.safeGet<Tuple>(); |
| return writeStruct(row_idx, tuple, parent_offset); |
| } |
| |
| throw Exception(ErrorCodes::UNKNOWN_TYPE, "Doesn't support type {} for BackingDataWriter", type_without_nullable->getName()); |
| } |
| |
| int64_t BackingDataLengthCalculator::getOffsetAndSize(int64_t cursor, int64_t size) |
| { |
| return (cursor << 32) | size; |
| } |
| |
| int64_t BackingDataLengthCalculator::extractOffset(int64_t offset_and_size) |
| { |
| return offset_and_size >> 32; |
| } |
| |
| int64_t BackingDataLengthCalculator::extractSize(int64_t offset_and_size) |
| { |
| return offset_and_size & 0xffffffff; |
| } |
| |
| int64_t VariableLengthDataWriter::writeUnalignedBytes(size_t row_idx, const char * src, size_t size, int64_t parent_offset) |
| { |
| memcpy(buffer_address + offsets[row_idx] + buffer_cursor[row_idx], src, size); |
| auto res = BackingDataLengthCalculator::getOffsetAndSize(buffer_cursor[row_idx] - parent_offset, size); |
| buffer_cursor[row_idx] += roundNumberOfBytesToNearestWord(size); |
| return res; |
| } |
| |
| |
| FixedLengthDataWriter::FixedLengthDataWriter(const DB::DataTypePtr & type_) |
| : type_without_nullable(removeNullable(type_)), which(type_without_nullable) |
| { |
| if (!BackingDataLengthCalculator::isFixedLengthDataType(type_without_nullable)) |
| throw Exception(ErrorCodes::UNKNOWN_TYPE, "FixedLengthWriter doesn't support type {}", type_without_nullable->getName()); |
| } |
| |
| void FixedLengthDataWriter::write(const DB::Field & field, char * buffer) |
| { |
| /// Skip null value |
| if (field.isNull()) |
| return; |
| |
| if (which.isUInt8()) |
| { |
| const auto value = static_cast<UInt8>(field.safeGet<UInt8>()); |
| memcpy(buffer, &value, 1); |
| } |
| else if (which.isUInt16() || which.isDate()) |
| { |
| const auto value = static_cast<UInt16>(field.safeGet<UInt16>()); |
| memcpy(buffer, &value, 2); |
| } |
| else if (which.isUInt32() || which.isDate32()) |
| { |
| const auto value = static_cast<UInt32>(field.safeGet<UInt32>()); |
| memcpy(buffer, &value, 4); |
| } |
| else if (which.isUInt64()) |
| { |
| const auto & value = field.safeGet<UInt64>(); |
| memcpy(buffer, &value, 8); |
| } |
| else if (which.isInt8()) |
| { |
| const auto value = static_cast<Int8>(field.safeGet<Int8>()); |
| memcpy(buffer, &value, 1); |
| } |
| else if (which.isInt16()) |
| { |
| const auto value = static_cast<Int16>(field.safeGet<Int16>()); |
| memcpy(buffer, &value, 2); |
| } |
| else if (which.isInt32()) |
| { |
| const auto value = static_cast<Int32>(field.safeGet<Int32>()); |
| memcpy(buffer, &value, 4); |
| } |
| else if (which.isInt64()) |
| { |
| const auto & value = field.safeGet<Int64>(); |
| memcpy(buffer, &value, 8); |
| } |
| else if (which.isFloat32()) |
| { |
| const auto value = static_cast<Float32>(field.safeGet<Float32>()); |
| memcpy(buffer, &value, 4); |
| } |
| else if (which.isFloat64()) |
| { |
| const auto & value = field.safeGet<Float64>(); |
| memcpy(buffer, &value, 8); |
| } |
| else if (which.isDecimal32()) |
| { |
| const auto & value = field.safeGet<Decimal32>(); |
| const Int64 decimal = static_cast<Int64>(value.getValue()); |
| memcpy(buffer, &decimal, 8); |
| } |
| else if (which.isDecimal64() || which.isDateTime64()) |
| { |
| const auto & value = field.safeGet<Decimal64>(); |
| const auto decimal = value.getValue(); |
| memcpy(buffer, &decimal, 8); |
| } |
| else |
| throw Exception(ErrorCodes::UNKNOWN_TYPE, "FixedLengthDataWriter doesn't support type {}", type_without_nullable->getName()); |
| } |
| |
| void FixedLengthDataWriter::unsafeWrite(const StringRef & str, char * buffer) |
| { |
| memcpy(buffer, str.data, str.size); |
| } |
| |
| void FixedLengthDataWriter::unsafeWrite(const char * __restrict src, char * __restrict buffer) |
| { |
| memcpy(buffer, src, type_without_nullable->getSizeOfValueInMemory()); |
| } |
| |
| namespace SparkRowInfoJNI |
| { |
| static jclass spark_row_info_class; |
| static jmethodID spark_row_info_constructor; |
| void init(JNIEnv * env) |
| { |
| spark_row_info_class = CreateGlobalClassReference(env, "Lorg/apache/gluten/row/SparkRowInfo;"); |
| spark_row_info_constructor = GetMethodID(env, spark_row_info_class, "<init>", "([J[JJJJ)V"); |
| } |
| void destroy(JNIEnv * env) |
| { |
| env->DeleteGlobalRef(spark_row_info_class); |
| } |
| jobject create(JNIEnv * env, const SparkRowInfo & spark_row_info) |
| { |
| auto * offsets_arr = env->NewLongArray(spark_row_info.getNumRows()); |
| const auto * offsets_src = spark_row_info.getOffsets().data(); |
| env->SetLongArrayRegion(offsets_arr, 0, spark_row_info.getNumRows(), reinterpret_cast<const jlong *>(offsets_src)); |
| auto * lengths_arr = env->NewLongArray(spark_row_info.getNumRows()); |
| const auto * lengths_src = spark_row_info.getLengths().data(); |
| env->SetLongArrayRegion(lengths_arr, 0, spark_row_info.getNumRows(), reinterpret_cast<const jlong *>(lengths_src)); |
| int64_t address = reinterpret_cast<int64_t>(spark_row_info.getBufferAddress()); |
| int64_t column_number = spark_row_info.getNumCols(); |
| int64_t total_size = spark_row_info.getTotalBytes(); |
| |
| return env->NewObject( |
| spark_row_info_class, spark_row_info_constructor, offsets_arr, lengths_arr, address, column_number, total_size); |
| } |
| } |
| |
| } |
