Skip to content

Commit d17b403

Browse files
wangzhaobo957-cloudwangzhaobo957-cloud
committed
[test](lance) Cover additional type edge cases #66496 
1 parent 1139417 commit d17b403

1 file changed

Lines changed: 139 additions & 0 deletions

File tree

‎be/test/format_v2/table/lance_reader_test.cpp‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -887,6 +887,24 @@ std::shared_ptr<arrow::Array> make_bfloat16_vector_array() {
887887
return array;
888888
}
889889

890+
// Creates a scalar Lance BFloat16 array from raw 16-bit values.
891+
std::shared_ptr<arrow::FixedSizeBinaryArray> make_bfloat16_array(
892+
const std::vector<std::optional<uint16_t>>& values) {
893+
arrow::FixedSizeBinaryBuilder builder(arrow::fixed_size_binary(2));
894+
for (const auto& value : values) {
895+
if (!value.has_value()) {
896+
EXPECT_TRUE(builder.AppendNull().ok());
897+
continue;
898+
}
899+
std::array<uint8_t, sizeof(uint16_t)> bytes {};
900+
LittleEndian::Store16(bytes.data(), *value);
901+
EXPECT_TRUE(builder.Append(bytes.data()).ok());
902+
}
903+
std::shared_ptr<arrow::FixedSizeBinaryArray> array;
904+
EXPECT_TRUE(builder.Finish(&array).ok());
905+
return array;
906+
}
907+
890908
std::pair<size_t, size_t> array_range(const ColumnArray& array, size_t row) {
891909
const auto& offsets = array.get_offsets();
892910
return {row == 0 ? 0 : static_cast<size_t>(offsets[row - 1]),
@@ -1115,6 +1133,127 @@ TEST(LanceTableReaderTypeTest, ReadsAdditionalArrowAndLanceTypes) {
11151133
EXPECT_TRUE(reader.close().ok());
11161134
}
11171135

1136+
// Verifies BFloat16 special values retain their exact Float32 bit patterns.
1137+
TEST(LanceTableReaderTypeTest, ConvertsBFloat16SpecialValuesExactly) {
1138+
const auto metadata =
1139+
arrow::KeyValueMetadata::Make({"ARROW:extension:name"}, {"lance.bfloat16"});
1140+
const auto field = arrow::field("value", arrow::fixed_size_binary(2))->WithMetadata(metadata);
1141+
const std::vector<std::optional<uint16_t>> input_bits {0x0000, 0x8000, 0x3F80, 0xBF80,
1142+
0x7F80, 0xFF80, 0x7FC1, std::nullopt};
1143+
const std::array<uint32_t, 7> expected_bits {0x00000000, 0x80000000, 0x3F800000, 0xBF800000,
1144+
0x7F800000, 0xFF800000, 0x7FC10000};
1145+
1146+
std::shared_ptr<arrow::Array> normalized;
1147+
const auto status = normalize_lance_arrow_array_for_test(field, make_bfloat16_array(input_bits),
1148+
&normalized);
1149+
ASSERT_TRUE(status.ok()) << status.to_string();
1150+
const auto floats = std::dynamic_pointer_cast<arrow::FloatArray>(normalized);
1151+
ASSERT_NE(nullptr, floats);
1152+
ASSERT_EQ(input_bits.size(), static_cast<size_t>(floats->length()));
1153+
for (size_t index = 0; index < expected_bits.size(); ++index) {
1154+
ASSERT_FALSE(floats->IsNull(index));
1155+
EXPECT_EQ(expected_bits[index], std::bit_cast<uint32_t>(floats->Value(index)));
1156+
}
1157+
EXPECT_TRUE(floats->IsNull(expected_bits.size()));
1158+
}
1159+
1160+
// Verifies nested BFloat16 fields are converted without rebuilding unaffected sibling data.
1161+
TEST(LanceTableReaderTypeTest, ConvertsBFloat16NestedInStruct) {
1162+
const auto metadata =
1163+
arrow::KeyValueMetadata::Make({"ARROW:extension:name"}, {"lance.bfloat16"});
1164+
const auto value_field =
1165+
arrow::field("value", arrow::fixed_size_binary(2))->WithMetadata(metadata);
1166+
const auto id_field = arrow::field("id", arrow::int32());
1167+
const auto bfloat16_values = make_bfloat16_array(
1168+
{std::optional<uint16_t> {0x3F80}, std::optional<uint16_t> {0xC020}});
1169+
arrow::Int32Builder id_builder;
1170+
ASSERT_TRUE(id_builder.AppendValues({7, 8}).ok());
1171+
std::shared_ptr<arrow::Int32Array> ids;
1172+
ASSERT_TRUE(id_builder.Finish(&ids).ok());
1173+
auto struct_result = arrow::StructArray::Make({bfloat16_values, ids}, {value_field, id_field});
1174+
ASSERT_TRUE(struct_result.ok()) << struct_result.status().ToString();
1175+
const auto input = std::move(struct_result).ValueUnsafe();
1176+
const auto field = arrow::field("record", input->type());
1177+
1178+
std::shared_ptr<arrow::Array> normalized;
1179+
const auto status = normalize_lance_arrow_array_for_test(field, input, &normalized);
1180+
ASSERT_TRUE(status.ok()) << status.to_string();
1181+
ASSERT_NE(input.get(), normalized.get());
1182+
const auto normalized_struct = std::dynamic_pointer_cast<arrow::StructArray>(normalized);
1183+
ASSERT_NE(nullptr, normalized_struct);
1184+
const auto floats = std::dynamic_pointer_cast<arrow::FloatArray>(normalized_struct->field(0));
1185+
ASSERT_NE(nullptr, floats);
1186+
EXPECT_FLOAT_EQ(1.0F, floats->Value(0));
1187+
EXPECT_FLOAT_EQ(-2.5F, floats->Value(1));
1188+
EXPECT_EQ(ids->data().get(), normalized_struct->field(1)->data().get());
1189+
}
1190+
1191+
// Verifies the Lance JSON extension reads LargeBinary values and nulls through JSON SerDe.
1192+
TEST(LanceTableReaderTypeTest, ReadsLanceJsonLargeBinaryValues) {
1193+
const auto metadata = arrow::KeyValueMetadata::Make({"ARROW:extension:name"}, {"lance.json"});
1194+
const auto schema = arrow::schema(
1195+
{arrow::field("json_value", arrow::large_binary())->WithMetadata(metadata)});
1196+
arrow::LargeBinaryBuilder builder;
1197+
ASSERT_TRUE(builder.Append(R"({"format":"lance","value":42})").ok());
1198+
ASSERT_TRUE(builder.AppendNull().ok());
1199+
std::shared_ptr<arrow::LargeBinaryArray> values;
1200+
ASSERT_TRUE(builder.Finish(&values).ok());
1201+
const auto record_batch = arrow::RecordBatch::Make(schema, 2, {values});
1202+
1203+
const Columns columns {projected_column("json_value", TYPE_JSONB, true)};
1204+
TQueryGlobals query_globals;
1205+
RuntimeState state(query_globals);
1206+
RuntimeProfile profile("lance_json_large_binary");
1207+
TFileScanRangeParams scan_params;
1208+
LanceTableReader reader;
1209+
ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok());
1210+
1211+
Block block;
1212+
add_output_columns(&block, columns);
1213+
size_t rows = 0;
1214+
ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok());
1215+
ASSERT_EQ(2, rows);
1216+
const auto& json_values = assert_cast<const ColumnNullable&>(*block.get_by_position(0).column);
1217+
EXPECT_EQ((ColumnUInt8::Container {0, 1}), json_values.get_null_map_data());
1218+
EXPECT_EQ(R"({"format":"lance","value":42})", columns[0].type->to_string(json_values, 0));
1219+
EXPECT_TRUE(reader.close().ok());
1220+
}
1221+
1222+
// Verifies Duration values preserve signed 64-bit boundaries and nullability.
1223+
TEST(LanceTableReaderTypeTest, ReadsDurationBoundaryValues) {
1224+
const auto duration_type = arrow::duration(arrow::TimeUnit::NANO);
1225+
const auto schema = arrow::schema({arrow::field("duration", duration_type)});
1226+
arrow::DurationBuilder builder(duration_type, arrow::default_memory_pool());
1227+
const std::array<int64_t, 4> expected {std::numeric_limits<int64_t>::min(), -1, 0,
1228+
std::numeric_limits<int64_t>::max()};
1229+
ASSERT_TRUE(builder.AppendValues(expected.data(), expected.size()).ok());
1230+
ASSERT_TRUE(builder.AppendNull().ok());
1231+
std::shared_ptr<arrow::DurationArray> values;
1232+
ASSERT_TRUE(builder.Finish(&values).ok());
1233+
const auto record_batch = arrow::RecordBatch::Make(schema, 5, {values});
1234+
1235+
const Columns columns {projected_column("duration", TYPE_BIGINT, true)};
1236+
TQueryGlobals query_globals;
1237+
RuntimeState state(query_globals);
1238+
RuntimeProfile profile("lance_duration_boundaries");
1239+
TFileScanRangeParams scan_params;
1240+
LanceTableReader reader;
1241+
ASSERT_TRUE(init_reader(&reader, columns, &state, &profile, &scan_params).ok());
1242+
1243+
Block block;
1244+
add_output_columns(&block, columns);
1245+
size_t rows = 0;
1246+
ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok());
1247+
ASSERT_EQ(5, rows);
1248+
const auto& durations = assert_cast<const ColumnNullable&>(*block.get_by_position(0).column);
1249+
const auto& duration_values = assert_cast<const ColumnInt64&>(durations.get_nested_column());
1250+
EXPECT_EQ((ColumnUInt8::Container {0, 0, 0, 0, 1}), durations.get_null_map_data());
1251+
for (size_t index = 0; index < expected.size(); ++index) {
1252+
EXPECT_EQ(expected[index], duration_values.get_data()[index]);
1253+
}
1254+
EXPECT_TRUE(reader.close().ok());
1255+
}
1256+
11181257
// Verifies ordinary nested columns bypass Arrow reconstruction when no BFloat16 exists.
11191258
TEST(LanceTableReaderTypeTest, KeepsOrdinaryNestedArrayUnchangedDuringNormalization) {
11201259
const auto nested_type = arrow::struct_({

0 commit comments

Comments
 (0)