Rebuild JSON index when JsonIndexConfig changes; fix equals/hashCode to include all fields

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/JsonIndexHandler.java

@@ -23,6 +23,8 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.commons.configuration2.PropertiesConfiguration;
 import org.apache.commons.io.FileUtils;
 import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType;
 import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler;
@@ -40,10 +42,14 @@
 import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader;
 import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext;
 import org.apache.pinot.segment.spi.store.SegmentDirectory;
+import org.apache.pinot.segment.spi.utils.SegmentMetadataUtils;
 import org.apache.pinot.spi.config.table.JsonIndexConfig;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.CommonsConfigurationUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.MapUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -65,12 +71,43 @@ public boolean needUpdateIndices(SegmentDirectory.Reader segmentReader) {
     String segmentName = _segmentDirectory.getSegmentMetadata().getName();
     Set<String> columnsToAddIdx = new HashSet<>(_jsonIndexConfigs.keySet());
     Set<String> existingColumns = segmentReader.toSegmentDirectory().getColumnsWithIndex(StandardIndexes.json());
-    // Check if any existing index need to be removed.
+    // Load metadata properties once to avoid repeated disk reads when checking multiple columns.
+    PropertiesConfiguration properties = loadMetadataProperties();
+    // Check if any existing index need to be removed or if the config changed.
     for (String column : existingColumns) {
       if (!columnsToAddIdx.remove(column)) {
         LOGGER.info("Need to remove existing json index from segment: {}, column: {}", segmentName, column);
         return true;
       }
+      // Column exists in both existing indexes and config; check if config changed.
+      // When properties != null but storedConfig is null the index predates config persistence (legacy segment):
+      // only backfill (rebuild) when the forward index is present or reconstructable. If the forward index is
+      // absent (e.g. forwardIndexDisabled=true with no dictionary/inverted pair), preserve the existing JSON
+      // index untouched — createForwardIndexIfNeeded() would fail with no dictionary/inverted index available.
+      JsonIndexConfig currentConfig = _jsonIndexConfigs.get(column);
+      JsonIndexConfig storedConfig = readStoredJsonIndexConfig(column, properties);
+      if (properties != null) {
+        // configChanged: stored config exists and differs from current — rebuild regardless of forward
+        // index presence (same pre-existing behaviour; createForwardIndexIfNeeded will fail if the
+        // forward index is truly unrecoverable, which requires an explicit segment refresh).
+        boolean configChanged = storedConfig != null && !storedConfig.equals(currentConfig);
+        // legacyBackfill: stored config is absent (either the index predates config persistence, or
+        // the stored value failed to deserialize and readStoredJsonIndexConfig returned null). Only
+        // backfill when the forward index is present or reconstructable; if it is absent
+        // (forwardIndexDisabled=true with no dictionary/inverted pair) createForwardIndexIfNeeded()
+        // would throw, so preserve the existing JSON index instead.
+        boolean legacyBackfill = storedConfig == null && isForwardIndexAvailable(segmentReader, column);
+        if (configChanged) {
+          LOGGER.info("Need to rebuild json index for segment: {}, column: {} due to config change", segmentName,
+              column);
+          return true;
+        }
+        if (legacyBackfill) {
+          LOGGER.info("Need to rebuild json index for segment: {}, column: {} to backfill stored config "
+              + "(legacy segment or unreadable stored config)", segmentName, column);
+          return true;
+        }
+      }
     }
     // Check if any new index need to be added.
     for (String column : columnsToAddIdx) {
@@ -90,21 +127,118 @@ public void updateIndices(SegmentDirectory.Writer segmentWriter)
     String segmentName = _segmentDirectory.getSegmentMetadata().getName();
     Set<String> columnsToAddIdx = new HashSet<>(_jsonIndexConfigs.keySet());
     Set<String> existingColumns = segmentWriter.toSegmentDirectory().getColumnsWithIndex(StandardIndexes.json());
+    // Load metadata properties once; track whether we need to save changes.
+    PropertiesConfiguration properties = loadMetadataProperties();
+    boolean metadataModified = false;
     for (String column : existingColumns) {
       if (!columnsToAddIdx.remove(column)) {
         LOGGER.info("Removing existing json index from segment: {}, column: {}", segmentName, column);
         segmentWriter.removeIndex(column, StandardIndexes.json());
         LOGGER.info("Removed existing json index from segment: {}, column: {}", segmentName, column);
+      } else {
+        // Column exists in both existing indexes and config; check if config changed and rebuild if needed.
+        // When properties != null but storedConfig is null the index predates config persistence (legacy segment):
+        // only backfill (rebuild) when the forward index is present or reconstructable. If the forward index is
+        // absent (e.g. forwardIndexDisabled=true with no dictionary/inverted pair), preserve the existing JSON
+        // index untouched — createForwardIndexIfNeeded() would fail with no dictionary/inverted index available.
+        JsonIndexConfig currentConfig = _jsonIndexConfigs.get(column);
+        JsonIndexConfig storedConfig = readStoredJsonIndexConfig(column, properties);
+        if (properties != null) {
+          // configChanged: stored config exists and differs — rebuild (see needUpdateIndices for rationale).
+          boolean configChanged = storedConfig != null && !storedConfig.equals(currentConfig);
+          // legacyBackfill: absent/unreadable stored config; only rebuild when forward index is recoverable.
+          boolean legacyBackfill = storedConfig == null && isForwardIndexAvailable(segmentWriter, column);
+          if (configChanged) {
+            LOGGER.info("Rebuilding json index for segment: {}, column: {} due to config change", segmentName, column);
+            segmentWriter.removeIndex(column, StandardIndexes.json());
+            columnsToAddIdx.add(column);
+          } else if (legacyBackfill) {
+            LOGGER.info("Rebuilding json index for segment: {}, column: {} to backfill stored config "
+                + "(legacy segment or unreadable stored config)", segmentName, column);
+            segmentWriter.removeIndex(column, StandardIndexes.json());
+            columnsToAddIdx.add(column);
+          }
+        }
       }
     }
     for (String column : columnsToAddIdx) {
       ColumnMetadata columnMetadata = _segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column);
       if (shouldCreateJsonIndex(columnMetadata)) {
         createJsonIndexForColumn(segmentWriter, columnMetadata);
+        if (properties != null) {
+          setStoredJsonIndexConfig(column, _jsonIndexConfigs.get(column), properties);
+          metadataModified = true;
+        }
+      }
+    }
+    if (metadataModified && properties != null) {
+      SegmentMetadataUtils.savePropertiesConfiguration(properties,
+          _segmentDirectory.getSegmentMetadata().getIndexDir());
+    }
+  }
+
+  /// Loads the segment {@code metadata.properties} file, or returns {@code null} if unavailable (in-memory segments).
+  @Nullable
+  private PropertiesConfiguration loadMetadataProperties() {
+    try {
+      return SegmentMetadataUtils.getPropertiesConfiguration(_segmentDirectory.getSegmentMetadata());
+    } catch (Exception e) {
+      LOGGER.warn("Cannot load segment metadata properties for segment: {}, config-change detection disabled",
+          _segmentDirectory.getSegmentMetadata().getName(), e);
+      return null;
+    }
+  }
+
+  /// Reads the stored {@link JsonIndexConfig} from the given {@code properties}, or returns {@code null} if not
+  /// present (e.g., index was built before config persistence was added).
+  @Nullable
+  private static JsonIndexConfig readStoredJsonIndexConfig(String columnName,
+      @Nullable PropertiesConfiguration properties) {
+    if (properties == null) {
+      return null;
+    }
+    try {
+      String key = V1Constants.MetadataKeys.Column.getKeyFor(columnName, "jsonIndexConfig");
+      String escaped = properties.getString(key, null);
+      if (escaped == null) {
+        return null;
+      }
+      String serialized = CommonsConfigurationUtils.recoverSpecialCharacterInPropertyValue(escaped);
+      return JsonUtils.stringToObject(serialized, JsonIndexConfig.class);
+    } catch (Exception e) {
+      LOGGER.warn("Failed to read stored json index config for column: {}, treating as legacy (no stored config)",
+          columnName, e);
+      return null;
+    }
+  }
+
+  /// Writes the {@link JsonIndexConfig} for {@code columnName} into {@code properties} (without saving to disk).
+  private static void setStoredJsonIndexConfig(String columnName, JsonIndexConfig config,
+      PropertiesConfiguration properties) {
+    try {
+      String serialized = JsonUtils.objectToString(config);
+      String escaped = CommonsConfigurationUtils.replaceSpecialCharacterInPropertyValue(serialized);
+      if (escaped == null) {
+        LOGGER.warn("Cannot persist json index config for column: {} due to unsupported characters", columnName);
+        return;
       }
+      String key = V1Constants.MetadataKeys.Column.getKeyFor(columnName, "jsonIndexConfig");
+      properties.setProperty(key, escaped);
+    } catch (Exception e) {
+      LOGGER.warn("Failed to serialize json index config for column: {}, config changes will not trigger rebuild",
+          columnName, e);
     }
   }
 
