|
21 | 21 | import static org.junit.jupiter.api.Assertions.assertFalse; |
22 | 22 | import static org.junit.jupiter.api.Assertions.assertNull; |
23 | 23 | import static org.junit.jupiter.api.Assertions.assertSame; |
| 24 | +import static org.junit.jupiter.api.Assertions.assertThrows; |
24 | 25 | import static org.junit.jupiter.api.Assertions.assertTrue; |
| 26 | +import static org.junit.jupiter.api.Assertions.fail; |
25 | 27 |
|
26 | 28 | import java.util.ArrayList; |
27 | 29 | import java.util.Arrays; |
28 | 30 | import java.util.List; |
29 | 31 | import java.util.UUID; |
30 | 32 | import org.apache.arrow.memory.ArrowBuf; |
31 | 33 | import org.apache.arrow.memory.BufferAllocator; |
| 34 | +import org.apache.arrow.memory.OutOfMemoryException; |
32 | 35 | import org.apache.arrow.vector.complex.BaseRepeatedValueVector; |
33 | 36 | import org.apache.arrow.vector.complex.LargeListVector; |
34 | 37 | import org.apache.arrow.vector.complex.ListVector; |
|
42 | 45 | import org.apache.arrow.vector.types.pojo.ArrowType; |
43 | 46 | import org.apache.arrow.vector.types.pojo.Field; |
44 | 47 | import org.apache.arrow.vector.types.pojo.FieldType; |
| 48 | +import org.apache.arrow.vector.util.OversizedAllocationException; |
45 | 49 | import org.apache.arrow.vector.util.TransferPair; |
46 | 50 | import org.apache.arrow.vector.util.UuidUtility; |
47 | 51 | import org.junit.jupiter.api.AfterEach; |
@@ -1127,4 +1131,89 @@ private void writeIntValues(UnionLargeListWriter writer, int[] values) { |
1127 | 1131 | } |
1128 | 1132 | writer.endList(); |
1129 | 1133 | } |
| 1134 | + |
| 1135 | + @Test |
| 1136 | + public void testSetValueCountRejectsNegativeChildValueCount() { |
| 1137 | + try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { |
| 1138 | + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); |
| 1139 | + vector.allocateNew(); |
| 1140 | + |
| 1141 | + // Write a negative value into the offset buffer to simulate a corrupted offset. |
| 1142 | + // The Preconditions check should reject any negative childValueCount. |
| 1143 | + vector.getOffsetBuffer().setLong(LargeListVector.OFFSET_WIDTH, -1L); |
| 1144 | + vector.setLastSet(0); |
| 1145 | + |
| 1146 | + assertThrows( |
| 1147 | + IllegalArgumentException.class, |
| 1148 | + () -> vector.setValueCount(1), |
| 1149 | + "setValueCount should reject negative childValueCount"); |
| 1150 | + } |
| 1151 | + } |
| 1152 | + |
| 1153 | + @Test |
| 1154 | + public void testSetValueCountRejectsOverflowChildValueCount() { |
| 1155 | + try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { |
| 1156 | + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); |
| 1157 | + vector.allocateNew(); |
| 1158 | + |
| 1159 | + // Write a value exceeding Integer.MAX_VALUE into the offset buffer |
| 1160 | + vector.getOffsetBuffer().setLong(LargeListVector.OFFSET_WIDTH, (long) Integer.MAX_VALUE + 1L); |
| 1161 | + vector.setLastSet(0); |
| 1162 | + |
| 1163 | + assertThrows( |
| 1164 | + IllegalArgumentException.class, |
| 1165 | + () -> vector.setValueCount(1), |
| 1166 | + "setValueCount should reject childValueCount exceeding Integer.MAX_VALUE"); |
| 1167 | + } |
| 1168 | + } |
| 1169 | + |
| 1170 | + @Test |
| 1171 | + public void testSetValueCountAcceptsZeroChildValueCount() { |
| 1172 | + try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { |
| 1173 | + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); |
| 1174 | + vector.allocateNew(); |
| 1175 | + |
| 1176 | + // childValueCount = 0 when valueCount = 0 (no elements) |
| 1177 | + vector.setValueCount(0); |
| 1178 | + assertEquals(0, vector.getValueCount()); |
| 1179 | + } |
| 1180 | + } |
| 1181 | + |
| 1182 | + @Test |
| 1183 | + public void testSetValueCountAcceptsValidChildValueCount() { |
| 1184 | + try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { |
| 1185 | + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); |
| 1186 | + vector.allocateNew(); |
| 1187 | + |
| 1188 | + // Write a valid childValueCount (5) into the offset buffer |
| 1189 | + vector.getOffsetBuffer().setLong(LargeListVector.OFFSET_WIDTH, 5L); |
| 1190 | + vector.setLastSet(0); |
| 1191 | + vector.setValueCount(1); |
| 1192 | + assertEquals(1, vector.getValueCount()); |
| 1193 | + } |
| 1194 | + } |
| 1195 | + |
| 1196 | + @Test |
| 1197 | + public void testSetValueCountAcceptsMaxIntChildValueCount() { |
| 1198 | + try (final LargeListVector vector = LargeListVector.empty("list", allocator)) { |
| 1199 | + vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())); |
| 1200 | + vector.allocateNew(); |
| 1201 | + |
| 1202 | + // Write Integer.MAX_VALUE into the offset buffer to verify the upper boundary is inclusive. |
| 1203 | + // The Preconditions check should accept this value (not throw IllegalArgumentException). |
| 1204 | + // The child vector may throw OversizedAllocationException due to memory limits, |
| 1205 | + // which is expected and unrelated to the precondition validation. |
| 1206 | + vector.getOffsetBuffer().setLong(LargeListVector.OFFSET_WIDTH, Integer.MAX_VALUE); |
| 1207 | + vector.setLastSet(0); |
| 1208 | + try { |
| 1209 | + vector.setValueCount(1); |
| 1210 | + } catch (IllegalArgumentException e) { |
| 1211 | + fail("setValueCount should not reject childValueCount = Integer.MAX_VALUE", e); |
| 1212 | + } catch (OversizedAllocationException | OutOfMemoryException e) { |
| 1213 | + // OversizedAllocationException or other allocation errors are expected |
| 1214 | + // when trying to allocate Integer.MAX_VALUE elements — this is fine, |
| 1215 | + // the precondition check itself passed. |
| 1216 | + } |
| 1217 | + } |
| 1218 | + } |
1130 | 1219 | } |
0 commit comments