Skip to content

Commit 7eb176c

Browse files
committed
Commit #1 reading chunks, checkpoint to talk to nate, next will try reading schema at beginning of stream
1 parent 6fc2c6c commit 7eb176c

8 files changed

Lines changed: 381 additions & 190 deletions

File tree

extensions/barrage/src/main/java/io/deephaven/extensions/barrage/chunk/ChunkInputStreamGenerator.java

Lines changed: 6 additions & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import io.deephaven.chunk.attributes.Values;
1111
import io.deephaven.chunk.util.pools.PoolableChunk;
1212
import io.deephaven.engine.rowset.RowSet;
13-
import io.deephaven.extensions.barrage.ColumnConversionMode;
1413
import io.deephaven.extensions.barrage.util.DefensiveDrainable;
1514
import io.deephaven.extensions.barrage.util.StreamReaderOptions;
1615
import io.deephaven.time.DateTimeUtils;
@@ -19,7 +18,6 @@
1918
import io.deephaven.chunk.Chunk;
2019
import io.deephaven.chunk.ChunkType;
2120
import io.deephaven.util.SafeCloseable;
22-
import io.deephaven.util.type.TypeUtils;
2321
import io.deephaven.vector.Vector;
2422
import org.jetbrains.annotations.Nullable;
2523

@@ -31,7 +29,6 @@
3129
import java.time.LocalDate;
3230
import java.time.LocalTime;
3331
import java.time.ZonedDateTime;
34-
import java.util.Arrays;
3532
import java.util.Iterator;
3633
import java.util.PrimitiveIterator;
3734

@@ -187,6 +184,7 @@ static <T> ChunkInputStreamGenerator makeInputStreamGenerator(
187184
}
188185
}
189186

187+
@Deprecated
190188
static WritableChunk<Values> extractChunkFromInputStream(
191189
final StreamReaderOptions options,
192190
final ChunkType chunkType, final Class<?> type, final Class<?> componentType,
@@ -195,10 +193,10 @@ static WritableChunk<Values> extractChunkFromInputStream(
195193
final DataInput is,
196194
final WritableChunk<Values> outChunk, final int offset, final int totalRows) throws IOException {
197195
return extractChunkFromInputStream(options, 1, chunkType, type, componentType, fieldNodeIter, bufferInfoIter,
198-
is,
199-
outChunk, offset, totalRows);
196+
is, outChunk, offset, totalRows);
200197
}
201198