+  /// Returns {@code true} if the forward index for {@code column} is already present in the segment or can be
+  /// reconstructed from the dictionary and inverted index (the path used by
+  /// {@link BaseIndexHandler#createForwardIndexIfNeeded}).
+  private static boolean isForwardIndexAvailable(SegmentDirectory.Reader segmentReader, String column) {
+    return segmentReader.hasIndexFor(column, StandardIndexes.forward())
+        || (segmentReader.hasIndexFor(column, StandardIndexes.dictionary())
+            && segmentReader.hasIndexFor(column, StandardIndexes.inverted()));
+  }
+
   private boolean shouldCreateJsonIndex(ColumnMetadata columnMetadata) {
     return columnMetadata != null;
   }
@@ -186,8 +320,13 @@ private void handleNonDictionaryBasedColumn(SegmentDirectory.Writer segmentWrite
         ForwardIndexReaderContext readerContext = forwardIndexReader.createContext();
         JsonIndexCreator jsonIndexCreator = StandardIndexes.json().createIndexCreator(context, config)) {
       int numDocs = columnMetadata.getTotalDocs();
+      boolean isMapType = columnMetadata.getDataType() == DataType.MAP;
       for (int i = 0; i < numDocs; i++) {
-        jsonIndexCreator.add(forwardIndexReader.getString(i, readerContext));
+        // MAP columns store binary-encoded data; getString() returns garbled bytes.
+        // Use getMap() + MapUtils.toString() to produce the JSON string the index creator expects.
+        String value = isMapType ? MapUtils.toString(forwardIndexReader.getMap(i, readerContext))
+            : forwardIndexReader.getString(i, readerContext);
+        jsonIndexCreator.add(value);
       }
       jsonIndexCreator.seal();
     }