support Bool

JuliaHEP · Oct 17, 2024 · 7e20fae · 7e20fae
1 parent 0719ae5
commit 7e20fae
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 4 deletions.
diff --git a/src/RNTuple/Writing/constants.jl b/src/RNTuple/Writing/constants.jl
@@ -1,6 +1,7 @@
 const RNTUPLE_WRITE_TYPE_IDX_DICT = Dict(
     Index64 => (0x01, sizeof(Index64) * 8),
     Index32 => (0x02, sizeof(Index32) * 8),
+    Bool => (0x06, 1),
     Float64 => (0x10, sizeof(UInt64) * 8),
     Float32 => (0x11, sizeof(UInt32) * 8),
     Float16 => (0x12, sizeof(UInt16) * 8),
@@ -14,6 +15,7 @@ const RNTUPLE_WRITE_TYPE_IDX_DICT = Dict(
 )
 
 const RNTUPLE_WRITE_TYPE_CPPNAME_DICT = Dict(
+    Bool => "bool",
     Float16 => "std::float16_t",
     Float32 => "float",
     Float64 => "double",

diff --git a/src/RNTuple/Writing/page_writing.jl b/src/RNTuple/Writing/page_writing.jl
@@ -24,6 +24,11 @@ Turns an AbstractVector into a page of an RNTuple. The element type must be prim
 """
 function rnt_ary_to_page(ary::AbstractVector, cr::ColumnRecord) end
 
+function rnt_ary_to_page(ary::AbstractVector{Bool}, cr::ColumnRecord)
+    chunks = BitVector(ary).chunks
+    Page_write(reinterpret(UInt8, chunks))
+end
+
 function rnt_ary_to_page(ary::AbstractVector{Float64}, cr::ColumnRecord)
     (;split, zigzag, delta) = _detect_encoding(cr.type)
     if split

diff --git a/src/RNTuple/fieldcolumn_reading.jl b/src/RNTuple/fieldcolumn_reading.jl
@@ -105,7 +105,8 @@ function read_field(io, field::LeafField{Bool}, page_list)
 
     # pad to nearest 8*k bytes because each chunk needs to be UInt64
     bytes = read_pagedesc(io, pages, cr)
-    append!(bytes, zeros(eltype(bytes), (64 - rem(total_num_elements, 64))÷8))
+    N_pad = 8 - mod1(length(bytes), 8)
+    append!(bytes, zeros(eltype(bytes), N_pad))
     chunks = reinterpret(UInt64, bytes)
 
     res = BitVector(undef, total_num_elements)

diff --git a/test/RNTupleWriting/lowlevel.jl b/test/RNTupleWriting/lowlevel.jl
@@ -236,8 +236,10 @@ end
     @test field_ids == [0,1,2,3,4,5]
 end
 
+const RNT_primitive_Ts = [Bool, Float64, Float32, Float16, Int64, Int32, Int16, Int8, UInt64, UInt32, UInt16]
+
 @testset "RNTuple Writing - Single colunm round trips" begin
-for _ = 1:10, T in [Float64, Float32, Float16, Int64, Int32, Int16, Int8, UInt64, UInt32, UInt16]
+for _ = 1:10, T in RNT_primitive_Ts
     newtable = Dict(randstring(rand(2:10)) => rand(T, rand(1:100)))
     newio = IOBuffer()
     UnROOT.write_rntuple(newio, newtable)
@@ -260,7 +262,7 @@ end
 end
 
 @testset "RNTuple Writing - Multiple colunm round trips" begin
-    Ts = rand([Float64, Float32, Float16, Int64, Int32, Int16, Int8, UInt64, UInt32, UInt16], 15)
+    Ts = rand(RNT_primitive_Ts, 15)
     Nitems = rand(10:1000)
     newtable = Dict(randstring(rand(2:10)) => rand(T, Nitems) for T in Ts)
     newio = IOBuffer()
@@ -284,7 +286,7 @@ end
 end
 
 @testset "RNTuple Writing - Vector colunms" begin
-    Ts = rand([Float64, Float32, Float16, Int64, Int32, Int16, Int8, UInt64, UInt32, UInt16], 15)
+    Ts = rand(RNT_primitive_Ts, 15)
     inner_Nitems = [3,4,0,0,1,2]
     newtable = Dict(randstring(rand(2:10)) => [rand(T, Nitems) for Nitems in inner_Nitems] for T in Ts)
     newio = IOBuffer()