Skip to content

Commit

Permalink
Add getInts api for HostMemoryBuffer and UnsafeMemoryAccessor (#17767)
Browse files Browse the repository at this point in the history
Add getInts api for HostMemoryBuffer and UnsafeMemoryAccessor.

This pr is part of an optimization for optimizing concating validity buffer in kudo. It has no side effects, there already exists similar apis for `getShorts`, `getLongs`, `getBytes` etc. It didn't exist before because it was  not used, now we need to use it.

Authors:
  - Renjie Liu (https://github.com/liurenjie1024)

Approvers:
  - Liangcai Li (https://github.com/firestarman)
  - Chong Gao (https://github.com/res-life)

URL: #17767
  • Loading branch information
liurenjie1024 authored Jan 23, 2025
1 parent d825cfe commit 32d3fb1
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 4 deletions.
19 changes: 18 additions & 1 deletion java/src/main/java/ai/rapids/cudf/HostMemoryBuffer.java
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*
*
* Copyright (c) 2019-2024, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -367,6 +367,23 @@ public final int getInt(long offset) {
return UnsafeMemoryAccessor.getInt(requestedAddress);
}

/**
* Copy a set of ints to an array from the buffer starting at offset.
* @param dst destination int array
* @param dstIndex starting index within the destination array
* @param srcOffset starting offset within this buffer
* @param count number of ints to copy
*/
public final void getInts(int[] dst, long dstIndex, long srcOffset, int count) {
assert count >= 0;
assert count <= dst.length - dstIndex;
assert srcOffset >= 0;
long requestedAddress = this.address + srcOffset;
addressOutOfBoundsCheck(requestedAddress, count * 4L, "getInts");
UnsafeMemoryAccessor.getInts(dst, dstIndex, requestedAddress, count);
}


/**
* Sets the Integer value at that offset
* @param offset - offset from the address
Expand Down
15 changes: 14 additions & 1 deletion java/src/main/java/ai/rapids/cudf/UnsafeMemoryAccessor.java
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*
*
* Copyright (c) 2019, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -154,6 +154,19 @@ public static int getInt(long address) {
return UNSAFE.getInt(address);
}

/**
* Copy out an array of ints.
* @param dst where to write the data
* @param dstIndex index into values to start writing at.
* @param address src memory address
* @param count the number of ints to copy
* @throws IndexOutOfBoundsException
*/
public static void getInts(int[] dst, long dstIndex, long address, int count) {
copyMemory(null, address,
dst, UnsafeMemoryAccessor.INT_ARRAY_OFFSET + (dstIndex * 4), count * 4);
}

/**
* Sets the Integer value at that address
* @param address - memory address
Expand Down
16 changes: 15 additions & 1 deletion java/src/test/java/ai/rapids/cudf/HostMemoryBufferTest.java
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*
*
* Copyright (c) 2019-2024, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -87,6 +87,20 @@ public void testGetInt() {
}
}

@Test
public void testGetInts() {
try (HostMemoryBuffer hostMemoryBuffer = HostMemoryBuffer.allocate(16)) {
hostMemoryBuffer.setInt(0, 1);
hostMemoryBuffer.setInt(4, 2);
hostMemoryBuffer.setInt(8, 3);
hostMemoryBuffer.setInt(12, 4);
int[] expectedInts = new int[] {1, 2, 3, 4};
int[] result = new int[expectedInts.length];
hostMemoryBuffer.getInts(result, 0, 0, 4);
assertArrayEquals(expectedInts, result);
}
}

@Test
public void testGetByte() {
try (HostMemoryBuffer hostMemoryBuffer = HostMemoryBuffer.allocate(16)) {
Expand Down
20 changes: 19 additions & 1 deletion java/src/test/java/ai/rapids/cudf/UnsafeMemoryAccessorTest.java
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
/*
*
* Copyright (c) 2019-2023, NVIDIA CORPORATION.
* Copyright (c) 2019-2025, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -66,6 +66,24 @@ public void setIntAndGetInt() {
}
}

@Test
public void setAndGetInts() {
int numInts = 289;
long address = UnsafeMemoryAccessor.allocate(numInts * 4);
try {
for (int i = 0; i < numInts; i++) {
UnsafeMemoryAccessor.setInt(address + i * 4, i);
}
int[] ints = new int[numInts];
UnsafeMemoryAccessor.getInts(ints, 0, address, numInts);
for (int i = 0; i < numInts; i++) {
assertEquals(i, ints[i]);
}
} finally {
UnsafeMemoryAccessor.free(address);
}
}

@Test
public void setMemoryValue() {
long address = UnsafeMemoryAccessor.allocate(4);
Expand Down

0 comments on commit 32d3fb1

Please sign in to comment.