estimateArchiveBatchMemoryUsage estimate the GPU memory usage for archive batch
(batch *memstore.ArchiveBatch, isFirstOrLast bool)
| 1056 | |
| 1057 | // estimateArchiveBatchMemoryUsage estimate the GPU memory usage for archive batch |
| 1058 | func (qc *AQLQueryContext) estimateArchiveBatchMemoryUsage(batch *memstore.ArchiveBatch, isFirstOrLast bool) int { |
| 1059 | if batch == nil { |
| 1060 | return 0 |
| 1061 | } |
| 1062 | |
| 1063 | columnMemUsage := 0 |
| 1064 | var firstColumnSize int |
| 1065 | startRow, endRow := 0, batch.Size |
| 1066 | var hostSlice memCom.HostVectorPartySlice |
| 1067 | |
| 1068 | matchedColumnUsages := columnUsedByAllBatches |
| 1069 | if isFirstOrLast { |
| 1070 | matchedColumnUsages |= columnUsedByFirstArchiveBatch | columnUsedByLastArchiveBatch |
| 1071 | } |
| 1072 | |
| 1073 | prefilterIndex := 0 |
| 1074 | // max number of rows after pre-filtering. used for non-agg query |
| 1075 | maxSizeAfterPreFilter := batch.Size |
| 1076 | for i := len(qc.TableScanners[0].Columns) - 1; i >= 0; i-- { |
| 1077 | columnID := qc.TableScanners[0].Columns[i] |
| 1078 | usage := qc.TableScanners[0].ColumnUsages[columnID] |
| 1079 | // TODO(cdavid): only read metadata when estimate query memory requirement. |
| 1080 | sourceVP := batch.RequestVectorParty(columnID) |
| 1081 | sourceVP.WaitForDiskLoad() |
| 1082 | |
| 1083 | if usage&matchedColumnUsages != 0 || usage&columnUsedByPrefilter != 0 { |
| 1084 | startRow, endRow, hostSlice = qc.prefilterSlice(sourceVP, prefilterIndex, startRow, endRow) |
| 1085 | if endRow-startRow < maxSizeAfterPreFilter { |
| 1086 | maxSizeAfterPreFilter = endRow - startRow |
| 1087 | } |
| 1088 | prefilterIndex++ |
| 1089 | if usage&matchedColumnUsages != 0 { |
| 1090 | columnMemUsage += hostSlice.ValueBytes + hostSlice.NullBytes + hostSlice.CountBytes |
| 1091 | firstColumnSize = hostSlice.Length |
| 1092 | } |
| 1093 | } |
| 1094 | sourceVP.Release() |
| 1095 | } |
| 1096 | if maxSizeAfterPreFilter > qc.maxBatchSizeAfterPrefilter { |
| 1097 | qc.maxBatchSizeAfterPrefilter = maxSizeAfterPreFilter |
| 1098 | } |
| 1099 | |
| 1100 | totalBytes := qc.estimateMemUsageForBatch(firstColumnSize, columnMemUsage, maxSizeAfterPreFilter) |
| 1101 | utils.GetQueryLogger().Debugf("Archive batch %d needs memory: %d", batch.BatchID, totalBytes) |
| 1102 | return totalBytes |
| 1103 | } |
| 1104 | |
| 1105 | // estimateMemUsageForBatch calculates memory usage including: |
| 1106 | // * Index vector |
no test coverage detected