Skip to content

Commit 75a076b

Browse files
niloc132niloc132
committed
Commit #2, mostly mechanical changes, splitting creation and reading
1 parent c746c34 commit 75a076b

8 files changed

Lines changed: 233 additions & 195 deletions

File tree

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ static WritableChunk<Values> extractChunkFromInputStream(
197197
}
198198

199199
@Deprecated
200-
static WritableChunk<Values> extractChunkFromInputStream(
200+
private static WritableChunk<Values> extractChunkFromInputStream(
201201
final StreamReaderOptions options,
202202
final int factor,
203203
final ChunkType chunkType, final Class<?> type, final Class<?> componentType,
@@ -206,8 +206,8 @@ static WritableChunk<Values> extractChunkFromInputStream(
206206
final DataInput is,
207207
final WritableChunk<Values> outChunk, final int outOffset, final int totalRows) throws IOException {
208208
return DefaultChunkReadingFactory.INSTANCE.extractChunkFromInputStream(options, factor,
209-
new ChunkReadingFactory.ChunkTypeInfo(chunkType, type, componentType, null), fieldNodeIter,
210-
bufferInfoIter, is, outChunk, outOffset, totalRows);
209+
new ChunkReadingFactory.ChunkTypeInfo(chunkType, type, componentType, null))
210+
.read(fieldNodeIter, bufferInfoIter, is, outChunk, outOffset, totalRows);
211211
}
212212

213213
/**

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/chunk/ChunkReader.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
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.WritableChunk;
7+
import io.deephaven.chunk.attributes.Values;
8+
9+
import java.io.DataInput;
10+
import java.io.IOException;
11+
import java.util.Iterator;
12+
import java.util.PrimitiveIterator;
13+
14+
/**
15+
* Consumes Flight/Barrage streams and transforms them into WritableChunks.
16+
*/
17+
public interface ChunkReader {
18+
/**
19+
*
20+
* @param fieldNodeIter
21+
* @param bufferInfoIter
22+
* @param is
23+
* @param outChunk
24+
* @param outOffset
25+
* @param totalRows
26+
* @return
27+
*/
28+
WritableChunk<Values> read(final Iterator<ChunkInputStreamGenerator.FieldNodeInfo> fieldNodeIter,
29+
final PrimitiveIterator.OfLong bufferInfoIter,
30+
final DataInput is,
31+
final WritableChunk<Values> outChunk,
32+
final int outOffset,
33+
final int totalRows) throws IOException;
34+
}

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/chunk/ChunkReadingFactory.java‎

Lines changed: 5 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -67,50 +67,24 @@ public Field componentArrowField() {
6767
* @param options
6868
* @param factor
6969
* @param typeInfo
70-
* @param fieldNodeIter
71-
* @param bufferInfoIter
72-
* @param is
73-
* @param outChunk
74-
* @param outOffset
75-
* @param totalRows
7670
* @return
7771
* @throws IOException
7872
*/
79-
WritableChunk<Values> extractChunkFromInputStream(
73+
ChunkReader extractChunkFromInputStream(
8074
final StreamReaderOptions options,
8175
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;
76+
final ChunkTypeInfo typeInfo) throws IOException;
8977

9078
/**
9179
*
9280
* @param options
9381
* @param typeInfo
94-
* @param fieldNodeIter
95-
* @param bufferInfoIter
96-
* @param is
97-
* @param outChunk
98-
* @param offset
99-
* @param totalRows
10082
* @return
10183
* @throws IOException
10284
*/
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);
85+
default ChunkReader extractChunkFromInputStream(final StreamReaderOptions options, final ChunkTypeInfo typeInfo)
86+
throws IOException {
87+
return extractChunkFromInputStream(options, 1, typeInfo);
11488
}
11589

11690
}

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/chunk/DefaultChunkReadingFactory.java‎

Lines changed: 145 additions & 111 deletions
Large diffs are not rendered by default.

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/chunk/VarListChunkInputStreamGenerator.java‎

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -244,10 +244,6 @@ static <T> WritableObjectChunk<T, Values> extractChunkFromInputStream(
244244
final int totalRows,
245245
ChunkReadingFactory chunkReadingFactory) throws IOException {
246246

247-
final FieldNodeInfo nodeInfo = fieldNodeIter.next();
248-
final long validityBuffer = bufferInfoIter.nextLong();
249-
final long offsetsBuffer = bufferInfoIter.nextLong();
250-
251247
final Class<?> componentType = typeInfo.type().getComponentType();
252248
final Class<?> innerComponentType = componentType != null ? componentType.getComponentType() : null;
253249

@@ -259,13 +255,18 @@ static <T> WritableObjectChunk<T, Values> extractChunkFromInputStream(
259255
chunkType = ChunkType.fromElementType(componentType);
260256
}
261257

258+
ChunkReader componentReader = chunkReadingFactory.extractChunkFromInputStream(
259+
options,
260+
new ChunkReadingFactory.ChunkTypeInfo(chunkType, componentType, innerComponentType,
261+
typeInfo.componentArrowField()));
262+
263+
final FieldNodeInfo nodeInfo = fieldNodeIter.next();
264+
final long validityBuffer = bufferInfoIter.nextLong();
265+
final long offsetsBuffer = bufferInfoIter.nextLong();
266+
262267
if (nodeInfo.numElements == 0) {
263-
try (final WritableChunk<Values> ignored = chunkReadingFactory.extractChunkFromInputStream(
264-
options,
265-
new ChunkReadingFactory.ChunkTypeInfo(chunkType, componentType, innerComponentType,
266-
typeInfo.componentArrowField()),
267-
fieldNodeIter,
268-
bufferInfoIter, is, null, 0, 0)) {
268+
try (final WritableChunk<Values> ignored =
269+
componentReader.read(fieldNodeIter, bufferInfoIter, is, null, 0, 0)) {
269270
return WritableObjectChunk.makeWritableChunk(nodeInfo.numElements);
270271
}
271272
}
@@ -303,11 +304,8 @@ static <T> WritableObjectChunk<T, Values> extractChunkFromInputStream(
303304
}
304305

305306
final ArrayExpansionKernel kernel = ArrayExpansionKernel.makeExpansionKernel(chunkType, componentType);
306-
try (final WritableChunk<Values> inner = chunkReadingFactory.extractChunkFromInputStream(
307-
options,
308-
new ChunkReadingFactory.ChunkTypeInfo(chunkType, componentType, innerComponentType,
309-
typeInfo.componentArrowField()),
310-
fieldNodeIter, bufferInfoIter, is, null, 0, 0)) {
307+
try (final WritableChunk<Values> inner =
308+
componentReader.read(fieldNodeIter, bufferInfoIter, is, null, 0, 0)) {
311309
chunk = kernel.contract(inner, offsets, outChunk, outOffset, totalRows);
312310

313311
long nextValid = 0;

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/chunk/VectorChunkInputStreamGenerator.java‎

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -244,22 +244,22 @@ static WritableObjectChunk<Vector<?>, Values> extractChunkFromInputStream(
244244
final int totalRows,
245245
ChunkReadingFactory chunkReadingFactory) throws IOException {
246246

247+
final Class<?> componentType =
248+
VectorExpansionKernel.getComponentType(typeInfo.type(), typeInfo.componentType());
249+
final ChunkType chunkType = ChunkType.fromElementType(componentType);
250+
ChunkReader componentReader = chunkReadingFactory.extractChunkFromInputStream(
251+
options,
252+
new ChunkReadingFactory.ChunkTypeInfo(chunkType, componentType, componentType.getComponentType(),
253+
typeInfo.componentArrowField()));
254+
247255
final FieldNodeInfo nodeInfo = fieldNodeIter.next();
248256
final long validityBuffer = bufferInfoIter.nextLong();
249257
final long offsetsBuffer = bufferInfoIter.nextLong();
250258

251-
final Class<?> componentType =
252-
VectorExpansionKernel.getComponentType(typeInfo.type(), typeInfo.componentType());
253-
final ChunkType chunkType = ChunkType.fromElementType(componentType);
254259

255260
if (nodeInfo.numElements == 0) {
256-
try (final WritableChunk<Values> ignored = chunkReadingFactory.extractChunkFromInputStream(
257-
options,
258-
new ChunkReadingFactory.ChunkTypeInfo(chunkType, componentType, componentType.getComponentType(),
259-
typeInfo.componentArrowField()),
260-
fieldNodeIter, bufferInfoIter,
261-
is,
262-
null, 0, 0)) {
261+
try (final WritableChunk<Values> ignored =
262+
componentReader.read(fieldNodeIter, bufferInfoIter, is, null, 0, 0)) {
263263
if (outChunk != null) {
264264
return outChunk.asWritableObjectChunk();
265265
}
@@ -300,13 +300,8 @@ static WritableObjectChunk<Vector<?>, Values> extractChunkFromInputStream(
300300
}
301301

302302
final VectorExpansionKernel kernel = VectorExpansionKernel.makeExpansionKernel(chunkType, componentType);
303-
try (final WritableChunk<Values> inner = chunkReadingFactory.extractChunkFromInputStream(
304-
options,
305-
new ChunkReadingFactory.ChunkTypeInfo(chunkType, componentType, componentType.getComponentType(),
306-
typeInfo.componentArrowField()),
307-
fieldNodeIter, bufferInfoIter,
308-
is,
309-
null, 0, 0)) {
303+
try (final WritableChunk<Values> inner =
304+
componentReader.read(fieldNodeIter, bufferInfoIter, is, null, 0, 0)) {
310305
chunk = kernel.contract(inner, offsets, outChunk, outOffset, totalRows);
311306

312307
long nextValid = 0;

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/ArrowToTableConverter.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import io.deephaven.engine.table.impl.util.BarrageMessage;
1414
import io.deephaven.extensions.barrage.BarrageSubscriptionOptions;
1515
import io.deephaven.extensions.barrage.chunk.ChunkInputStreamGenerator;
16+
import io.deephaven.extensions.barrage.chunk.ChunkReader;
1617
import io.deephaven.extensions.barrage.chunk.ChunkReadingFactory;
1718
import io.deephaven.extensions.barrage.chunk.DefaultChunkReadingFactory;
1819
import io.deephaven.extensions.barrage.table.BarrageTable;
@@ -198,11 +199,10 @@ protected BarrageMessage createBarrageMessage(BarrageProtoUtil.MessageInfo mi, i
198199
msg.addColumnData[ci].data = new ArrayList<>();
199200
final int factor = (columnConversionFactors == null) ? 1 : columnConversionFactors[ci];
200201
try {
201-
acd.data.add(DefaultChunkReadingFactory.INSTANCE.extractChunkFromInputStream(options, factor,
202+
ChunkReader reader = DefaultChunkReadingFactory.INSTANCE.extractChunkFromInputStream(options, factor,
202203
new ChunkReadingFactory.ChunkTypeInfo(columnChunkTypes[ci], columnTypes[ci], componentTypes[ci],
203-
schema.fields(ci)),
204-
fieldNodeIter,
205-
bufferInfoIter, mi.inputStream, null, 0, 0));
204+
schema.fields(ci)));
205+
acd.data.add(reader.read(fieldNodeIter, bufferInfoIter, mi.inputStream, null, 0, 0));
206206
} catch (final IOException unexpected) {
207207
throw new UncheckedDeephavenException(unexpected);
208208
}

‎extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/BarrageStreamReader.java‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
import io.deephaven.engine.rowset.RowSetShiftData;
2020
import io.deephaven.engine.table.impl.util.*;
2121
import io.deephaven.extensions.barrage.chunk.ChunkInputStreamGenerator;
22+
import io.deephaven.extensions.barrage.chunk.ChunkReader;
2223
import io.deephaven.extensions.barrage.chunk.ChunkReadingFactory;
2324
import io.deephaven.extensions.barrage.chunk.DefaultChunkReadingFactory;
2425
import io.deephaven.util.datastructures.LongSizedDataStructure;
2526
import io.deephaven.chunk.ChunkType;
2627
import io.deephaven.internal.log.LoggerFactory;
2728
import io.deephaven.io.logger.Logger;
29+
import org.apache.arrow.flatbuf.Field;
2830
import org.apache.arrow.flatbuf.Message;
2931
import org.apache.arrow.flatbuf.MessageHeader;
3032
import org.apache.arrow.flatbuf.RecordBatch;
@@ -37,6 +39,7 @@
3739
import java.util.Arrays;
3840
import java.util.BitSet;
3941
import java.util.Iterator;
42+
import java.util.List;
4043
import java.util.PrimitiveIterator;
4144
import java.util.function.LongConsumer;
4245

@@ -54,10 +57,10 @@ public class BarrageStreamReader implements StreamReader {
5457
private long numModRowsRead = 0;
5558
private long numModRowsTotal = 0;
5659

57-
private Schema schema;
5860
private BarrageMessage msg = null;
5961

6062
private final ChunkReadingFactory chunkReadingFactory = DefaultChunkReadingFactory.INSTANCE;
63+
private final List<ChunkReader> readers = new ArrayList<>();
6164

6265
public BarrageStreamReader(final LongConsumer deserializeTmConsumer) {
6366
this.deserializeTmConsumer = deserializeTmConsumer;
@@ -244,12 +247,8 @@ public BarrageMessage safelyParseFrom(final StreamReaderOptions options,
244247
}
245248

246249
// fill the chunk with data and assign back into the array
247-
acd.data.set(lastChunkIndex,
248-
chunkReadingFactory.extractChunkFromInputStream(options,
249-
new ChunkReadingFactory.ChunkTypeInfo(columnChunkTypes[ci],
250-
columnTypes[ci], componentTypes[ci], schema.fields(ci)),
251-
fieldNodeIter, bufferInfoIter, ois,
252-
chunk, chunk.size(), (int) batch.length()));
250+
acd.data.set(lastChunkIndex, readers.get(ci).read(fieldNodeIter, bufferInfoIter, ois, chunk,
251+
chunk.size(), (int) batch.length()));
253252
chunk.setSize(chunk.size() + (int) batch.length());
254253
}
255254
numAddRowsRead += batch.length();
@@ -277,12 +276,8 @@ public BarrageMessage safelyParseFrom(final StreamReaderOptions options,
277276
}
278277

279278
// fill the chunk with data and assign back into the array
280-
mcd.data.set(lastChunkIndex,
281-
chunkReadingFactory.extractChunkFromInputStream(options,
282-
new ChunkReadingFactory.ChunkTypeInfo(columnChunkTypes[ci],
283-
columnTypes[ci], componentTypes[ci], null),
284-
fieldNodeIter, bufferInfoIter, ois,
285-
chunk, chunk.size(), numRowsToRead));
279+
mcd.data.set(lastChunkIndex, readers.get(ci).read(fieldNodeIter, bufferInfoIter, ois, chunk,
280+
chunk.size(), numRowsToRead));
286281
chunk.setSize(chunk.size() + numRowsToRead);
287282
}
288283
numModRowsRead += batch.length();
@@ -292,7 +287,15 @@ public BarrageMessage safelyParseFrom(final StreamReaderOptions options,
292287

293288
if (header != null && header.headerType() == MessageHeader.Schema) {
294289
// there is no body and our clients do not want to see schema messages
295-
this.schema = (Schema) header.header(new Schema());
290+
Schema schema = new Schema();
291+
header.header(schema);
292+
for (int i = 0; i < schema.fieldsLength(); i++) {
293+
Field field = schema.fields(i);
294+
ChunkReader chunkReader = chunkReadingFactory.extractChunkFromInputStream(options,
295+
new ChunkReadingFactory.ChunkTypeInfo(columnChunkTypes[i], columnTypes[i],
296+
componentTypes[i], field));
297+
readers.add(chunkReader);
298+
}
296299
return null;
297300
}
298301

0 commit comments

Comments
 (0)