Fix segment close logic

Kartik Khare · Kartik Khare · commit d157124f437d · 2022-11-30T21:44:16.000+05:30
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java
@@ -81,6 +81,7 @@ public static DimensionTableDataManager getInstanceByTableName(String tableNameW
 
   private volatile DimensionTable _dimensionTable;
   private boolean _disablePreload = false;
+  private List<SegmentDataManager> _segmentDataManagers;
 
   @Override
   protected void doInit() {
@@ -101,9 +102,9 @@ protected void doInit() {
     }
 
     if (_disablePreload) {
-      _dimensionTable = new MemoryOptimizedDimensionTable(schema, primaryKeyColumns);
+      _dimensionTable = new MemoryOptimizedDimensionTable(schema, primaryKeyColumns, new HashMap<>());
     } else {
-      _dimensionTable = new FastLookupDimensionTable(schema, primaryKeyColumns);
+      _dimensionTable = new FastLookupDimensionTable(schema, primaryKeyColumns, new HashMap<>());
     }
   }
 
@@ -136,16 +137,15 @@ public void removeSegment(String segmentName) {
 
   @Override
   protected void doShutdown() {
-    closeDimensionTable(_dimensionTable, true);
+    closeDimensionTable(_dimensionTable, _disablePreload);
   }
 
   private void closeDimensionTable(DimensionTable dimensionTable, boolean releaseSegments) {
     if (dimensionTable != null) {
       try {
         dimensionTable.close();
-        if (releaseSegments) {
-          List<SegmentDataManager> segmentDataManagers = acquireAllSegments();
-          for (SegmentDataManager segmentDataManager: segmentDataManagers) {
+        if (releaseSegments && _segmentDataManagers != null) {
+          for (SegmentDataManager segmentDataManager: _segmentDataManagers) {
             releaseSegment(segmentDataManager);
           }
         }
@@ -182,9 +182,9 @@ private DimensionTable createFastLookupDimensionTable() {
         "Primary key columns must be configured for dimension table: %s", _tableNameWithType);
 
     Map<PrimaryKey, GenericRow> lookupTable = new HashMap<>();
-    List<SegmentDataManager> segmentManagers = acquireAllSegments();
+    List<SegmentDataManager> segmentDataManagers = acquireAllSegments();
     try {
-      for (SegmentDataManager segmentManager : segmentManagers) {
+      for (SegmentDataManager segmentManager : segmentDataManagers) {
         IndexSegment indexSegment = segmentManager.getSegment();
         int numTotalDocs = indexSegment.getSegmentMetadata().getTotalDocs();
         if (numTotalDocs > 0) {
@@ -203,7 +203,7 @@ private DimensionTable createFastLookupDimensionTable() {
       }
       return new FastLookupDimensionTable(schema, primaryKeyColumns, lookupTable);
     } finally {
-      for (SegmentDataManager segmentManager : segmentManagers) {
+      for (SegmentDataManager segmentManager : segmentDataManagers) {
         releaseSegment(segmentManager);
       }
     }
@@ -218,8 +218,8 @@ private DimensionTable createMemOptimisedDimensionTable() {
         "Primary key columns must be configured for dimension table: %s", _tableNameWithType);
 
     Map<PrimaryKey, LookupRecordLocation> lookupTable = new HashMap<>();
-    List<SegmentDataManager> segmentManagers = acquireAllSegments();
-    for (SegmentDataManager segmentManager : segmentManagers) {
+    _segmentDataManagers = acquireAllSegments();
+    for (SegmentDataManager segmentManager : _segmentDataManagers) {
       IndexSegment indexSegment = segmentManager.getSegment();
       int numTotalDocs = indexSegment.getSegmentMetadata().getTotalDocs();
       if (numTotalDocs > 0) {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/FastLookupDimensionTable.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/FastLookupDimensionTable.java
@@ -19,7 +19,6 @@
 package org.apache.pinot.core.data.manager.offline;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.pinot.spi.data.FieldSpec;
@@ -34,10 +33,6 @@ class FastLookupDimensionTable implements DimensionTable {
   private final Schema _tableSchema;
   private final List<String> _primaryKeyColumns;
 
-  FastLookupDimensionTable(Schema tableSchema, List<String> primaryKeyColumns) {
-    this(tableSchema, primaryKeyColumns, new HashMap<>());
-  }
-
   FastLookupDimensionTable(Schema tableSchema, List<String> primaryKeyColumns,
       Map<PrimaryKey, GenericRow> lookupTable) {
     _lookupTable = lookupTable;
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java
@@ -19,7 +19,6 @@
 package org.apache.pinot.core.data.manager.offline;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.pinot.core.data.manager.BaseTableDataManager;
@@ -39,10 +38,6 @@ class MemoryOptimizedDimensionTable implements DimensionTable {
   private final List<String> _primaryKeyColumns;
   private final GenericRow _reuse = new GenericRow();
 
-  MemoryOptimizedDimensionTable(Schema tableSchema, List<String> primaryKeyColumns) {
-    this(tableSchema, primaryKeyColumns, new HashMap<>());
-  }
-
   MemoryOptimizedDimensionTable(Schema tableSchema, List<String> primaryKeyColumns,
       Map<PrimaryKey, LookupRecordLocation> lookupTable) {
     _lookupTable = lookupTable;
