/
UnsafeArrayData.java
320 lines (285 loc) · 11.1 KB
/
UnsafeArrayData.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.spark.sql.catalyst.expressions;
import java.math.BigDecimal;
import java.math.BigInteger;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.types.*;
import org.apache.spark.unsafe.Platform;
import org.apache.spark.unsafe.array.ByteArrayMethods;
import org.apache.spark.unsafe.hash.Murmur3_x86_32;
import org.apache.spark.unsafe.types.CalendarInterval;
import org.apache.spark.unsafe.types.UTF8String;
/**
* An Unsafe implementation of Array which is backed by raw memory instead of Java objects.
*
* Each tuple has two parts: [offsets] [values]
*
* In the `offsets` region, we store 4 bytes per element, represents the start address of this
* element in `values` region. We can get the length of this element by subtracting next offset.
* Note that offset can by negative which means this element is null.
*
* In the `values` region, we store the content of elements. As we can get length info, so elements
* can be variable-length.
*
* Note that when we write out this array, we should write out the `numElements` at first 4 bytes,
* then follows content. When we read in an array, we should read first 4 bytes as `numElements`
* and take the rest as content.
*
* Instances of `UnsafeArrayData` act as pointers to row data stored in this format.
*/
// todo: there is a lof of duplicated code between UnsafeRow and UnsafeArrayData.
public class UnsafeArrayData extends ArrayData {
private Object baseObject;
private long baseOffset;
// The number of elements in this array
private int numElements;
// The size of this array's backing data, in bytes
private int sizeInBytes;
private int getElementOffset(int ordinal) {
return Platform.getInt(baseObject, baseOffset + ordinal * 4L);
}
private int getElementSize(int offset, int ordinal) {
if (ordinal == numElements - 1) {
return sizeInBytes - offset;
} else {
return Math.abs(getElementOffset(ordinal + 1)) - offset;
}
}
private void assertIndexIsValid(int ordinal) {
assert ordinal >= 0 : "ordinal (" + ordinal + ") should >= 0";
assert ordinal < numElements : "ordinal (" + ordinal + ") should < " + numElements;
}
/**
* Construct a new UnsafeArrayData. The resulting UnsafeArrayData won't be usable until
* `pointTo()` has been called, since the value returned by this constructor is equivalent
* to a null pointer.
*/
public UnsafeArrayData() { }
public Object getBaseObject() { return baseObject; }
public long getBaseOffset() { return baseOffset; }
public int getSizeInBytes() { return sizeInBytes; }
@Override
public int numElements() { return numElements; }
/**
* Update this UnsafeArrayData to point to different backing data.
*
* @param baseObject the base object
* @param baseOffset the offset within the base object
* @param sizeInBytes the size of this row's backing data, in bytes
*/
public void pointTo(Object baseObject, long baseOffset, int numElements, int sizeInBytes) {
assert numElements >= 0 : "numElements (" + numElements + ") should >= 0";
this.numElements = numElements;
this.baseObject = baseObject;
this.baseOffset = baseOffset;
this.sizeInBytes = sizeInBytes;
}
@Override
public boolean isNullAt(int ordinal) {
assertIndexIsValid(ordinal);
return getElementOffset(ordinal) < 0;
}
@Override
public Object get(int ordinal, DataType dataType) {
if (isNullAt(ordinal) || dataType instanceof NullType) {
return null;
} else if (dataType instanceof BooleanType) {
return getBoolean(ordinal);
} else if (dataType instanceof ByteType) {
return getByte(ordinal);
} else if (dataType instanceof ShortType) {
return getShort(ordinal);
} else if (dataType instanceof IntegerType) {
return getInt(ordinal);
} else if (dataType instanceof LongType) {
return getLong(ordinal);
} else if (dataType instanceof FloatType) {
return getFloat(ordinal);
} else if (dataType instanceof DoubleType) {
return getDouble(ordinal);
} else if (dataType instanceof DecimalType) {
DecimalType dt = (DecimalType) dataType;
return getDecimal(ordinal, dt.precision(), dt.scale());
} else if (dataType instanceof DateType) {
return getInt(ordinal);
} else if (dataType instanceof TimestampType) {
return getLong(ordinal);
} else if (dataType instanceof BinaryType) {
return getBinary(ordinal);
} else if (dataType instanceof StringType) {
return getUTF8String(ordinal);
} else if (dataType instanceof CalendarIntervalType) {
return getInterval(ordinal);
} else if (dataType instanceof StructType) {
return getStruct(ordinal, ((StructType) dataType).size());
} else if (dataType instanceof ArrayType) {
return getArray(ordinal);
} else if (dataType instanceof MapType) {
return getMap(ordinal);
} else {
throw new UnsupportedOperationException("Unsupported data type " + dataType.simpleString());
}
}
@Override
public boolean getBoolean(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return false;
return Platform.getBoolean(baseObject, baseOffset + offset);
}
@Override
public byte getByte(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return 0;
return Platform.getByte(baseObject, baseOffset + offset);
}
@Override
public short getShort(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return 0;
return Platform.getShort(baseObject, baseOffset + offset);
}
@Override
public int getInt(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return 0;
return Platform.getInt(baseObject, baseOffset + offset);
}
@Override
public long getLong(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return 0;
return Platform.getLong(baseObject, baseOffset + offset);
}
@Override
public float getFloat(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return 0;
return Platform.getFloat(baseObject, baseOffset + offset);
}
@Override
public double getDouble(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return 0;
return Platform.getDouble(baseObject, baseOffset + offset);
}
@Override
public Decimal getDecimal(int ordinal, int precision, int scale) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
if (precision <= Decimal.MAX_LONG_DIGITS()) {
final long value = Platform.getLong(baseObject, baseOffset + offset);
return Decimal.apply(value, precision, scale);
} else {
final byte[] bytes = getBinary(ordinal);
final BigInteger bigInteger = new BigInteger(bytes);
final BigDecimal javaDecimal = new BigDecimal(bigInteger, scale);
return Decimal.apply(new scala.math.BigDecimal(javaDecimal), precision, scale);
}
}
@Override
public UTF8String getUTF8String(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
final int size = getElementSize(offset, ordinal);
return UTF8String.fromAddress(baseObject, baseOffset + offset, size);
}
@Override
public byte[] getBinary(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
final int size = getElementSize(offset, ordinal);
final byte[] bytes = new byte[size];
Platform.copyMemory(baseObject, baseOffset + offset, bytes, Platform.BYTE_ARRAY_OFFSET, size);
return bytes;
}
@Override
public CalendarInterval getInterval(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
final int months = (int) Platform.getLong(baseObject, baseOffset + offset);
final long microseconds = Platform.getLong(baseObject, baseOffset + offset + 8);
return new CalendarInterval(months, microseconds);
}
@Override
public InternalRow getStruct(int ordinal, int numFields) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
final int size = getElementSize(offset, ordinal);
final UnsafeRow row = new UnsafeRow();
row.pointTo(baseObject, baseOffset + offset, numFields, size);
return row;
}
@Override
public ArrayData getArray(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
final int size = getElementSize(offset, ordinal);
return UnsafeReaders.readArray(baseObject, baseOffset + offset, size);
}
@Override
public MapData getMap(int ordinal) {
assertIndexIsValid(ordinal);
final int offset = getElementOffset(ordinal);
if (offset < 0) return null;
final int size = getElementSize(offset, ordinal);
return UnsafeReaders.readMap(baseObject, baseOffset + offset, size);
}
@Override
public int hashCode() {
int result = 37;
for (int i = 0; i < sizeInBytes; i++) {
result = 37 * result + Platform.getByte(baseObject, baseOffset + i);
}
return result;
}
@Override
public boolean equals(Object other) {
if (other instanceof UnsafeArrayData) {
UnsafeArrayData o = (UnsafeArrayData) other;
return (sizeInBytes == o.sizeInBytes) &&
ByteArrayMethods.arrayEquals(baseObject, baseOffset, o.baseObject, o.baseOffset,
sizeInBytes);
}
return false;
}
public void writeToMemory(Object target, long targetOffset) {
Platform.copyMemory(baseObject, baseOffset, target, targetOffset, sizeInBytes);
}
@Override
public UnsafeArrayData copy() {
UnsafeArrayData arrayCopy = new UnsafeArrayData();
final byte[] arrayDataCopy = new byte[sizeInBytes];
Platform.copyMemory(
baseObject, baseOffset, arrayDataCopy, Platform.BYTE_ARRAY_OFFSET, sizeInBytes);
arrayCopy.pointTo(arrayDataCopy, Platform.BYTE_ARRAY_OFFSET, numElements, sizeInBytes);
return arrayCopy;
}
}