@@ -367,36 +367,56 @@ private static SaveBlockResult saveStripeBlocks(
367367 final var sparseVals = blocks .sparseVals ();
368368 final var denseFlats = blocks .denseFlats ();
369369
370+ long stripeSparseLen = 0L ;
371+ for (int i = 0 ; i < colStripes .length ; i ++) {
372+ if (!denseFlags [i ]) {
373+ stripeSparseLen += lengths [i ];
374+ }
375+ }
376+
377+ final long [] offsetRow = new long [stripeCount ];
378+ final long [] lengthRow = new long [stripeCount ];
379+ final long [] stripeRows = stripeSparseLen > 0 ? new long [(int ) stripeSparseLen ] : new long [0 ];
380+ final long [] stripeCols = stripeSparseLen > 0 ? new long [(int ) stripeSparseLen ] : new long [0 ];
381+ final long [] stripeVals = stripeSparseLen > 0 ? new long [(int ) stripeSparseLen ] : new long [0 ];
382+ int stripePos = 0 ;
383+
370384 for (int i = 0 ; i < colStripes .length ; i ++) {
371385 final int blockLen = lengths [i ];
372386 if (blockLen <= 0 ) {
373387 continue ;
374388 }
375- final long blockIndex = ( long ) rowStripe * stripeCount + colStripes [i ];
389+ final int colStripe = colStripes [i ];
376390
377391 if (denseFlags [i ]) {
378- dst .int64 ().writeArrayBlockWithOffset (blockOffsetPath , new long []{-denseOffset - 1L }, 1 , blockIndex );
379- dst .int64 ().writeArrayBlockWithOffset (blockLengthPath , new long []{blockLen }, 1 , blockIndex );
380-
392+ offsetRow [colStripe ] = -denseOffset - 1L ;
393+ lengthRow [colStripe ] = blockLen ;
381394 final var denseMd = new MDLongArray (denseFlats [i ], new int []{1 , 1 , SUBMATRIX_SIZE , SUBMATRIX_SIZE });
382395 dst .int64 ().writeMDArrayBlockWithOffset (denseBlocksPath , denseMd , new long []{denseOffset , 0L , 0L , 0L });
383396 denseOffset ++;
384397 } else {
385- dst .int64 ().writeArrayBlockWithOffset (blockOffsetPath , new long []{sparseOffset }, 1 , blockIndex );
386- dst .int64 ().writeArrayBlockWithOffset (blockLengthPath , new long []{blockLen }, 1 , blockIndex );
387-
398+ offsetRow [colStripe ] = sparseOffset ;
399+ lengthRow [colStripe ] = blockLen ;
388400 final var blockRows = sparseRows [i ];
389401 final var blockCols = sparseCols [i ];
390402 final var blockVals = sparseVals [i ];
391-
392- dst .int64 ().writeArrayBlockWithOffset (blockRowsPath , blockRows , blockLen , sparseOffset );
393- dst .int64 ().writeArrayBlockWithOffset (blockColsPath , blockCols , blockLen , sparseOffset );
394- dst .int64 ().writeArrayBlockWithOffset (blockValsPath , blockVals , blockLen , sparseOffset );
395-
403+ System .arraycopy (blockRows , 0 , stripeRows , stripePos , blockLen );
404+ System .arraycopy (blockCols , 0 , stripeCols , stripePos , blockLen );
405+ System .arraycopy (blockVals , 0 , stripeVals , stripePos , blockLen );
406+ stripePos += blockLen ;
396407 sparseOffset += blockLen ;
397408 }
398409 }
399410
411+ final long rowBase = (long ) rowStripe * stripeCount ;
412+ dst .int64 ().writeArrayBlockWithOffset (blockOffsetPath , offsetRow , stripeCount , rowBase );
413+ dst .int64 ().writeArrayBlockWithOffset (blockLengthPath , lengthRow , stripeCount , rowBase );
414+ if (stripeSparseLen > 0 ) {
415+ dst .int64 ().writeArrayBlockWithOffset (blockRowsPath , stripeRows , (int ) stripeSparseLen , stripeSparseOffset );
416+ dst .int64 ().writeArrayBlockWithOffset (blockColsPath , stripeCols , (int ) stripeSparseLen , stripeSparseOffset );
417+ dst .int64 ().writeArrayBlockWithOffset (blockValsPath , stripeVals , (int ) stripeSparseLen , stripeSparseOffset );
418+ }
419+
400420 return new SaveBlockResult (sparseOffset , denseOffset );
401421 }
402422
0 commit comments