@@ -161,24 +161,27 @@ public void testWriteReadComplex() throws IOException {
161161 @ Test
162162 public void testWriteReadMultipleRBs () throws IOException {
163163 File file = new File ("target/mytest_multiple.arrow" );
164- int count = COUNT ;
164+ int [] counts = { 10 , 5 } ;
165165
166166 // write
167167 try (
168168 BufferAllocator originalVectorAllocator = allocator .newChildAllocator ("original vectors" , 0 , Integer .MAX_VALUE );
169169 MapVector parent = new MapVector ("parent" , originalVectorAllocator , null );
170170 FileOutputStream fileOutputStream = new FileOutputStream (file );) {
171- writeData (count , parent );
172- VectorUnloader vectorUnloader = newVectorUnloader (parent .getChild ("root" ));
173- Schema schema = vectorUnloader .getSchema ();
171+ writeData (counts [ 0 ] , parent );
172+ VectorUnloader vectorUnloader0 = newVectorUnloader (parent .getChild ("root" ));
173+ Schema schema = vectorUnloader0 .getSchema ();
174174 Assert .assertEquals (2 , schema .getFields ().size ());
175175 try (ArrowWriter arrowWriter = new ArrowWriter (fileOutputStream .getChannel (), schema );) {
176- try (ArrowRecordBatch recordBatch = vectorUnloader .getRecordBatch ()) {
176+ try (ArrowRecordBatch recordBatch = vectorUnloader0 .getRecordBatch ()) {
177+ Assert .assertEquals ("RB #0" , counts [0 ], recordBatch .getLength ());
177178 arrowWriter .writeRecordBatch (recordBatch );
178179 }
179180 parent .allocateNew ();
180- writeData (count , parent );
181- try (ArrowRecordBatch recordBatch = vectorUnloader .getRecordBatch ()) {
181+ writeData (counts [1 ], parent ); // if we write the same data we don't catch that the metadata is stored in the wrong order.
182+ VectorUnloader vectorUnloader1 = newVectorUnloader (parent .getChild ("root" ));
183+ try (ArrowRecordBatch recordBatch = vectorUnloader1 .getRecordBatch ()) {
184+ Assert .assertEquals ("RB #1" , counts [1 ], recordBatch .getLength ());
182185 arrowWriter .writeRecordBatch (recordBatch );
183186 }
184187 }
@@ -195,21 +198,27 @@ public void testWriteReadMultipleRBs() throws IOException {
195198 ArrowFooter footer = arrowReader .readFooter ();
196199 Schema schema = footer .getSchema ();
197200 LOGGER .debug ("reading schema: " + schema );
201+ int i = 0 ;
198202 try (VectorSchemaRoot root = new VectorSchemaRoot (schema , vectorAllocator );) {
199203 VectorLoader vectorLoader = new VectorLoader (root );
200204 List <ArrowBlock > recordBatches = footer .getRecordBatches ();
201205 Assert .assertEquals (2 , recordBatches .size ());
206+ long previousOffset = 0 ;
202207 for (ArrowBlock rbBlock : recordBatches ) {
208+ Assert .assertTrue (rbBlock .getOffset () + " > " + previousOffset , rbBlock .getOffset () > previousOffset );
209+ previousOffset = rbBlock .getOffset ();
203210 Assert .assertEquals (0 , rbBlock .getOffset () % 8 );
204211 Assert .assertEquals (0 , rbBlock .getMetadataLength () % 8 );
205212 try (ArrowRecordBatch recordBatch = arrowReader .readRecordBatch (rbBlock )) {
213+ Assert .assertEquals ("RB #" + i , counts [i ], recordBatch .getLength ());
206214 List <ArrowBuffer > buffersLayout = recordBatch .getBuffersLayout ();
207215 for (ArrowBuffer arrowBuffer : buffersLayout ) {
208216 Assert .assertEquals (0 , arrowBuffer .getOffset () % 8 );
209217 }
210218 vectorLoader .load (recordBatch );
211- validateContent (count , root );
219+ validateContent (counts [ i ] , root );
212220 }
221+ ++i ;
213222 }
214223 }
215224 }
0 commit comments