This is an automated email from the ASF dual-hosted git repository.

xiangfu0 pushed a commit to branch xiangfu0/data-3221-9-column-arrays
in repository https://gitbox.apache.org/repos/asf/pinot.git

commit a2313efa6a3e1af0d382253fdb7a406052ee6960
Author: Xiang Fu <[email protected]>
AuthorDate: Sat Sep 5 22:03:28 2026 -0700

    DATA-3221 (10): publish the segment column arrays as one immutable pair
    
    Review follow-up on the sorted-array column metadata store.
    
    - The names and the metadata were two independent volatile fields replaced 
one
      after the other, so a reader could see the names of one version beside the
      metadata of another. Both now live in one immutable Columns holder behind 
a
      single volatile field, so every publication is atomic, and init() fills 
the
      arrays before publishing them instead of publishing them empty and filling
      them afterwards (the index_map loop looks its column up in the local 
array).
    - addColumnMetadata() replaced the metadata of an existing column in place,
      which a Collection returned earlier by getAllColumnMetadata() -- 
documented as
      a snapshot -- observed. It is copy-on-write now, like the insert branch.
    - removeColumn() on a CONSUMING segment fell through to 
invalidateDerivedViews()
      and nulled the caller-supplied schema, leaving a metadata with neither a
      schema nor column metadata (the next getSchema() then NPE'd). It rejects 
that
      segment, as addColumnMetadata() already did, and returns early -- without
      dropping the derived views -- when the column is not there.
    - The two mutators and the two derived-view builders now share the instance
      monitor, so a schema or map built from columns that have already been 
replaced
      can no longer be cached.
    - SortedStringArraySet's range views did not enforce their own bounds, so
      headSet("f").tailSet("h") returned an empty set where a TreeSet range view
      throws. They carry their bounds and reject an out-of-range argument the 
way
      TreeSet does.
    - Documented that a segment holding no column metadata (CONSUMING) reports 
its
      schema's columns from getAllColumns()/getNumColumns() while
      getAllColumnMetadata() is empty and forEachColumn() visits nothing, so 
the two
      families must not be paired; PhysicalColumnNames was pairing them, and now
      counts its size from the same column metadata it iterates.
    
    Heap: no new per-column term. Per segment the Columns holder costs ~20 B 
(one
    object of 24 B replacing one of the two reference fields) and the retained
    SortedStringArraySet grows 32 -> 40 B for its bounds, i.e. ~+28 B/segment, 
which
    at 1000 columns is +0.03 B/column against the measured 166 B/column.
    
    Co-Authored-By: Claude Opus 5 <[email protected]>
---
 .../immutable/PhysicalColumnNames.java             |  17 +--
 .../segment/index/SegmentMetadataImplTest.java     |  55 +++++++-
 .../apache/pinot/segment/spi/SegmentMetadata.java  |  18 ++-
 .../spi/index/metadata/SegmentMetadataImpl.java    | 147 +++++++++++++--------
 .../spi/index/metadata/SortedStringArraySet.java   |  59 ++++++++-
 .../index/metadata/SortedStringArraySetTest.java   |  47 +++++++
 6 files changed, 271 insertions(+), 72 deletions(-)

diff --git 
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java
 
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java
index 3ca7967ecbc..9dd977a9d90 100644
--- 
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java
+++ 
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/immutable/PhysicalColumnNames.java
@@ -30,23 +30,24 @@ import org.apache.pinot.segment.spi.SegmentMetadata;
 ///
 /// It is a view rather than a copy so a segment retains nothing per column 
for it: the segment schema this replaces
 /// held a `TreeMap` entry per column, and a cached `TreeSet` would hold the 
same. `contains` is one column lookup and
-/// iteration is a filtered pass over the column metadata. The virtual column 
count is taken once at construction,
-/// which is sound because the column metadata is fixed once the segment is 
loaded.
+/// iteration is a filtered pass over the column metadata. The size is counted 
once at construction, from the same
+/// column metadata the iteration reads, which is sound because the column 
metadata is fixed once the segment is
+/// loaded (the loader registers the virtual columns before the segment is 
built).
 ///
 /// Thread-safe for reads, like the underlying segment metadata once loaded.
 final class PhysicalColumnNames extends AbstractSet<String> {
   private final SegmentMetadata _segmentMetadata;
-  private final int _numVirtualColumns;
+  private final int _numPhysicalColumns;
 
   PhysicalColumnNames(SegmentMetadata segmentMetadata) {
     _segmentMetadata = segmentMetadata;
-    int numVirtualColumns = 0;
+    int numPhysicalColumns = 0;
     for (ColumnMetadata columnMetadata : 
segmentMetadata.getAllColumnMetadata()) {
-      if (!isPhysical(columnMetadata)) {
-        numVirtualColumns++;
+      if (isPhysical(columnMetadata)) {
+        numPhysicalColumns++;
       }
     }
-    _numVirtualColumns = numVirtualColumns;
+    _numPhysicalColumns = numPhysicalColumns;
   }
 
   private static boolean isPhysical(ColumnMetadata columnMetadata) {
@@ -64,7 +65,7 @@ final class PhysicalColumnNames extends AbstractSet<String> {
 
   @Override
   public int size() {
-    return _segmentMetadata.getNumColumns() - _numVirtualColumns;
+    return _numPhysicalColumns;
   }
 
   @Override
diff --git 
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java
 
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java
index d6f922fdd23..31ede949bbe 100644
--- 
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java
+++ 
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/SegmentMetadataImplTest.java
@@ -24,6 +24,7 @@ import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
@@ -80,6 +81,7 @@ import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertSame;
 import static org.testng.Assert.assertThrows;
 import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
 
 
 public class SegmentMetadataImplTest {
@@ -447,7 +449,7 @@ public class SegmentMetadataImplTest {
     assertEquals(metadata.getAllColumns().first(), column, "must be inserted 
in natural order, not appended");
     assertTrue(metadata.getSchema().hasColumn(column));
 
-    // Re-registering replaces in place rather than duplicating the column
+    // Re-registering replaces the column rather than duplicating it
     ColumnMetadata replacement =
         new EmptyColumnMetadata(new DimensionFieldSpec(column, 
FieldSpec.DataType.LONG, true), null, null);
     metadata.addColumnMetadata(column, replacement);
@@ -455,6 +457,57 @@ public class SegmentMetadataImplTest {
     assertSame(metadata.getColumnMetadataFor(column), replacement);
   }
 
+  /// The column arrays are replaced as a whole, never written in place, so a 
collection handed out earlier stays the
+  /// snapshot it is documented to be — through an insertion and through a 
replacement of a column already there —
+  /// and every name stays paired with its own metadata.
+  @Test
+  public void testColumnMetadataViewIsASnapshot()
+      throws Exception {
+    SegmentMetadataImpl metadata = new SegmentMetadataImpl(_segmentDirectory);
+    int numColumns = metadata.getNumColumns();
+    Collection<ColumnMetadata> snapshot = metadata.getAllColumnMetadata();
+    String column = metadata.getAllColumns().first();
+    ColumnMetadata original = metadata.getColumnMetadataFor(column);
+    assertSame(new ArrayList<>(snapshot).get(0), original);
+
+    metadata.addColumnMetadata(column,
+        new EmptyColumnMetadata(new DimensionFieldSpec(column, 
FieldSpec.DataType.INT, true), null, null));
+    metadata.addColumnMetadata("$aVirtualColumn",
+        new EmptyColumnMetadata(new DimensionFieldSpec("$aVirtualColumn", 
FieldSpec.DataType.INT, true), null, null));
+    assertEquals(snapshot.size(), numColumns);
+    assertSame(new ArrayList<>(snapshot).get(0), original, "a replacement must 
not reach the earlier snapshot");
+    assertNotSame(metadata.getColumnMetadataFor(column), original);
+
+    assertEquals(metadata.getNumColumns(), numColumns + 1);
+    assertEquals(metadata.getAllColumnMetadata().size(), numColumns + 1);
+    metadata.forEachColumn((name, columnMetadata) -> 
assertEquals(columnMetadata.getColumnName(), name));
+  }
+
+  /// A CONSUMING segment holds no column metadata: its column names come from 
the explicit schema, the column
+  /// metadata accessors are empty, and both mutators reject it rather than 
drop the schema it was given.
+  @Test
+  public void testConsumingSegmentHoldsNoColumnMetadata() {
+    Schema schema = new Schema.SchemaBuilder().setSchemaName("consuming")
+        .addSingleValueDimension("dim", FieldSpec.DataType.STRING)
+        .addMetric("metric", FieldSpec.DataType.LONG)
+        .build();
+    SegmentMetadataImpl metadata =
+        new SegmentMetadataImpl("testTable", 
"testTable__0__0__20240101T0000Z", schema, 123L);
+    assertEquals(metadata.getAllColumns(), schema.getColumnNames());
+    assertEquals(metadata.getNumColumns(), schema.size());
+    assertTrue(metadata.getAllColumnMetadata().isEmpty());
+    assertNull(metadata.getColumnMetadataMap());
+    assertNull(metadata.getColumnMetadataFor("dim"));
+    metadata.forEachColumn((column, columnMetadata) -> fail("no column 
metadata to visit, got: " + column));
+
+    ColumnMetadata added =
+        new EmptyColumnMetadata(new DimensionFieldSpec("added", 
FieldSpec.DataType.INT, true), null, null);
+    assertThrows(IllegalStateException.class, () -> 
metadata.addColumnMetadata("added", added));
+    assertThrows(IllegalStateException.class, () -> 
metadata.removeColumn("dim"));
+    assertSame(metadata.getSchema(), schema, "the explicit schema survives a 
rejected mutation");
+    assertEquals(metadata.getAllColumns(), schema.getColumnNames());
+  }
+
   /// The metadata JSON is a public REST payload: it must list the columns in 
the same natural order the map view
   /// does, filter included.
   @Test
diff --git 
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java
 
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java
index c01f2cf39e6..142bf688dca 100644
--- 
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java
+++ 
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/SegmentMetadata.java
@@ -122,18 +122,24 @@ public interface SegmentMetadata {
     return getSchema().getColumnNames();
   }
 
-  /// Number of columns the segment metadata holds.
+  /// Number of columns in [#getAllColumns()].
+  ///
+  /// A segment that holds no column metadata (a CONSUMING one, built from an 
explicit schema) still reports its
+  /// schema's columns here, so this is not the size of 
[#getAllColumnMetadata()]: do not pair the two.
   default int getNumColumns() {
     return getColumnMetadataMap().size();
   }
 
-  /// The column metadata of every column, in the natural column-name order of 
[#getAllColumns()].
+  /// The column metadata of every column that has some, in the natural 
column-name order of [#getAllColumns()], and
+  /// empty for a segment that holds none (a CONSUMING one, which answers 
[#getColumnMetadataFor(String)] with `null`
+  /// for every column of its schema).
   default Collection<ColumnMetadata> getAllColumnMetadata() {
     return getColumnMetadataMap().values();
   }
 
   /// Applies `action` to every (column name, column metadata) pair, in the 
natural column-name order of
-  /// [#getAllColumns()].
+  /// [#getAllColumns()], and to nothing at all for a segment that holds no 
column metadata, exactly as
+  /// [#getAllColumnMetadata()] is empty for one.
   default void forEachColumn(BiConsumer<String, ColumnMetadata> action) {
     getColumnMetadataMap().forEach(action);
   }
@@ -155,12 +161,14 @@ public interface SegmentMetadata {
     return getColumnMetadataMap().get(column);
   }
 
-  /// Registers the metadata of a column, replacing any metadata already 
registered under the same name.
+  /// Registers the metadata of a column, replacing any metadata already 
registered under the same name. An
+  /// implementation that holds no column metadata (a CONSUMING segment) may 
reject this.
   default void addColumnMetadata(String column, ColumnMetadata columnMetadata) 
{
     getColumnMetadataMap().put(column, columnMetadata);
   }
 
-  /// Removes a column from the segment metadata.
+  /// Removes a column from the segment metadata. An implementation that holds 
no column metadata (a CONSUMING
+  /// segment) may reject this.
   void removeColumn(String column);
 
   /// Converts segment metadata to json.
diff --git 
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java
 
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java
index 9cdb5b17775..9d8211d41d6 100644
--- 
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java
+++ 
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SegmentMetadataImpl.java
@@ -90,11 +90,16 @@ import org.slf4j.LoggerFactory;
 ///
 /// Once the loader has registered the built-in virtual columns through 
[#addColumnMetadata(String, ColumnMetadata)],
 /// the derived schema includes them, exactly as the eagerly built one did. 
[#removeColumn(String)] and
-/// [#addColumnMetadata(String, ColumnMetadata)] replace the arrays and drop 
both derived views. The explicit-schema
-/// constructor keeps the caller's Schema as is and holds no column metadata 
at all.
+/// [#addColumnMetadata(String, ColumnMetadata)] replace both arrays at once 
and drop both derived views. The
+/// explicit-schema constructor keeps the caller's Schema as is and holds no 
column metadata at all, so those two
+/// mutators reject such a metadata rather than drop the schema it was given: 
a CONSUMING segment answers
+/// [#getAllColumns()] and [#getNumColumns()] from that schema and reports no 
column metadata at all
+/// ([#getColumnMetadataFor(String)] `null`, [#getAllColumnMetadata()] empty, 
[#forEachColumn(BiConsumer)] a no-op,
+/// [#getColumnMetadataMap()] `null`).
 ///
-/// Thread-safe for the derived schema and map (double-checked on volatiles, 
so one instance per metadata); the
-/// columns themselves are populated at load before the metadata is published.
+/// Thread-safe: the two arrays are published together in one immutable 
holder, so no reader can see the names of
+/// one version beside the metadata of another, and the derived schema and map 
are built under the instance monitor
+/// the mutators hold as well, so neither can be cached from columns that have 
already been replaced.
 public class SegmentMetadataImpl implements SegmentMetadata {
   private static final Logger LOGGER = 
LoggerFactory.getLogger(SegmentMetadataImpl.class);
 
@@ -106,15 +111,13 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
   private static final AtomicLong NUM_COLUMN_METADATA_MAP_MATERIALIZATIONS = 
new AtomicLong();
 
   private final File _indexDir;
-  /// Column names in natural order, and their metadata at the same index. 
Both `null` for a CONSUMING segment, which
-  /// is constructed with an explicit schema and holds no column metadata. 
Replaced (never written in place) by
+  /// The columns of a metadata-backed segment, or `null` for a CONSUMING 
segment, which is constructed with an
+  /// explicit schema and holds no column metadata. Replaced as a whole (never 
written in place) by
   /// [#addColumnMetadata(String, ColumnMetadata)] and 
[#removeColumn(String)], so a view handed out earlier stays a
   /// consistent snapshot, and volatile so a metadata published without other 
synchronization is seen with its
   /// columns.
   @Nullable
-  private volatile String[] _columnNames;
-  @Nullable
-  private volatile ColumnMetadata[] _columnMetadata;
+  private volatile Columns _columns;
   /// The explicit schema of a CONSUMING segment, or the lazily derived schema 
of a metadata-backed segment (null
   /// until [#getSchema()] builds it, and again whenever the columns change).
   @Nullable
@@ -275,12 +278,11 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
 
     // Build the sorted column arrays (the map view and the schema are derived 
from them on demand, see
     // getColumnMetadataMap() and getSchema()). Empty segments use a 
stripped-down [EmptyColumnMetadata] since the
-    // shape stats (cardinality, element lengths, etc.) are meaningless when 
there are no rows.
+    // shape stats (cardinality, element lengths, etc.) are meaningless when 
there are no rows. Both arrays are
+    // filled before they are published below, so a reader never sees a 
half-built one.
     String[] columns = physicalColumns.toArray(new String[0]);
     Arrays.sort(columns);
     ColumnMetadata[] columnMetadata = new ColumnMetadata[columns.length];
-    _columnNames = columns;
-    _columnMetadata = columnMetadata;
     if (_totalDocs > 0) {
       for (int i = 0; i < columns.length; i++) {
         columnMetadata[i] = 
ColumnMetadataImpl.fromPropertiesConfiguration(segmentMetadata, _totalDocs, 
columns[i]);
@@ -301,8 +303,10 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
               String[] parsedKeys = ColumnIndexUtils.parseIndexMapKeys(key, 
_indexDir.getPath());
               if (parsedKeys[2].equals(ColumnIndexUtils.MAP_KEY_NAME_SIZE)) {
                 short indexType = indexService.getNumericId(parsedKeys[1]);
-                ((ColumnMetadataImpl) 
getColumnMetadataFor(parsedKeys[0])).addIndexSize(indexType,
-                    mapConfig.getLong(key));
+                // The arrays are not published yet, so this looks the column 
up in the local one
+                int index = Arrays.binarySearch(columns, parsedKeys[0]);
+                Preconditions.checkState(index >= 0, "Column: %s is not in the 
segment metadata", parsedKeys[0]);
+                ((ColumnMetadataImpl) 
columnMetadata[index]).addIndexSize(indexType, mapConfig.getLong(key));
               }
             } catch (Exception e) {
               LOGGER.debug("Unable to load index metadata in {} for {}!", 
indexMapFile, key, e);
@@ -315,6 +319,7 @@ public class SegmentMetadataImpl implements SegmentMetadata 
{
         columnMetadata[i] = 
EmptyColumnMetadata.fromPropertiesConfiguration(segmentMetadata, columns[i]);
       }
     }
+    _columns = new Columns(columns, columnMetadata);
 
     // Build star-tree v2 metadata
     int starTreeV2Count =
@@ -453,8 +458,11 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
 
   private Schema buildSchema() {
     NUM_SCHEMA_MATERIALIZATIONS.incrementAndGet();
+    // Only a metadata-backed segment gets here: a CONSUMING one is 
constructed with its schema, so getSchema()
+    // returns before building one
+    Columns columns = Preconditions.checkNotNull(_columns, "Segment: %s holds 
no column metadata", _segmentName);
     Schema schema = new Schema();
-    for (ColumnMetadata columnMetadata : _columnMetadata) {
+    for (ColumnMetadata columnMetadata : columns._metadata) {
       schema.addField(columnMetadata.getFieldSpec());
     }
     return schema;
@@ -479,48 +487,48 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
   /// metadata. The view is a snapshot: it does not reflect columns added or 
removed after this call.
   @Override
   public NavigableSet<String> getAllColumns() {
-    String[] columnNames = _columnNames;
-    return columnNames != null ? new SortedStringArraySet(columnNames) : 
getSchema().getColumnNames();
+    Columns columns = _columns;
+    return columns != null ? new SortedStringArraySet(columns._names) : 
getSchema().getColumnNames();
   }
 
   @Override
   public int getNumColumns() {
-    String[] columnNames = _columnNames;
-    return columnNames != null ? columnNames.length : getSchema().size();
+    Columns columns = _columns;
+    return columns != null ? columns._names.length : getSchema().size();
   }
 
   /// An unmodifiable view of the column metadata array, in the natural 
column-name order of [#getAllColumns()], and
-  /// empty for a CONSUMING segment. Like [#getAllColumns()] it is a snapshot.
+  /// empty for a CONSUMING segment, which holds no column metadata (see the 
class documentation). Like
+  /// [#getAllColumns()] it is a snapshot.
   @Override
   public Collection<ColumnMetadata> getAllColumnMetadata() {
-    ColumnMetadata[] columnMetadata = _columnMetadata;
-    return columnMetadata != null ? 
Collections.unmodifiableList(Arrays.asList(columnMetadata)) : List.of();
+    Columns columns = _columns;
+    return columns != null ? 
Collections.unmodifiableList(Arrays.asList(columns._metadata)) : List.of();
   }
 
+  /// Visits every column and its metadata in natural column-name order, and 
visits nothing for a CONSUMING segment,
+  /// which holds no column metadata (see the class documentation). The pair 
comes from one snapshot of the columns,
+  /// so a concurrent change cannot pair a name with another column's metadata.
   @Override
   public void forEachColumn(BiConsumer<String, ColumnMetadata> action) {
-    String[] columnNames = _columnNames;
-    if (columnNames == null) {
+    Columns columns = _columns;
+    if (columns == null) {
       return;
     }
-    ColumnMetadata[] columnMetadata = _columnMetadata;
-    for (int i = 0; i < columnNames.length; i++) {
-      action.accept(columnNames[i], columnMetadata[i]);
+    for (int i = 0; i < columns._names.length; i++) {
+      action.accept(columns._names[i], columns._metadata[i]);
     }
   }
 
   @Nullable
   @Override
   public ColumnMetadata getColumnMetadataFor(String column) {
-    int index = indexOf(column);
-    return index >= 0 ? _columnMetadata[index] : null;
-  }
-
-  /// Index of the column in the two arrays, or `-(insertion point) - 1`; 
always negative when there is no column
-  /// metadata at all (CONSUMING segment).
-  private int indexOf(String column) {
-    String[] columnNames = _columnNames;
-    return columnNames != null ? Arrays.binarySearch(columnNames, column) : -1;
+    Columns columns = _columns;
+    if (columns == null) {
+      return null;
+    }
+    int index = columns.indexOf(column);
+    return index >= 0 ? columns._metadata[index] : null;
   }
 
   @Override
@@ -626,7 +634,7 @@ public class SegmentMetadataImpl implements SegmentMetadata 
{
   @Nullable
   @Override
   public TreeMap<String, ColumnMetadata> getColumnMetadataMap() {
-    if (_columnNames == null) {
+    if (_columns == null) {
       return null;
     }
     TreeMap<String, ColumnMetadata> columnMetadataMap = _columnMetadataMapView;
@@ -664,35 +672,48 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
 
   /// {@inheritDoc}
   ///
-  /// Inserts the column in natural order, which is a copy of both arrays; the 
loader adds a handful of virtual
-  /// columns once per segment, so this is not a hot path.
+  /// Inserts the column in natural order, or replaces the metadata already 
registered under the name, either way by
+  /// copying both arrays: a view handed out earlier is documented as a 
snapshot, and the loader adds a handful of
+  /// virtual columns once per segment, so this is not a hot path.
+  ///
+  /// Throws for a CONSUMING segment, which was given an explicit schema and 
holds no column metadata to add to.
   @Override
-  public void addColumnMetadata(String column, ColumnMetadata columnMetadata) {
-    String[] columnNames = _columnNames;
-    Preconditions.checkState(columnNames != null, "Segment: %s holds no column 
metadata", _segmentName);
-    int index = Arrays.binarySearch(columnNames, column);
+  public synchronized void addColumnMetadata(String column, ColumnMetadata 
columnMetadata) {
+    Columns columns = _columns;
+    Preconditions.checkState(columns != null, "Segment: %s holds no column 
metadata", _segmentName);
+    int index = columns.indexOf(column);
     if (index >= 0) {
-      _columnMetadata[index] = columnMetadata;
+      ColumnMetadata[] metadata = columns._metadata.clone();
+      metadata[index] = columnMetadata;
+      _columns = new Columns(columns._names, metadata);
     } else {
       int insertionPoint = -index - 1;
-      _columnMetadata = insert(_columnMetadata, insertionPoint, 
columnMetadata);
-      _columnNames = insert(columnNames, insertionPoint, column);
+      _columns = new Columns(insert(columns._names, insertionPoint, column),
+          insert(columns._metadata, insertionPoint, columnMetadata));
     }
     invalidateDerivedViews();
   }
 
+  /// {@inheritDoc}
+  ///
+  /// Throws for a CONSUMING segment, which holds no column metadata: dropping 
its explicit schema instead would
+  /// leave it with neither.
   @Override
-  public void removeColumn(String column) {
+  public synchronized void removeColumn(String column) {
     Preconditions.checkState(!column.equals(_timeColumn), "Cannot remove time 
column: %s", _timeColumn);
-    int index = indexOf(column);
-    if (index >= 0) {
-      _columnMetadata = delete(_columnMetadata, index);
-      _columnNames = delete(_columnNames, index);
+    Columns columns = _columns;
+    Preconditions.checkState(columns != null, "Segment: %s holds no column 
metadata", _segmentName);
+    int index = columns.indexOf(column);
+    if (index < 0) {
+      return;
     }
+    _columns = new Columns(delete(columns._names, index), 
delete(columns._metadata, index));
     invalidateDerivedViews();
   }
 
-  /// Drops the schema and the map derived from the columns, so the next 
caller rebuilds them from the current arrays.
+  /// Drops the schema and the map derived from the columns, so the next 
caller rebuilds them from the current
+  /// arrays. Called while holding the instance monitor, which [#getSchema()] 
and [#getColumnMetadataMap()] also hold
+  /// while they build and cache, so a view derived from the replaced columns 
cannot survive this.
   private void invalidateDerivedViews() {
     _schema = null;
     _columnMetadataMapView = null;
@@ -756,7 +777,7 @@ public class SegmentMetadataImpl implements SegmentMetadata 
{
     segmentMetadata.put("startOffset", _startOffset);
     segmentMetadata.put("endOffset", _endOffset);
 
-    if (_columnNames != null) {
+    if (_columns != null) {
       ArrayNode columnsMetadata = JsonUtils.newArrayNode();
       forEachColumn((column, columnMetadata) -> {
         if (columnFilter == null || columnFilter.contains(column)) {
@@ -773,4 +794,24 @@ public class SegmentMetadataImpl implements 
SegmentMetadata {
   public String toString() {
     return toJson(null).toString();
   }
+
+  /// The columns of a metadata-backed segment: the names in natural order, 
and their metadata at the same index.
+  ///
+  /// The two arrays live in one immutable object so that every publication is 
atomic — a reader that sees a name
+  /// array never sees the metadata array of another version beside it — and 
so that the arrays a view was handed
+  /// stay exactly as they were.
+  private static final class Columns {
+    final String[] _names;
+    final ColumnMetadata[] _metadata;
+
+    Columns(String[] names, ColumnMetadata[] metadata) {
+      _names = names;
+      _metadata = metadata;
+    }
+
+    /// Index of the column in both arrays, or `-(insertion point) - 1`.
+    int indexOf(String column) {
+      return Arrays.binarySearch(_names, column);
+    }
+  }
 }
diff --git 
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java
 
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java
index d8be4387e2d..897f4ad7252 100644
--- 
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java
+++ 
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySet.java
@@ -43,20 +43,37 @@ import static 
com.google.common.base.Preconditions.checkArgument;
 /// writes: [SegmentMetadataImpl] replaces its arrays when its columns change, 
which leaves an already-returned view
 /// as the snapshot taken at the time of the call.
 ///
+/// `subSet`/`headSet`/`tailSet` are ranges of the same array, and like every 
[NavigableSet] range view they reject
+/// an argument outside their own range rather than silently widening it.
+///
 /// Immutable and thread-safe as long as the backing array is not written in 
place.
 final class SortedStringArraySet extends AbstractSet<String> implements 
NavigableSet<String> {
   private final String[] _elements;
   private final int _from;
   private final int _to;
+  /// The bounds this view was created with, `null` on the side it is 
unbounded on. A range view has to reject an
+  /// argument outside its own range, as [NavigableSet] requires, which the 
array indices alone cannot tell: an
+  /// exclusive endpoint that is not in the array leaves no trace in them.
+  @Nullable
+  private final String _low;
+  private final boolean _lowInclusive;
+  @Nullable
+  private final String _high;
+  private final boolean _highInclusive;
 
   SortedStringArraySet(String[] elements) {
-    this(elements, 0, elements.length);
+    this(elements, 0, elements.length, null, false, null, false);
   }
 
-  private SortedStringArraySet(String[] elements, int from, int to) {
+  private SortedStringArraySet(String[] elements, int from, int to, @Nullable 
String low, boolean lowInclusive,
+      @Nullable String high, boolean highInclusive) {
     _elements = elements;
     _from = from;
     _to = to;
+    _low = low;
+    _lowInclusive = lowInclusive;
+    _high = high;
+    _highInclusive = highInclusive;
   }
 
   @Override
@@ -193,9 +210,11 @@ final class SortedStringArraySet extends 
AbstractSet<String> implements Navigabl
   @Override
   public NavigableSet<String> subSet(String from, boolean fromInclusive, 
String to, boolean toInclusive) {
     checkArgument(from.compareTo(to) <= 0, "from: %s > to: %s", from, to);
+    checkInRange(from, fromInclusive);
+    checkInRange(to, toInclusive);
     int start = fromInclusive ? ceilingIndex(from) : higherIndex(from);
     int end = toInclusive ? higherIndex(to) : ceilingIndex(to);
-    return new SortedStringArraySet(_elements, start, Math.max(start, end));
+    return new SortedStringArraySet(_elements, start, Math.max(start, end), 
from, fromInclusive, to, toInclusive);
   }
 
   @Override
@@ -205,7 +224,9 @@ final class SortedStringArraySet extends 
AbstractSet<String> implements Navigabl
 
   @Override
   public NavigableSet<String> headSet(String to, boolean inclusive) {
-    return new SortedStringArraySet(_elements, _from, inclusive ? 
higherIndex(to) : ceilingIndex(to));
+    checkInRange(to, inclusive);
+    return new SortedStringArraySet(_elements, _from, inclusive ? 
higherIndex(to) : ceilingIndex(to), _low,
+        _lowInclusive, to, inclusive);
   }
 
   @Override
@@ -215,7 +236,9 @@ final class SortedStringArraySet extends 
AbstractSet<String> implements Navigabl
 
   @Override
   public NavigableSet<String> tailSet(String from, boolean inclusive) {
-    return new SortedStringArraySet(_elements, inclusive ? ceilingIndex(from) 
: higherIndex(from), _to);
+    checkInRange(from, inclusive);
+    return new SortedStringArraySet(_elements, inclusive ? ceilingIndex(from) 
: higherIndex(from), _to, from,
+        inclusive, _high, _highInclusive);
   }
 
   @Override
@@ -223,6 +246,32 @@ final class SortedStringArraySet extends 
AbstractSet<String> implements Navigabl
     return tailSet(from, true);
   }
 
+  /// Rejects an argument that a further range call cannot reach from this 
view, as [java.util.TreeSet]'s range views
+  /// do: an endpoint the view excludes is still a legal *exclusive* argument, 
since the range it asks for is empty
+  /// on that side rather than wider.
+  private void checkInRange(String element, boolean inclusive) {
+    boolean inRange = inclusive ? !tooLow(element) && !tooHigh(element)
+        : (_low == null || element.compareTo(_low) >= 0) && (_high == null || 
_high.compareTo(element) >= 0);
+    checkArgument(inRange, "element: %s is out of range: %s%s, %s%s", element, 
_lowInclusive ? "[" : "(", _low, _high,
+        _highInclusive ? "]" : ")");
+  }
+
+  private boolean tooLow(String element) {
+    if (_low == null) {
+      return false;
+    }
+    int comparison = element.compareTo(_low);
+    return comparison < 0 || (comparison == 0 && !_lowInclusive);
+  }
+
+  private boolean tooHigh(String element) {
+    if (_high == null) {
+      return false;
+    }
+    int comparison = element.compareTo(_high);
+    return comparison > 0 || (comparison == 0 && !_highInclusive);
+  }
+
   /// Unlike the range sets above this is a copy, not a view, since the 
backing array is ascending. Nothing on the
   /// segment paths calls it; it exists so the [NavigableSet] contract holds.
   @Override
diff --git 
a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java
 
b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java
index 3564bc9500c..b241ce6548a 100644
--- 
a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java
+++ 
b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/index/metadata/SortedStringArraySetTest.java
@@ -23,6 +23,7 @@ import java.util.List;
 import java.util.NavigableSet;
 import java.util.NoSuchElementException;
 import java.util.TreeSet;
+import java.util.function.Function;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
@@ -103,4 +104,50 @@ public class SortedStringArraySetTest {
     assertThrows(UnsupportedOperationException.class, () -> 
set.iterator().remove());
     assertThrows(IllegalArgumentException.class, () -> set.subSet("f", true, 
"b", true));
   }
+
+  /// A range view is a [NavigableSet] in its own right, so its own range 
calls must answer — and refuse — exactly
+  /// what the same range of a `TreeSet` does, rather than reaching back 
outside their bounds.
+  @Test
+  public void testRangeOfARangeMatchesTreeSet() {
+    for (boolean lowInclusive : List.of(true, false)) {
+      for (boolean highInclusive : List.of(true, false)) {
+        NavigableSet<String> view = set().subSet("c", lowInclusive, "g", 
highInclusive);
+        NavigableSet<String> reference = reference().subSet("c", lowInclusive, 
"g", highInclusive);
+        assertEquals(new ArrayList<>(view), new ArrayList<>(reference));
+        for (String probe : List.of("a", "b", "c", "d", "e", "f", "g", "h", 
"i")) {
+          assertRangeMatches(reference, view, probe, true);
+          assertRangeMatches(reference, view, probe, false);
+        }
+      }
+    }
+
+    NavigableSet<String> head = set().headSet("f", false);
+    assertEquals(new ArrayList<>(head.tailSet("b", true)), List.of("b", "d"));
+    assertThrows(IllegalArgumentException.class, () -> head.tailSet("h", 
true));
+    assertThrows(IllegalArgumentException.class, () -> head.subSet("b", true, 
"h", false));
+    NavigableSet<String> tail = set().tailSet("d", true);
+    assertEquals(new ArrayList<>(tail.headSet("h", true)), List.of("d", "f", 
"h"));
+    assertThrows(IllegalArgumentException.class, () -> tail.headSet("b", 
true));
+  }
+
+  /// Asserts that `view` answers the three range calls the way `reference` 
does, an [IllegalArgumentException] for an
+  /// out-of-range argument included.
+  private static void assertRangeMatches(NavigableSet<String> reference, 
NavigableSet<String> view, String probe,
+      boolean inclusive) {
+    assertRangeMatches(reference, view, probe, set -> set.headSet(probe, 
inclusive));
+    assertRangeMatches(reference, view, probe, set -> set.tailSet(probe, 
inclusive));
+    assertRangeMatches(reference, view, probe, set -> set.subSet(probe, 
inclusive, "h", true));
+  }
+
+  private static void assertRangeMatches(NavigableSet<String> reference, 
NavigableSet<String> view, String probe,
+      Function<NavigableSet<String>, NavigableSet<String>> range) {
+    List<String> expected;
+    try {
+      expected = new ArrayList<>(range.apply(reference));
+    } catch (IllegalArgumentException e) {
+      assertThrows(IllegalArgumentException.class, () -> range.apply(view));
+      return;
+    }
+    assertEquals(new ArrayList<>(range.apply(view)), expected, probe);
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to