199+
@Deprecated
202200
static WritableChunk<Values> extractChunkFromInputStream(
203201
final StreamReaderOptions options,
204202
final int factor,
@@ -207,168 +205,9 @@ static WritableChunk<Values> extractChunkFromInputStream(
207205
final PrimitiveIterator.OfLong bufferInfoIter,
208206
final DataInput is,
209207
final WritableChunk<Values> outChunk, final int outOffset, final int totalRows) throws IOException {
210-
// TODO (deephaven-core#5453): pass in ArrowType to enable ser/deser of single java class in multiple formats
211-
switch (chunkType) {
212-
case Boolean:
213-
throw new UnsupportedOperationException("Booleans are reinterpreted as bytes");
214-
case Char:
215-
return CharChunkInputStreamGenerator.extractChunkFromInputStream(
216-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
217-
case Byte:
218-
if (type == Boolean.class || type == boolean.class) {
219-
return BooleanChunkInputStreamGenerator.extractChunkFromInputStream(
220-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
221-
}
222-
return ByteChunkInputStreamGenerator.extractChunkFromInputStream(
223-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
224-
case Short:
225-
return ShortChunkInputStreamGenerator.extractChunkFromInputStream(
226-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
227-
case Int:
228-
return IntChunkInputStreamGenerator.extractChunkFromInputStream(
229-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
230-
case Long:
231-
if (factor == 1) {
232-
return LongChunkInputStreamGenerator.extractChunkFromInputStream(
233-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
234-
}
235-
return LongChunkInputStreamGenerator.extractChunkFromInputStreamWithConversion(
236-
options, (long v) -> v == QueryConstants.NULL_LONG ? QueryConstants.NULL_LONG : (v * factor),
237-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
238-
case Float:
239-
return FloatChunkInputStreamGenerator.extractChunkFromInputStream(
240-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
241-
case Double:
242-
return DoubleChunkInputStreamGenerator.extractChunkFromInputStream(
243-
options, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
244-
case Object:
245-
if (type.isArray()) {
246-
if (componentType == byte.class) {
247-
return VarBinaryChunkInputStreamGenerator.extractChunkFromInputStream(
248-
is,
249-
fieldNodeIter,
250-
bufferInfoIter,
251-
(buf, off, len) -> Arrays.copyOfRange(buf, off, off + len),
252-
outChunk, outOffset, totalRows);
253-
} else {
254-
return VarListChunkInputStreamGenerator.extractChunkFromInputStream(
255-
options, type, fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
256-
}
257-
}
258-
if (Vector.class.isAssignableFrom(type)) {
259-
// noinspection unchecked
260-
return VectorChunkInputStreamGenerator.extractChunkFromInputStream(
261-
options, (Class<Vector<?>>) type, componentType, fieldNodeIter, bufferInfoIter, is,
262-
outChunk, outOffset, totalRows);
263-
}
264-
if (type == BigInteger.class) {
265-
return VarBinaryChunkInputStreamGenerator.extractChunkFromInputStream(
266-
is,
267-
fieldNodeIter,
268-
bufferInfoIter,
269-
BigInteger::new,
270-
outChunk, outOffset, totalRows);
271-
}
272-
if (type == BigDecimal.class) {
273-
return VarBinaryChunkInputStreamGenerator.extractChunkFromInputStream(
274-
is,
275-
fieldNodeIter,
276-
bufferInfoIter,
277-
(final byte[] buf, final int offset, final int length) -> {
278-
// read the int scale value as little endian, arrow's endianness.
279-
final byte b1 = buf[offset];
280-
final byte b2 = buf[offset + 1];
281-
final byte b3 = buf[offset + 2];
282-
final byte b4 = buf[offset + 3];
283-
final int scale = b4 << 24 | (b3 & 0xFF) << 16 | (b2 & 0xFF) << 8 | (b1 & 0xFF);
284-
return new BigDecimal(new BigInteger(buf, offset + 4, length - 4), scale);
285-
},
286-
outChunk, outOffset, totalRows);
287-
}
288-
if (type == Instant.class) {
289-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
290-
Long.BYTES, options, io -> {
291-
final long value = io.readLong();
292-
if (value == QueryConstants.NULL_LONG) {
293-
return null;
294-
}
295-
return DateTimeUtils.epochNanosToInstant(value * factor);
296-
},
297-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
298-
}
299-
if (type == ZonedDateTime.class) {
300-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
301-
Long.BYTES, options, io -> {
302-
final long value = io.readLong();
303-
if (value == QueryConstants.NULL_LONG) {
304-
return null;
305-
}
306-
return DateTimeUtils.epochNanosToZonedDateTime(
307-
value * factor, DateTimeUtils.timeZone());
308-
},
309-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
310-
}
311-
if (type == Byte.class) {
312-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
313-
Byte.BYTES, options, io -> TypeUtils.box(io.readByte()),
314-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
315-
}
316-
if (type == Character.class) {
317-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
318-
Character.BYTES, options, io -> TypeUtils.box(io.readChar()),
319-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
320-
}
321-
if (type == Double.class) {
322-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
323-
Double.BYTES, options, io -> TypeUtils.box(io.readDouble()),
324-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
325-
}
326-
if (type == Float.class) {
327-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
328-
Float.BYTES, options, io -> TypeUtils.box(io.readFloat()),
329-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
330-
}
331-
if (type == Integer.class) {
332-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
333-
Integer.BYTES, options, io -> TypeUtils.box(io.readInt()),
334-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
335-
}
336-
if (type == Long.class) {
337-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
338-
Long.BYTES, options, io -> TypeUtils.box(io.readLong()),
339-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
340-
}
341-
if (type == Short.class) {
342-
return FixedWidthChunkInputStreamGenerator.extractChunkFromInputStreamWithTypeConversion(
343-
Short.BYTES, options, io -> TypeUtils.box(io.readShort()),
344-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
345-
}
346-
if (type == LocalDate.class) {
347-
return LongChunkInputStreamGenerator.extractChunkFromInputStreamWithTransform(
348-
options,
349-
value -> value == QueryConstants.NULL_LONG
350-
? null
351-
: LocalDate.ofEpochDay(value / MS_PER_DAY),
352-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
353-
}
354-
if (type == LocalTime.class) {
355-
return LongChunkInputStreamGenerator.extractChunkFromInputStreamWithTransform(
356-
options,
357-
value -> value == QueryConstants.NULL_LONG ? null : LocalTime.ofNanoOfDay(value),
358-
fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
359-
}
360-
if (type == String.class ||
361-
options.columnConversionMode().equals(ColumnConversionMode.Stringify)) {
362-
return VarBinaryChunkInputStreamGenerator.extractChunkFromInputStream(is, fieldNodeIter,
363-
bufferInfoIter,
364-
(buf, off, len) -> new String(buf, off, len, Charsets.UTF_8), outChunk, outOffset,
365-
totalRows);
366-
}
367-
throw new UnsupportedOperationException(
368-
"Do not yet support column conversion mode: " + options.columnConversionMode());
369-
default:
370-
throw new UnsupportedOperationException();
371-
}
208+
return DefaultChunkReadingFactory.INSTANCE.extractChunkFromInputStream(options, factor,
209+
new ChunkReadingFactory.ChunkTypeInfo(chunkType, type, componentType, null), fieldNodeIter,
210+
bufferInfoIter, is, outChunk, outOffset, totalRows);
372211
}
373212

374213
/**
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
//
2+
// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
3+
//
4+
package io.deephaven.extensions.barrage.chunk;
5+
6+
import io.deephaven.chunk.ChunkType;
7+
import io.deephaven.chunk.WritableChunk;
8+
import io.deephaven.chunk.attributes.Values;
9+
import io.deephaven.extensions.barrage.util.StreamReaderOptions;
10+
import org.apache.arrow.flatbuf.Field;
11+
import org.apache.arrow.flatbuf.Type;
12+
13+
import java.io.DataInput;
14+
import java.io.IOException;
15+
import java.util.Iterator;
16+
import java.util.PrimitiveIterator;
17+
18+
/**
19+
*
20+
*/
21+
public interface ChunkReadingFactory {
22+
/**
23+
*
24+
*/
25+
class ChunkTypeInfo {
26+
private final ChunkType chunkType;
27+
private final Class<?> type;
28+
private final Class<?> componentType;
29+
private final Field arrowField;
30+
31+
public ChunkTypeInfo(ChunkType chunkType, Class<?> type, Class<?> componentType, Field arrowField) {
32+
this.chunkType = chunkType;
33+
this.type = type;
34+
this.componentType = componentType;
35+
this.arrowField = arrowField;
36+
}
37+
38+
public ChunkType chunkType() {
39+
return chunkType;
40+
}
41+
42+
public Class<?> type() {
43+
return type;
44+
}
45+
46+
public Class<?> componentType() {
47+
return componentType;
48+
}
49+
50+
public Field arrowField() {
51+
return arrowField;
52+
}
53+
54+
public Field componentArrowField() {
55+
if (arrowField.typeType() != Type.List) {
56+
throw new IllegalStateException("Not a flight List");
57+
}
58+
if (arrowField.childrenLength() != 1) {
59+
throw new IllegalStateException("Incorrect number of child Fields");
60+
}
61+
return arrowField.children(0);
62+
}
63+
}
64+
65+
/**
66+
*
67+
* @param options
68+
* @param factor
69+
* @param typeInfo
70+
* @param fieldNodeIter
71+
* @param bufferInfoIter
72+
* @param is
73+
* @param outChunk
74+
* @param outOffset
75+
* @param totalRows
76+
* @return
77+
* @throws IOException
78+
*/
79+
WritableChunk<Values> extractChunkFromInputStream(
80+
final StreamReaderOptions options,
81+
final int factor,
82+
final ChunkTypeInfo typeInfo,
83+
final Iterator<ChunkInputStreamGenerator.FieldNodeInfo> fieldNodeIter,
84+
final PrimitiveIterator.OfLong bufferInfoIter,
85+
final DataInput is,
86+
final WritableChunk<Values> outChunk,
87+
final int outOffset,
88+
final int totalRows) throws IOException;
89+
90+
/**
91+
*
92+
* @param options
93+
* @param typeInfo
94+
* @param fieldNodeIter
95+
* @param bufferInfoIter
96+
* @param is
97+
* @param outChunk
98+
* @param offset
99+
* @param totalRows
100+
* @return
101+
* @throws IOException
102+
*/
103+
default WritableChunk<Values> extractChunkFromInputStream(
104+
final StreamReaderOptions options,
105+
final ChunkTypeInfo typeInfo,
106+
final Iterator<ChunkInputStreamGenerator.FieldNodeInfo> fieldNodeIter,
107+
final PrimitiveIterator.OfLong bufferInfoIter,
108+
final DataInput is,
109+
final WritableChunk<Values> outChunk,
110+
final int offset,
111+
final int totalRows) throws IOException {
112+
return extractChunkFromInputStream(options, 1, typeInfo, fieldNodeIter, bufferInfoIter, is, outChunk, offset,
113+
totalRows);
114+
}
115+
116+
}

0 commit comments

Comments
 (0)