Skip to content

Commit fea57e4

Browse files
kszucskszucs
committed
style: format code
1 parent de5bd14 commit fea57e4

File tree

11 files changed

+198
-64
lines changed

11 files changed

+198
-64
lines changed

‎marrow/arrays.mojo‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,11 @@ struct PrimitiveArray[T: DataType](Movable, Sized, Writable):
253253
"""Upload array data to the GPU."""
254254
var bm: Optional[Bitmap] = None
255255
if self.bitmap:
256-
bm = Bitmap(self.bitmap.value()._buffer.to_device(ctx), 0, self.bitmap.value()._length)
256+
bm = Bitmap(
257+
self.bitmap.value()._buffer.to_device(ctx),
258+
0,
259+
self.bitmap.value()._length,
260+
)
257261
return PrimitiveArray[Self.T](
258262
length=self.length,
259263
nulls=self.nulls,
@@ -266,7 +270,11 @@ struct PrimitiveArray[T: DataType](Movable, Sized, Writable):
266270
"""Download array data from the GPU to owned CPU heap buffers."""
267271
var bm: Optional[Bitmap] = None
268272
if self.bitmap:
269-
bm = Bitmap(self.bitmap.value()._buffer.to_cpu(ctx), 0, self.bitmap.value()._length)
273+
bm = Bitmap(
274+
self.bitmap.value()._buffer.to_cpu(ctx),
275+
0,
276+
self.bitmap.value()._length,
277+
)
270278
return PrimitiveArray[Self.T](
271279
length=self.length,
272280
nulls=self.nulls,

‎marrow/bitmap.mojo‎

Lines changed: 40 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,9 @@ struct Bitmap(ImplicitlyCopyable, Movable, Sized, Writable):
9393
fn is_valid(self, index: Int) -> Bool:
9494
"""Return True if bit at (_offset + index) is set (value is valid)."""
9595
var bit_index = self.bit_offset() + index
96-
return Bool((self._buffer.ptr[bit_index >> 3] >> UInt8(bit_index & 7)) & 1)
96+
return Bool(
97+
(self._buffer.ptr[bit_index >> 3] >> UInt8(bit_index & 7)) & 1
98+
)
9799

98100
@always_inline
99101
fn is_null(self, index: Int) -> Bool:
@@ -125,9 +127,17 @@ struct Bitmap(ImplicitlyCopyable, Movable, Sized, Writable):
125127
var acc0 = SIMD[DType.uint8, width](0)
126128
var acc1 = SIMD[DType.uint8, width](0)
127129
comptime for j in range(t1_iters):
128-
acc0 += pop_count((ptr + i + (j * 2) * width).load[width=width]())
129-
acc1 += pop_count((ptr + i + (j * 2 + 1) * width).load[width=width]())
130-
count += Int((acc0.cast[DType.uint16]() + acc1.cast[DType.uint16]()).reduce_add())
130+
acc0 += pop_count(
131+
(ptr + i + (j * 2) * width).load[width=width]()
132+
)
133+
acc1 += pop_count(
134+
(ptr + i + (j * 2 + 1) * width).load[width=width]()
135+
)
136+
count += Int(
137+
(
138+
acc0.cast[DType.uint16]() + acc1.cast[DType.uint16]()
139+
).reduce_add()
140+
)
131141

132142
# Tier 2: 64-byte blocks for the remainder.
133143
# NEON: 4 iters, max 32/lane ≤ 255 ✓.
@@ -144,13 +154,17 @@ struct Bitmap(ImplicitlyCopyable, Movable, Sized, Writable):
144154
for i in range(lead_bytes):
145155
count -= Int(pop_count(ptr[i]))
146156
if lead_sub_byte:
147-
count -= Int(pop_count(ptr[lead_bytes] & UInt8((1 << lead_sub_byte) - 1)))
157+
count -= Int(
158+
pop_count(ptr[lead_bytes] & UInt8((1 << lead_sub_byte) - 1))
159+
)
148160
if trail_bits:
149161
var trail_bytes = trail_bits >> 3
150162
var trail_sub_byte = trail_bits & 7
151163
var first_trail = total_bytes - trail_bytes
152164
if trail_sub_byte:
153-
count -= Int(pop_count(ptr[first_trail - 1] >> UInt8(trail_sub_byte)))
165+
count -= Int(
166+
pop_count(ptr[first_trail - 1] >> UInt8(trail_sub_byte))
167+
)
154168
for i in range(first_trail, total_bytes):
155169
count -= Int(pop_count(ptr[i]))
156170

@@ -160,7 +174,6 @@ struct Bitmap(ImplicitlyCopyable, Movable, Sized, Writable):
160174
"""Return a zero-copy view of `length` bits starting at `offset`."""
161175
return Bitmap(self._buffer, self.bit_offset() + offset, length)
162176

163-
164177
# --- Bulk SIMD operations ---
165178
#
166179
# Arrow buffers are 64-byte aligned and zero-padded to multiples of 64 B.
@@ -212,7 +225,9 @@ struct Bitmap(ImplicitlyCopyable, Movable, Sized, Writable):
212225

213226
@always_inline
214227
fn _binop[
215-
op: fn[W: Int](SIMD[DType.uint8, W], SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]
228+
op: fn[W: Int](SIMD[DType.uint8, W], SIMD[DType.uint8, W]) -> SIMD[
229+
DType.uint8, W
230+
]
216231
](self, other: Bitmap) raises -> Bitmap:
217232
"""Apply a byte-level binary operation across two bitmaps.
218233
@@ -260,30 +275,37 @@ struct Bitmap(ImplicitlyCopyable, Movable, Sized, Writable):
260275

261276
return Bitmap(builder.finish(), lead_bits_a, self._length)
262277

263-
264278
fn write_to[W: Writer](self, mut writer: W):
265279
writer.write(
266280
"Bitmap(offset=", self.bit_offset(), ", length=", self._length, ")"
267281
)
268282

269283

270284
@always_inline
271-
fn _and[W: Int](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
285+
fn _and[
286+
W: Int
287+
](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
272288
return a & b
273289

274290

275291
@always_inline
276-
fn _or[W: Int](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
292+
fn _or[
293+
W: Int
294+
](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
277295
return a | b
278296

279297

280298
@always_inline
281-
fn _xor[W: Int](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
299+
fn _xor[
300+
W: Int
301+
](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
282302
return a ^ b
283303

284304

285305
@always_inline
286-
fn _and_not[W: Int](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
306+
fn _and_not[
307+
W: Int
308+
](a: SIMD[DType.uint8, W], b: SIMD[DType.uint8, W]) -> SIMD[DType.uint8, W]:
287309
return a & ~b
288310

289311

@@ -317,7 +339,8 @@ struct BitmapBuilder(Movable):
317339

318340
@always_inline
319341
fn unsafe_ptr(self) -> UnsafePointer[UInt8, MutExternalOrigin]:
320-
"""Return the raw mutable byte pointer (for low-level bit operations)."""
342+
"""Return the raw mutable byte pointer (for low-level bit operations).
343+
"""
321344
return self._builder.ptr
322345

323346
@always_inline
@@ -344,7 +367,9 @@ struct BitmapBuilder(Movable):
344367

345368
if start_byte == end_byte:
346369
# All bits in one byte
347-
var mask = UInt8((1 << end_bit) - 1) & (UInt8(0xFF) << UInt8(start_bit))
370+
var mask = UInt8((1 << end_bit) - 1) & (
371+
UInt8(0xFF) << UInt8(start_bit)
372+
)
348373
if value:
349374
ptr[start_byte] = ptr[start_byte] | mask
350375
else:

‎marrow/buffers.mojo‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,7 @@ struct BufferBuilder(Movable):
485485
# Buffer — immutable buffer for read-only array data
486486
# ---------------------------------------------------------------------------
487487

488+
488489
# TODO: add assertions to ensure alignment and padding invariants hold
489490
struct Buffer(ImplicitlyCopyable, Movable, Writable):
490491
"""Immutable contiguous memory region.
@@ -720,5 +721,3 @@ struct Buffer(ImplicitlyCopyable, Movable, Writable):
720721
fn write_to[W: Writer](self, mut writer: W):
721722
"""Write the buffer's bytes to a Writer."""
722723
writer.write("Buffer(ptr={}, size={})".format(self.ptr, self.size))
723-
724-

‎marrow/kernels/__init__.mojo‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,9 @@ from marrow.dtypes import DataType, all_numeric_dtypes
6060
# ---------------------------------------------------------------------------
6161

6262

63-
fn bitmap_and(a: Optional[Bitmap], b: Optional[Bitmap]) raises -> Optional[Bitmap]:
63+
fn bitmap_and(
64+
a: Optional[Bitmap], b: Optional[Bitmap]
65+
) raises -> Optional[Bitmap]:
6466
"""Compute the output validity bitmap as the bitwise AND of two input bitmaps.
6567
6668
Output bit i is True iff both a[i] and b[i] are True (valid).
@@ -355,7 +357,9 @@ fn reduce_simd[
355357
comptime width = simd_byte_width() // size_of[native]()
356358
var length = len(array)
357359
var has_nulls = array.bitmap.__bool__()
358-
var bp = array.bitmap.value()._buffer.unsafe_ptr() if has_nulls else UnsafePointer[UInt8, ImmutExternalOrigin]()
360+
var bp = array.bitmap.value()._buffer.unsafe_ptr() if has_nulls else UnsafePointer[
361+
UInt8, ImmutExternalOrigin
362+
]()
359363

360364
var acc = SIMD[native, width](initial)
361365
var identity_vec = SIMD[native, width](initial)

‎marrow/kernels/arithmetic.mojo‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,9 @@ fn add[
9393
Null if either input is null at that position.
9494
"""
9595
if ctx:
96-
return binary_gpu[T, func=_add[T.native, _], name="add"](left, right, ctx.value())
96+
return binary_gpu[T, func=_add[T.native, _], name="add"](
97+
left, right, ctx.value()
98+
)
9799
else:
98100
return binary_simd[T, func=_add[T.native, _], name="add"](left, right)
99101

@@ -127,7 +129,9 @@ fn sub[
127129
Null if either input is null at that position.
128130
"""
129131
if ctx:
130-
return binary_gpu[T, func=_sub[T.native, _], name="sub"](left, right, ctx.value())
132+
return binary_gpu[T, func=_sub[T.native, _], name="sub"](
133+
left, right, ctx.value()
134+
)
131135
else:
132136
return binary_simd[T, func=_sub[T.native, _], name="sub"](left, right)
133137

@@ -161,7 +165,9 @@ fn mul[
161165
Null if either input is null at that position.
162166
"""
163167
if ctx:
164-
return binary_gpu[T, func=_mul[T.native, _], name="mul"](left, right, ctx.value())
168+
return binary_gpu[T, func=_mul[T.native, _], name="mul"](
169+
left, right, ctx.value()
170+
)
165171
else:
166172
return binary_simd[T, func=_mul[T.native, _], name="mul"](left, right)
167173

@@ -195,7 +201,9 @@ fn div[
195201
Null if either input is null at that position.
196202
"""
197203
if ctx:
198-
return binary_gpu[T, func=_div[T.native, _], name="div"](left, right, ctx.value())
204+
return binary_gpu[T, func=_div[T.native, _], name="div"](
205+
left, right, ctx.value()
206+
)
199207
else:
200208
return binary_simd[T, func=_div[T.native, _], name="div"](left, right)
201209

@@ -233,7 +241,9 @@ fn floordiv[
233241
left, right, ctx.value()
234242
)
235243
else:
236-
return binary_simd[T, func=_floordiv[T.native, _], name="floordiv"](left, right)
244+
return binary_simd[T, func=_floordiv[T.native, _], name="floordiv"](
245+
left, right
246+
)
237247

238248

239249
fn floordiv(left: Array, right: Array) raises -> Array:
@@ -265,7 +275,9 @@ fn mod[
265275
Null if either input is null at that position.
266276
"""
267277
if ctx:
268-
return binary_gpu[T, func=_mod[T.native, _], name="mod"](left, right, ctx.value())
278+
return binary_gpu[T, func=_mod[T.native, _], name="mod"](
279+
left, right, ctx.value()
280+
)
269281
else:
270282
return binary_simd[T, func=_mod[T.native, _], name="mod"](left, right)
271283

@@ -299,7 +311,9 @@ fn min_[
299311
Null if either input is null at that position.
300312
"""
301313
if ctx:
302-
return binary_gpu[T, func=_min[T.native, _], name="min_"](left, right, ctx.value())
314+
return binary_gpu[T, func=_min[T.native, _], name="min_"](
315+
left, right, ctx.value()
316+
)
303317
else:
304318
return binary_simd[T, func=_min[T.native, _], name="min_"](left, right)
305319

@@ -333,7 +347,9 @@ fn max_[
333347
Null if either input is null at that position.
334348
"""
335349
if ctx:
336-
return binary_gpu[T, func=_max[T.native, _], name="max_"](left, right, ctx.value())
350+
return binary_gpu[T, func=_max[T.native, _], name="max_"](
351+
left, right, ctx.value()
352+
)
337353
else:
338354
return binary_simd[T, func=_max[T.native, _], name="max_"](left, right)
339355

‎marrow/kernels/filter.mojo‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,11 @@ fn filter[
168168
# All selected: single bulk copy of data and validity.
169169
var bm_bytes = math.ceildiv(n, 8)
170170
if array.bitmap:
171-
memcpy(dest=bm.unsafe_ptr(), src=array.bitmap.value()._buffer.unsafe_ptr(), count=bm_bytes)
171+
memcpy(
172+
dest=bm.unsafe_ptr(),
173+
src=array.bitmap.value()._buffer.unsafe_ptr(),
174+
count=bm_bytes,
175+
)
172176
else:
173177
bm.set_range(0, n, True)
174178
comptime if native == DType.bool:
@@ -189,7 +193,9 @@ fn filter[
189193
# bulk-copy each run. Threshold: out_len / n > 0.8 ↔ out_len * 10 > n * 8
190194
if out_len * 10 > n * 8:
191195
var has_nulls = array.nulls > 0
192-
var src_bm_ptr = array.bitmap.value()._buffer.unsafe_ptr() if has_nulls else UnsafePointer[UInt8, ImmutExternalOrigin]()
196+
var src_bm_ptr = array.bitmap.value()._buffer.unsafe_ptr() if has_nulls else UnsafePointer[
197+
UInt8, ImmutExternalOrigin
198+
]()
193199
for word_i in range(word_count):
194200
var word = word_ptr[word_i]
195201
if word == 0:
@@ -233,7 +239,9 @@ fn filter[
233239
var si = run_start + i
234240
if Bool((src_bm_ptr[si >> 3] >> UInt8(si & 7)) & 1):
235241
var di = out_idx + i
236-
bm.unsafe_ptr()[di >> 3] |= UInt8(1) << UInt8(di & 7)
242+
bm.unsafe_ptr()[di >> 3] |= UInt8(1) << UInt8(
243+
di & 7
244+
)
237245
else:
238246
bm.set_range(out_idx, run_len, True)
239247
out_idx += run_len

‎marrow/kernels/sum.mojo‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,9 @@ fn any_(array: PrimitiveArray[bool_dt]) -> Bool:
187187
"""True if any valid element is True. False if empty or all null."""
188188
return Bool(
189189
reduce_simd[
190-
bool_dt, combine=_or[DType.bool, _], horizontal=_horizontal_or[DType.bool, _]
190+
bool_dt,
191+
combine=_or[DType.bool, _],
192+
horizontal=_horizontal_or[DType.bool, _],
191193
](array, Scalar[DType.bool](False))
192194
)
193195

@@ -196,6 +198,8 @@ fn all_(array: PrimitiveArray[bool_dt]) -> Bool:
196198
"""True if all valid elements are True. True if empty or all null."""
197199
return Bool(
198200
reduce_simd[
199-
bool_dt, combine=_and[DType.bool, _], horizontal=_horizontal_and[DType.bool, _]
201+
bool_dt,
202+
combine=_and[DType.bool, _],
203+
horizontal=_horizontal_and[DType.bool, _],
200204
](array, Scalar[DType.bool](True))
201205
)

‎marrow/tests/bench_bitmap.mojo‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ fn bench_count_set_bits_aligned(mut b: Bencher, size: Int) raises:
8888

8989
@parameter
9090
fn bench_count_set_bits_unaligned(mut b: Bencher, size: Int) raises:
91-
"""count_set_bits with byte_offset=96 (NOT 64-byte aligned, lead_bytes=32)."""
91+
"""count_set_bits with byte_offset=96 (NOT 64-byte aligned, lead_bytes=32).
92+
"""
9293
var bm = _make_alternating(size + 2048).slice(96 << 3, size)
9394

9495
@always_inline
@@ -171,7 +172,8 @@ fn bench_set_range(mut b: Bencher, size: Int) raises:
171172

172173
@parameter
173174
fn bench_invert_cache_aligned(mut b: Bencher, size: Int) raises:
174-
"""Invert with a byte_offset that is already 64-byte aligned (lead_bytes=0)."""
175+
"""Invert with a byte_offset that is already 64-byte aligned (lead_bytes=0).
176+
"""
175177
# 128-byte = 1024-bit offset → byte_offset=128, 128 & 63 == 0, no lead bytes.
176178
var bm = _make_alternating(size + 2048).slice(128 << 3, size)
177179

@@ -225,7 +227,8 @@ fn bench_and_cache_unaligned(mut b: Bencher, size: Int) raises:
225227

226228
@parameter
227229
fn bench_and_same_offset(mut b: Bencher, size: Int) raises:
228-
"""AND of two bitmaps sharing the same non-zero sub-byte offset (pure SIMD, no shift)."""
230+
"""AND of two bitmaps sharing the same non-zero sub-byte offset (pure SIMD, no shift).
231+
"""
229232
var a = _make_half_set(size).slice(3, size - 8)
230233
var bm = _make_alternating(size).slice(3, size - 8)
231234

@@ -240,7 +243,8 @@ fn bench_and_same_offset(mut b: Bencher, size: Int) raises:
240243

241244
@parameter
242245
fn bench_and_diff_offset(mut b: Bencher, size: Int) raises:
243-
"""AND of two bitmaps with different sub-byte offsets (one-sided shift-combine)."""
246+
"""AND of two bitmaps with different sub-byte offsets (one-sided shift-combine).
247+
"""
244248
var a = _make_half_set(size).slice(3, size - 8)
245249
var bm = _make_alternating(size).slice(5, size - 8)
246250

@@ -261,7 +265,14 @@ fn bench_and_diff_offset(mut b: Bencher, size: Int) raises:
261265
def main() raises:
262266
var m = Bench(BenchConfig(num_repetitions=3))
263267

264-
comptime sizes = (1_000, 10_000, 100_000, 1_000_000, 10_000_000, 100_000_000)
268+
comptime sizes = (
269+
1_000,
270+
10_000,
271+
100_000,
272+
1_000_000,
273+
10_000_000,
274+
100_000_000,
275+
)
265276
comptime labels = ("1k", "10k", "100k", "1M", "10M", "100M")
266277

267278
comptime for si in range(6):

0 commit comments

Comments
 (0)