@@ -892,24 +892,33 @@ std::shared_ptr<arrow::RecordBatch> MergeColumns(const std::vector<std::shared_p
892892}
893893
894894std::vector<std::shared_ptr<arrow::RecordBatch>> SliceToRecordBatches (const std::shared_ptr<arrow::Table>& t) {
895- std::set<ui32> splitPositions;
896- const ui32 numRows = t->num_rows ();
897- for (auto && i : t->columns ()) {
898- ui32 pos = 0 ;
899- for (auto && arr : i->chunks ()) {
900- splitPositions.emplace (pos);
901- pos += arr->length ();
895+ if (!t->num_rows ()) {
896+ return {};
897+ }
898+ std::vector<ui32> positions;
899+ {
900+ for (auto && i : t->columns ()) {
901+ ui32 pos = 0 ;
902+ for (auto && arr : i->chunks ()) {
903+ positions.emplace_back (pos);
904+ pos += arr->length ();
905+ }
906+ AFL_VERIFY (pos == t->num_rows ());
902907 }
903- AFL_VERIFY (pos == t->num_rows ());
908+ positions. emplace_back ( t->num_rows ());
904909 }
910+ std::sort (positions.begin (), positions.end ());
911+ positions.erase (std::unique (positions.begin (), positions.end ()), positions.end ());
912+
905913 std::vector<std::vector<std::shared_ptr<arrow::Array>>> slicedData;
906- slicedData.resize (splitPositions.size ());
907- std::vector<ui32> positions (splitPositions.begin (), splitPositions.end ());
908- for (auto && i : t->columns ()) {
909- for (ui32 idx = 0 ; idx < positions.size (); ++idx) {
910- auto slice = i->Slice (positions[idx], ((idx + 1 == positions.size ()) ? numRows : positions[idx + 1 ]) - positions[idx]);
911- AFL_VERIFY (slice->num_chunks () == 1 );
912- slicedData[idx].emplace_back (slice->chunks ().front ());
914+ slicedData.resize (positions.size () - 1 );
915+ {
916+ for (auto && i : t->columns ()) {
917+ for (ui32 idx = 0 ; idx + 1 < positions.size (); ++idx) {
918+ auto slice = i->Slice (positions[idx], positions[idx + 1 ] - positions[idx]);
919+ AFL_VERIFY (slice->num_chunks () == 1 );
920+ slicedData[idx].emplace_back (slice->chunks ().front ());
921+ }
913922 }
914923 }
915924 std::vector<std::shared_ptr<arrow::RecordBatch>> result;
0 commit comments