This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 823229d0c31 Add a compact reader for table ideal states and external
views in the broker (#19786)
823229d0c31 is described below
commit 823229d0c314e6805cee2ea0836d28843e6da3b4
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Thu Oct 15 12:11:54 2026 +0200
Add a compact reader for table ideal states and external views in the
broker (#19786)
Part of #19782.
---
.../broker/broker/helix/BaseBrokerStarter.java | 3 +-
.../helix/CompactAssignmentZkHelixManager.java | 117 +++++
.../helix/MultiClusterHelixBrokerStarter.java | 7 +-
.../routing/manager/BaseBrokerRoutingManager.java | 6 +-
.../helix/CompactAssignmentZkHelixManagerTest.java | 219 ++++++++
.../CompactAssignmentInstanceSelectorTest.java | 179 +++++++
.../CompactAssignmentRoutingManagerTest.java | 410 +++++++++++++++
.../utils/helix/CompactZNRecordSerializer.java | 468 +++++++++++++++++
.../pinot/common/utils/helix/HelixHelper.java | 41 ++
.../utils/helix/ImmutableSortedArrayMap.java | 562 +++++++++++++++++++++
.../utils/helix/CompactZNRecordSerializerTest.java | 360 +++++++++++++
.../pinot/common/utils/helix/HelixHelperTest.java | 103 ++++
.../utils/helix/ImmutableSortedArrayMapTest.java | 207 ++++++++
.../apache/pinot/spi/utils/CommonConstants.java | 8 +
14 files changed, 2683 insertions(+), 7 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
index ed117e44a81..878d08f2db8 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
@@ -811,7 +811,8 @@ public abstract class BaseBrokerStarter implements
ServiceStartable {
protected void initSpectatorHelixManager() throws Exception {
_spectatorHelixManager =
- HelixManagerFactory.getZKHelixManager(_clusterName, _instanceId,
InstanceType.SPECTATOR, _zkServers);
+
CompactAssignmentZkHelixManager.createSpectatorHelixManager(_brokerConf,
_clusterName, _instanceId,
+ _zkServers);
_spectatorHelixManager.connect();
_helixAdmin = _spectatorHelixManager.getClusterManagmentTool();
_propertyStore = _spectatorHelixManager.getHelixPropertyStore();
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManager.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManager.java
new file mode 100644
index 00000000000..4797f892c92
--- /dev/null
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManager.java
@@ -0,0 +1,117 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.broker.helix;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import org.apache.helix.HelixManager;
+import org.apache.helix.HelixManagerFactory;
+import org.apache.helix.InstanceType;
+import org.apache.helix.PropertyKey;
+import org.apache.helix.manager.zk.ZKHelixManager;
+import org.apache.helix.zookeeper.api.client.RealmAwareZkClient;
+import org.apache.helix.zookeeper.datamodel.serializer.ChainedPathZkSerializer;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.zkclient.serialize.PathBasedZkSerializer;
+import org.apache.pinot.common.utils.helix.CompactZNRecordSerializer;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
+import org.apache.pinot.spi.utils.CommonConstants.Helix;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/// A spectator [ZKHelixManager] that reads table ideal states and external
views with
+/// [CompactZNRecordSerializer].
+///
+/// Helix creates the ZooKeeper client of the manager inside [#connect], with
a default
+/// [ChainedPathZkSerializer] that sends every path to [ZNRecordSerializer].
Right after that, this class replaces the
+/// serializer of the client with one that sends `/<cluster>/IDEALSTATES/...`
and `/<cluster>/EXTERNALVIEW/...` to
+/// [CompactZNRecordSerializer]. The broker resource keeps
[ZNRecordSerializer]: it is small and its readers are
+/// outside of routing. The client keeps the serializer across session expiry,
because Helix creates a new client only
+/// in [#connect], and this class installs the serializer again after each
[#connect].
+///
+/// The serializer applies to everything that reads through this manager: its
data accessor, its `HelixAdmin`, its
+/// property store (no path match) and the Helix callback handlers. All of
them only read the records under those
+/// paths. Writes still go through [ZNRecordSerializer].
+///
+/// The map fields of a compact record are immutable, and the `IdealState` and
`ExternalView` constructors share the
+/// inner instance-state maps. So every table ideal state or external view
read through this manager, by any
+/// reader, has immutable instance-state maps. A read-modify-write of a table
ideal state through this manager (for
+/// example `HelixHelper.updateIdealState`) fails with
[UnsupportedOperationException]. No broker code does that; the
+/// participant manager of the broker keeps its own client and the default
serializer.
+///
+/// Thread-safety: same as [ZKHelixManager]. The serializer is installed from
the thread that calls [#connect],
+/// before the caller registers any listener.
+public class CompactAssignmentZkHelixManager extends ZKHelixManager {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(CompactAssignmentZkHelixManager.class);
+
+ public CompactAssignmentZkHelixManager(String clusterName, String
instanceName, String zkAddress) {
+ super(clusterName, instanceName, InstanceType.SPECTATOR, zkAddress);
+ }
+
+ /// Creates the spectator Helix manager for the broker. Returns a
[CompactAssignmentZkHelixManager] when
+ /// [Broker#CONFIG_OF_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED] is on,
otherwise the default Helix manager.
+ public static HelixManager createSpectatorHelixManager(PinotConfiguration
brokerConf, String clusterName,
+ String instanceId, String zkServers) {
+ if
(brokerConf.getProperty(Broker.CONFIG_OF_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED,
+ Broker.DEFAULT_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED)) {
+ LOGGER.info("Using the compact ideal state and external view reader for
cluster: {}", clusterName);
+ return new CompactAssignmentZkHelixManager(clusterName, instanceId,
zkServers);
+ }
+ return HelixManagerFactory.getZKHelixManager(clusterName, instanceId,
InstanceType.SPECTATOR, zkServers);
+ }
+
+ @Override
+ public void connect()
+ throws Exception {
+ super.connect();
+ installCompactSerializer();
+ }
+
+ /// Replaces the serializer of the ZooKeeper client created by [#connect].
+ private void installCompactSerializer() {
+ RealmAwareZkClient zkClient = _zkclient;
+ Preconditions.checkState(zkClient != null, "ZooKeeper client is not
created for cluster: %s", getClusterName());
+ zkClient.setZkSerializer(createPathBasedSerializer(getClusterName()));
+ }
+
+ /// Returns the ZooKeeper client created by [#connect], for tests.
+ @VisibleForTesting
+ RealmAwareZkClient getZkClient() {
+ return _zkclient;
+ }
+
+ /// Returns a serializer that reads the table ideal states and external
views of the given cluster with
+ /// [CompactZNRecordSerializer], and everything else (including the broker
resource) with [ZNRecordSerializer].
+ private static PathBasedZkSerializer createPathBasedSerializer(String
clusterName) {
+ ZNRecordSerializer defaultSerializer = new ZNRecordSerializer();
+ CompactZNRecordSerializer compactSerializer = new
CompactZNRecordSerializer(defaultSerializer);
+ PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
+ String idealStatesPath = keyBuilder.idealStates().getPath();
+ String externalViewsPath = keyBuilder.externalViews().getPath();
+ // ChainedPathZkSerializer picks the longest matching path, so the broker
resource paths override the parents
+ return ChainedPathZkSerializer.builder(defaultSerializer)
+ .serialize(idealStatesPath, compactSerializer)
+ .serialize(externalViewsPath, compactSerializer)
+
.serialize(keyBuilder.idealStates(Helix.BROKER_RESOURCE_INSTANCE).getPath(),
defaultSerializer)
+
.serialize(keyBuilder.externalView(Helix.BROKER_RESOURCE_INSTANCE).getPath(),
defaultSerializer)
+ .build();
+ }
+}
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
index a1e05842c37..36ce21bc594 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
@@ -27,8 +27,6 @@ import java.util.Map;
import java.util.Set;
import org.apache.helix.HelixConstants.ChangeType;
import org.apache.helix.HelixManager;
-import org.apache.helix.HelixManagerFactory;
-import org.apache.helix.InstanceType;
import org.apache.pinot.broker.routing.manager.MultiClusterRoutingManager;
import
org.apache.pinot.broker.routing.manager.RemoteClusterBrokerRoutingManager;
import org.apache.pinot.common.config.provider.TableCache;
@@ -137,8 +135,9 @@ public class MultiClusterHelixBrokerStarter extends
BaseBrokerStarter {
String clusterName = entry.getKey();
String zkServers = entry.getValue();
try {
- HelixManager helixManager = HelixManagerFactory.getZKHelixManager(
- clusterName, _instanceId, InstanceType.SPECTATOR, zkServers);
+ HelixManager helixManager =
+
CompactAssignmentZkHelixManager.createSpectatorHelixManager(_brokerConf,
clusterName, _instanceId,
+ zkServers);
helixManager.connect();
_remoteSpectatorHelixManager.put(clusterName, helixManager);
LOGGER.info("[multi-cluster] Connected to remote cluster '{}' at ZK:
{}", clusterName, zkServers);
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
index 5aabc29e179..bb28278db52 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
@@ -81,6 +81,7 @@ import org.apache.pinot.common.metrics.BrokerTimer;
import org.apache.pinot.common.request.BrokerRequest;
import org.apache.pinot.common.utils.HashUtil;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
+import org.apache.pinot.common.utils.helix.HelixHelper;
import org.apache.pinot.core.routing.RoutingManager;
import org.apache.pinot.core.routing.RoutingTable;
import org.apache.pinot.core.routing.SegmentsToQuery;
@@ -366,7 +367,8 @@ public abstract class BaseBrokerRoutingManager
ZNRecord znRecord = _zkDataAccessor.get(idealStatePath, stat,
AccessOption.PERSISTENT);
if (znRecord != null) {
znRecord.setVersion(stat.getVersion());
- return new IdealState(znRecord);
+ // Skips the copy of the outer segment map when the record comes from
the compact reader
+ return HelixHelper.toIdealState(znRecord);
} else {
return null;
}
@@ -378,7 +380,7 @@ public abstract class BaseBrokerRoutingManager
ZNRecord znRecord = _zkDataAccessor.get(externalViewPath, stat,
AccessOption.PERSISTENT);
if (znRecord != null) {
znRecord.setVersion(stat.getVersion());
- return new ExternalView(znRecord);
+ return HelixHelper.toExternalView(znRecord);
} else {
return null;
}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManagerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManagerTest.java
new file mode 100644
index 00000000000..f363d4bc275
--- /dev/null
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/helix/CompactAssignmentZkHelixManagerTest.java
@@ -0,0 +1,219 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.broker.helix;
+
+import java.util.Map;
+import java.util.TreeMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.AccessOption;
+import org.apache.helix.BaseDataAccessor;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixManager;
+import org.apache.helix.PropertyKey;
+import org.apache.helix.manager.zk.ZKHelixAdmin;
+import org.apache.helix.model.ExternalView;
+import org.apache.helix.model.IdealState;
+import org.apache.helix.zookeeper.api.client.RealmAwareZkClient;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.impl.client.ZkClient;
+import org.apache.helix.zookeeper.zkclient.ZkConnection;
+import org.apache.helix.zookeeper.zkclient.serialize.PathBasedZkSerializer;
+import org.apache.pinot.common.utils.ZkStarter;
+import org.apache.pinot.common.utils.helix.ImmutableSortedArrayMap;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
+import org.apache.pinot.spi.utils.CommonConstants.Helix;
+import org.apache.pinot.util.TestUtils;
+import org.apache.zookeeper.Watcher;
+import org.apache.zookeeper.ZooKeeper;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that [CompactAssignmentZkHelixManager] installs
+/// [org.apache.pinot.common.utils.helix.CompactZNRecordSerializer] on the
client that Helix creates in `connect()`, for
+/// table ideal states and external views only.
+///
+/// This test fails when a Helix upgrade changes how
[org.apache.helix.manager.zk.ZKHelixManager] creates its ZooKeeper
+/// client or its serializer, for example by renaming the protected
`_zkclient` field (compile error) or by creating the
+/// client outside of `connect()` (the records read through the manager would
no longer be compact).
+public class CompactAssignmentZkHelixManagerTest {
+ private static final String CLUSTER_NAME =
"CompactAssignmentZkHelixManagerTest";
+ private static final String TABLE_NAME = "myTable_OFFLINE";
+
+ private ZkStarter.ZookeeperInstance _zookeeperInstance;
+ private ZkClient _zkClient;
+
+ @BeforeClass
+ public void setUp() {
+ _zookeeperInstance = ZkStarter.startLocalZkServer();
+ String zkUrl = _zookeeperInstance.getZkUrl();
+ HelixAdmin helixAdmin = new
ZKHelixAdmin.Builder().setZkAddress(zkUrl).build();
+ try {
+ helixAdmin.addCluster(CLUSTER_NAME);
+ } finally {
+ helixAdmin.close();
+ }
+ _zkClient = new ZkClient.Builder().setZkServer(zkUrl).setZkSerializer(new
ZNRecordSerializer()).build();
+ PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
+ _zkClient.createPersistent(keyBuilder.idealStates(TABLE_NAME).getPath(),
assignment(TABLE_NAME));
+ _zkClient.createPersistent(keyBuilder.externalView(TABLE_NAME).getPath(),
assignment(TABLE_NAME));
+
_zkClient.createPersistent(keyBuilder.idealStates(Helix.BROKER_RESOURCE_INSTANCE).getPath(),
+ assignment(Helix.BROKER_RESOURCE_INSTANCE));
+
_zkClient.createPersistent(keyBuilder.externalView(Helix.BROKER_RESOURCE_INSTANCE).getPath(),
+ assignment(Helix.BROKER_RESOURCE_INSTANCE));
+ }
+
+ @AfterClass
+ public void tearDown() {
+ if (_zkClient != null) {
+ _zkClient.close();
+ }
+ ZkStarter.stopLocalZkServer(_zookeeperInstance);
+ }
+
+ private static ZNRecord assignment(String id) {
+ ZNRecord record = new ZNRecord(id);
+ // Compressed, because the compact serializer only parses gzip-compressed
znodes itself
+ record.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD, true);
+ record.setSimpleField("REBALANCE_MODE", "CUSTOMIZED");
+ record.setMapField("segment_1", Map.of("Server_b_8098", "ONLINE",
"Server_a_8098", "ONLINE"));
+ record.setMapField("segment_0", Map.of("Server_b_8098", "ONLINE",
"Server_a_8098", "ONLINE"));
+ return record;
+ }
+
+ @Test
+ public void testFactoryHonorsFlag() {
+ String zkUrl = _zookeeperInstance.getZkUrl();
+ HelixManager defaultManager =
+ CompactAssignmentZkHelixManager.createSpectatorHelixManager(new
PinotConfiguration(), CLUSTER_NAME,
+ "Broker_localhost_1", zkUrl);
+ assertFalse(defaultManager instanceof CompactAssignmentZkHelixManager);
+ PinotConfiguration conf = new PinotConfiguration();
+
conf.setProperty(Broker.CONFIG_OF_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED,
true);
+ HelixManager compactManager =
+ CompactAssignmentZkHelixManager.createSpectatorHelixManager(conf,
CLUSTER_NAME, "Broker_localhost_1", zkUrl);
+ assertTrue(compactManager instanceof CompactAssignmentZkHelixManager);
+ }
+
+ @Test
+ public void testReadsThroughManagerAreCompact()
+ throws Exception {
+ CompactAssignmentZkHelixManager manager =
+ new CompactAssignmentZkHelixManager(CLUSTER_NAME,
"Broker_localhost_2", _zookeeperInstance.getZkUrl());
+ manager.connect();
+ try {
+ assertReads(manager);
+ // A new connect() creates a new client; the serializer must be
installed on it again
+ manager.disconnect();
+ manager.connect();
+ assertReads(manager);
+ } finally {
+ manager.disconnect();
+ }
+ }
+
+ @Test
+ public void testSerializerSurvivesSessionExpiry()
+ throws Exception {
+ CompactAssignmentZkHelixManager manager =
+ new CompactAssignmentZkHelixManager(CLUSTER_NAME,
"Broker_localhost_3", _zookeeperInstance.getZkUrl());
+ manager.connect();
+ try {
+ assertReads(manager);
+ RealmAwareZkClient zkClient = manager.getZkClient();
+ PathBasedZkSerializer serializer =
manager.getZkClient().getZkSerializer();
+ long oldSessionId = ((ZkClient) zkClient).getSessionId();
+ expireSession((ZkClient) zkClient);
+ TestUtils.waitForCondition(aVoid -> {
+ long sessionId = ((ZkClient) zkClient).getSessionId();
+ return sessionId != 0 && sessionId != oldSessionId &&
zkClient.waitUntilConnected(1, TimeUnit.SECONDS);
+ }, 30_000L, "Failed to get a new ZooKeeper session");
+ // Helix keeps the client (and so the serializer) across a session expiry
+ assertSame(manager.getZkClient(), zkClient);
+ assertSame(manager.getZkClient().getZkSerializer(), serializer);
+ assertReads(manager);
+ } finally {
+ manager.disconnect();
+ }
+ }
+
+ /// Expires the session of the given client: opens a second connection with
the same session id and password, then
+ /// closes it. The server then expires the session, and the client gets a
new one.
+ private void expireSession(ZkClient zkClient)
+ throws Exception {
+ ZooKeeper zooKeeper = ((ZkConnection)
zkClient.getConnection()).getZookeeper();
+ CountDownLatch connected = new CountDownLatch(1);
+ ZooKeeper duplicate = new ZooKeeper(_zookeeperInstance.getZkUrl(), 10_000,
event -> {
+ if (event.getState() == Watcher.Event.KeeperState.SyncConnected) {
+ connected.countDown();
+ }
+ }, zooKeeper.getSessionId(), zooKeeper.getSessionPasswd());
+ try {
+ assertTrue(connected.await(30, TimeUnit.SECONDS), "Failed to connect
with the duplicate session");
+ } finally {
+ duplicate.close();
+ }
+ }
+
+ private static void assertReads(CompactAssignmentZkHelixManager manager) {
+ assertNotNull(manager.getZkClient().getZkSerializer());
+ BaseDataAccessor<ZNRecord> accessor =
manager.getHelixDataAccessor().getBaseDataAccessor();
+ PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
+
+ // Table ideal state and external view: compact, equal to the written
content, inner maps shared
+ for (String path : new String[]{
+ keyBuilder.idealStates(TABLE_NAME).getPath(),
keyBuilder.externalView(TABLE_NAME).getPath()
+ }) {
+ ZNRecord record = accessor.get(path, null, AccessOption.PERSISTENT);
+ assertTrue(record.getMapFields() instanceof ImmutableSortedArrayMap,
path);
+ assertEquals(record, assignment(TABLE_NAME));
+ assertSame(record.getMapField("segment_0"),
record.getMapField("segment_1"));
+ }
+
+ // Broker resource: default deserializer
+ for (String path : new String[]{
+ keyBuilder.idealStates(Helix.BROKER_RESOURCE_INSTANCE).getPath(),
+ keyBuilder.externalView(Helix.BROKER_RESOURCE_INSTANCE).getPath()
+ }) {
+ ZNRecord record = accessor.get(path, null, AccessOption.PERSISTENT);
+ assertFalse(record.getMapFields() instanceof ImmutableSortedArrayMap,
path);
+ assertEquals(record, assignment(Helix.BROKER_RESOURCE_INSTANCE));
+ }
+
+ // Readers through the data accessor and the admin of the manager see the
compact records too
+ IdealState idealState =
manager.getHelixDataAccessor().getProperty(keyBuilder.idealStates(TABLE_NAME));
+ assertTrue(idealState.getInstanceStateMap("segment_0") instanceof
ImmutableSortedArrayMap);
+ ExternalView externalView =
manager.getClusterManagmentTool().getResourceExternalView(CLUSTER_NAME,
TABLE_NAME);
+ assertTrue(externalView.getStateMap("segment_0") instanceof
ImmutableSortedArrayMap);
+ assertEquals(new TreeMap<>(externalView.getRecord().getMapFields()),
assignment(TABLE_NAME).getMapFields());
+ ExternalView brokerResource =
+
manager.getClusterManagmentTool().getResourceExternalView(CLUSTER_NAME,
Helix.BROKER_RESOURCE_INSTANCE);
+ assertFalse(brokerResource.getStateMap("segment_0") instanceof
ImmutableSortedArrayMap);
+ }
+}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/CompactAssignmentInstanceSelectorTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/CompactAssignmentInstanceSelectorTest.java
new file mode 100644
index 00000000000..7a8ff4b11af
--- /dev/null
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/instanceselector/CompactAssignmentInstanceSelectorTest.java
@@ -0,0 +1,179 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.routing.instanceselector;
+
+import java.time.Clock;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+import org.apache.helix.model.ExternalView;
+import org.apache.helix.model.IdealState;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.util.GZipCompressionUtil;
+import org.apache.helix.zookeeper.zkclient.serialize.ZkSerializer;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.common.utils.helix.CompactZNRecordSerializer;
+import org.apache.pinot.common.utils.helix.HelixHelper;
+import org.apache.pinot.common.utils.helix.ImmutableSortedArrayMap;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING;
+import static
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ERROR;
+import static
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.OFFLINE;
+import static
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that the instance selectors route the same way, and do not mutate
the assignment, when the ideal state and
+/// external view come from [CompactZNRecordSerializer] instead of
[ZNRecordSerializer].
+public class CompactAssignmentInstanceSelectorTest {
+ private static final String TABLE_NAME = "testTable_OFFLINE";
+ private static final InstanceSelectorConfig CONFIG = new
InstanceSelectorConfig(false, 300, false);
+
+ @DataProvider
+ public Object[][] selectorTypes() {
+ return new Object[][]{{"balanced", false}, {"replicaGroup", false},
{"strictReplicaGroup", true}};
+ }
+
+ private static InstanceSelector newSelector(String type) {
+ switch (type) {
+ case "balanced":
+ return new BalancedInstanceSelector();
+ case "replicaGroup":
+ return new ReplicaGroupInstanceSelector();
+ default:
+ return new StrictReplicaGroupInstanceSelector();
+ }
+ }
+
+ @Test(dataProvider = "selectorTypes")
+ public void testSameSelection(String type, boolean upsert) {
+ Random random = new Random(7);
+ List<String> instances = new ArrayList<>();
+ for (int i = 0; i < 6; i++) {
+ instances.add("Server_host" + i + "_8098");
+ }
+ ZNRecord idealStateRecord = new ZNRecord(TABLE_NAME);
+ idealStateRecord.setSimpleField("REBALANCE_MODE", "CUSTOMIZED");
+ ZNRecord externalViewRecord = new ZNRecord(TABLE_NAME);
+ // Compressed, because the compact serializer only parses gzip-compressed
znodes itself
+
idealStateRecord.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD,
true);
+
externalViewRecord.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD,
true);
+ List<String> segments = new ArrayList<>();
+ for (int i = 0; i < 300; i++) {
+ String segment = "segment_" + i;
+ segments.add(segment);
+ // Replica groups: (0, 3), (1, 4), (2, 5); insertion order is not sorted
on purpose
+ int group = i % 3;
+ Map<String, String> idealStateMap = new LinkedHashMap<>();
+ idealStateMap.put(instances.get(group + 3), i % 10 == 0 ? CONSUMING :
ONLINE);
+ idealStateMap.put(instances.get(group), ONLINE);
+ idealStateRecord.setMapField(segment, idealStateMap);
+ if (i % 50 == 1) {
+ // Missing from the external view
+ continue;
+ }
+ Map<String, String> externalViewMap = new LinkedHashMap<>();
+ for (Map.Entry<String, String> entry : idealStateMap.entrySet()) {
+ int r = random.nextInt(20);
+ externalViewMap.put(entry.getKey(), r == 0 ? ERROR : r == 1 ? OFFLINE
: entry.getValue());
+ }
+ externalViewRecord.setMapField(segment, externalViewMap);
+ }
+ ZNRecordSerializer defaultSerializer = new ZNRecordSerializer();
+ byte[] idealStateBytes = defaultSerializer.serialize(idealStateRecord);
+ byte[] externalViewBytes = defaultSerializer.serialize(externalViewRecord);
+ assertTrue(GZipCompressionUtil.isCompressed(idealStateBytes));
+ assertTrue(GZipCompressionUtil.isCompressed(externalViewBytes));
+ Set<String> enabledInstances = new HashSet<>(instances);
+
+ List<Map<String, String>> expected =
+ select(type, upsert, defaultSerializer, idealStateBytes,
externalViewBytes, enabledInstances, segments, false);
+ List<Map<String, String>> actual =
+ select(type, upsert, new CompactZNRecordSerializer(), idealStateBytes,
externalViewBytes, enabledInstances,
+ segments, true);
+ assertEquals(actual, expected);
+ }
+
+ /// Initializes a selector from the given znode bytes, runs a few
selections, and returns their results.
+ private static List<Map<String, String>> select(String type, boolean upsert,
ZkSerializer serializer,
+ byte[] idealStateBytes, byte[] externalViewBytes, Set<String>
enabledInstances, List<String> segments,
+ boolean expectCompact) {
+ // Wrap the records as the routing manager does: a compact record is
shared, a default one is copied
+ ZNRecord idealStateRecord = (ZNRecord)
serializer.deserialize(idealStateBytes);
+ ZNRecord externalViewRecord = (ZNRecord)
serializer.deserialize(externalViewBytes);
+ IdealState idealState = HelixHelper.toIdealState(idealStateRecord);
+ ExternalView externalView = HelixHelper.toExternalView(externalViewRecord);
+ assertEquals(HelixHelper.isCompact(idealStateRecord), expectCompact);
+ assertEquals(HelixHelper.isCompact(externalViewRecord), expectCompact);
+ if (expectCompact) {
+ // The selectors see the immutable outer maps of the parsed records, not
copies
+ assertSame(idealState.getRecord().getMapFields(),
idealStateRecord.getMapFields());
+ assertSame(externalView.getRecord().getMapFields(),
externalViewRecord.getMapFields());
+ assertTrue(idealState.getRecord().getMapFields() instanceof
ImmutableSortedArrayMap);
+ assertTrue(idealState.getInstanceStateMap("segment_0") instanceof
ImmutableSortedArrayMap);
+ assertTrue(externalView.getStateMap("segment_0") instanceof
ImmutableSortedArrayMap);
+ } else {
+ assertNotSame(idealState.getRecord().getMapFields(),
idealStateRecord.getMapFields());
+ }
+ Set<String> onlineSegments = new HashSet<>();
+ for (Map.Entry<String, Map<String, String>> entry :
idealState.getRecord().getMapFields().entrySet()) {
+ if (entry.getValue().containsValue(ONLINE) ||
entry.getValue().containsValue(CONSUMING)) {
+ onlineSegments.add(entry.getKey());
+ }
+ }
+
+ TableConfig tableConfig = mock(TableConfig.class);
+ when(tableConfig.getTableName()).thenReturn(TABLE_NAME);
+ when(tableConfig.isUpsertEnabled()).thenReturn(upsert);
+ @SuppressWarnings("unchecked")
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
mock(ZkHelixPropertyStore.class);
+ BrokerRequest brokerRequest = mock(BrokerRequest.class);
+ when(brokerRequest.getPinotQuery()).thenReturn(mock(PinotQuery.class));
+
+ InstanceSelector selector = newSelector(type);
+ selector.init(tableConfig, propertyStore, mock(BrokerMetrics.class), null,
Clock.systemUTC(), CONFIG,
+ enabledInstances, Map.of(), idealState, externalView, onlineSegments);
+ List<Map<String, String>> results = new ArrayList<>();
+ for (int requestId = 0; requestId < 6; requestId++) {
+ InstanceSelector.SelectionResult result = selector.select(brokerRequest,
segments, requestId);
+ results.add(result.getSegmentToInstanceMap());
+ results.add(Map.of("unavailable",
String.valueOf(result.getUnavailableSegments())));
+ }
+ // A second assignment change on the same selector
+ selector.onAssignmentChange(idealState, externalView, onlineSegments);
+ results.add(selector.select(brokerRequest, segments,
0).getSegmentToInstanceMap());
+ return results;
+ }
+}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/CompactAssignmentRoutingManagerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/CompactAssignmentRoutingManagerTest.java
new file mode 100644
index 00000000000..bdc5024e98b
--- /dev/null
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/CompactAssignmentRoutingManagerTest.java
@@ -0,0 +1,410 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.routing.manager;
+
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.AccessOption;
+import org.apache.helix.HelixAdmin;
+import org.apache.helix.HelixConstants.ChangeType;
+import org.apache.helix.HelixManager;
+import org.apache.helix.HelixManagerFactory;
+import org.apache.helix.InstanceType;
+import org.apache.helix.PropertyKey;
+import org.apache.helix.manager.zk.ZKHelixAdmin;
+import org.apache.helix.model.InstanceConfig;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.impl.client.ZkClient;
+import org.apache.pinot.broker.broker.helix.CompactAssignmentZkHelixManager;
+import org.apache.pinot.common.metadata.ZKMetadataProvider;
+import org.apache.pinot.common.metadata.segment.SegmentPartitionMetadata;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.utils.ZkStarter;
+import org.apache.pinot.common.utils.helix.HelixHelper;
+import org.apache.pinot.core.routing.RoutingTable;
+import org.apache.pinot.core.routing.SegmentsToQuery;
+import org.apache.pinot.core.routing.TablePartitionInfo;
+import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo;
+import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo;
+import org.apache.pinot.core.transport.ServerInstance;
+import
org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
+import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
+import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.RoutingConfig;
+import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
+import
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that a [BrokerRoutingManager] on a
[CompactAssignmentZkHelixManager] routes exactly like one on the default
+/// Helix manager. The routing manager wraps the compact records without
copying them (see
+/// [HelixHelper#toIdealState]), so every routing component (instance
selector, segment pruners, time
+/// boundary manager, segment partition metadata manager) reads the immutable
outer map of the parsed record.
+///
+/// The test writes a hybrid table into a local ZooKeeper, builds the routing
in both managers, compares what they
+/// return, then changes the assignment and compares again after the
assignment change.
+public class CompactAssignmentRoutingManagerTest {
+ private static final String CLUSTER_NAME =
"CompactAssignmentRoutingManagerTest";
+ private static final String RAW_TABLE_NAME = "myTable";
+ private static final String OFFLINE_TABLE_NAME =
TableNameBuilder.OFFLINE.tableNameWithType(RAW_TABLE_NAME);
+ private static final String REALTIME_TABLE_NAME =
TableNameBuilder.REALTIME.tableNameWithType(RAW_TABLE_NAME);
+ private static final String TIME_COLUMN = "ts";
+ private static final String PARTITION_COLUMN = "memberId";
+ private static final int NUM_PARTITIONS = 4;
+ private static final int NUM_SERVERS = 4;
+ private static final long BASE_TIME_MS = 1_700_000_000_000L;
+ private static final long HOUR_MS = TimeUnit.HOURS.toMillis(1);
+ private static final List<String> QUERIES = List.of("SELECT * FROM myTable",
+ "SELECT * FROM myTable WHERE memberId = 2",
+ "SELECT * FROM myTable WHERE ts > " + (BASE_TIME_MS + 10 * HOUR_MS),
+ "SELECT * FROM myTable WHERE memberId IN (1, 3) AND ts < " +
(BASE_TIME_MS + 6 * HOUR_MS));
+
+ private ZkStarter.ZookeeperInstance _zookeeperInstance;
+ private ZkClient _zkClient;
+ private HelixManager _defaultHelixManager;
+ private HelixManager _compactHelixManager;
+ private ZkHelixPropertyStore<ZNRecord> _propertyStore;
+ private BrokerRoutingManager _defaultRoutingManager;
+ private BrokerRoutingManager _compactRoutingManager;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ _zookeeperInstance = ZkStarter.startLocalZkServer();
+ String zkUrl = _zookeeperInstance.getZkUrl();
+ HelixAdmin helixAdmin = new
ZKHelixAdmin.Builder().setZkAddress(zkUrl).build();
+ try {
+ helixAdmin.addCluster(CLUSTER_NAME);
+ for (int i = 0; i < NUM_SERVERS; i++) {
+ InstanceConfig instanceConfig = new InstanceConfig(server(i));
+ instanceConfig.setHostName("host" + i);
+ instanceConfig.setPort(Integer.toString(8000 + i));
+ helixAdmin.addInstance(CLUSTER_NAME, instanceConfig);
+ }
+ } finally {
+ helixAdmin.close();
+ }
+ _zkClient = new ZkClient.Builder().setZkServer(zkUrl).setZkSerializer(new
ZNRecordSerializer()).build();
+
+ _defaultHelixManager =
+ HelixManagerFactory.getZKHelixManager(CLUSTER_NAME,
"Broker_default_8099", InstanceType.SPECTATOR, zkUrl);
+ _defaultHelixManager.connect();
+ _compactHelixManager =
+ new CompactAssignmentZkHelixManager(CLUSTER_NAME,
"Broker_compact_8099", zkUrl);
+ _compactHelixManager.connect();
+ _propertyStore = _defaultHelixManager.getHelixPropertyStore();
+
+ Schema schema = new Schema.SchemaBuilder().setSchemaName(RAW_TABLE_NAME)
+ .addSingleValueDimension(PARTITION_COLUMN, FieldSpec.DataType.INT)
+ .addDateTime(TIME_COLUMN, FieldSpec.DataType.LONG,
"1:MILLISECONDS:EPOCH", "1:MILLISECONDS")
+ .build();
+ ZKMetadataProvider.setSchema(_propertyStore, schema);
+ ZKMetadataProvider.setTableConfig(_propertyStore,
tableConfig(TableType.OFFLINE));
+ ZKMetadataProvider.setTableConfig(_propertyStore,
tableConfig(TableType.REALTIME));
+
+ // Offline: 20 segments on 2 replicas each; one replica in ERROR and one
segment missing from the external view
+ Map<String, Map<String, String>> offlineIdealState = new TreeMap<>();
+ Map<String, Map<String, String>> offlineExternalView = new TreeMap<>();
+ for (int i = 0; i < 20; i++) {
+ String segment = offlineSegment(i);
+ setSegmentZkMetadata(OFFLINE_TABLE_NAME, segment, i % NUM_PARTITIONS,
BASE_TIME_MS + i * HOUR_MS, false);
+ Map<String, String> instanceStateMap = replicas(i,
SegmentStateModel.ONLINE);
+ offlineIdealState.put(segment, instanceStateMap);
+ if (i == 7) {
+ continue;
+ }
+ Map<String, String> externalViewStateMap = new
TreeMap<>(instanceStateMap);
+ if (i == 5) {
+ externalViewStateMap.put(server(i % NUM_SERVERS),
SegmentStateModel.ERROR);
+ }
+ offlineExternalView.put(segment, externalViewStateMap);
+ }
+ writeAssignment(OFFLINE_TABLE_NAME, offlineIdealState,
offlineExternalView, true);
+
+ // Realtime: 3 sequences per partition, the last one CONSUMING
+ Map<String, Map<String, String>> realtimeAssignment = new TreeMap<>();
+ for (int partition = 0; partition < NUM_PARTITIONS; partition++) {
+ for (int sequence = 0; sequence < 3; sequence++) {
+ String segment = RAW_TABLE_NAME + "__" + partition + "__" + sequence +
"__20240101T0000Z";
+ boolean consuming = sequence == 2;
+ setSegmentZkMetadata(REALTIME_TABLE_NAME, segment, partition,
+ BASE_TIME_MS + (15 + sequence) * HOUR_MS, consuming);
+ realtimeAssignment.put(segment,
+ replicas(partition, consuming ? SegmentStateModel.CONSUMING :
SegmentStateModel.ONLINE));
+ }
+ }
+ writeAssignment(REALTIME_TABLE_NAME, realtimeAssignment,
realtimeAssignment, true);
+
+ _defaultRoutingManager = newRoutingManager(_defaultHelixManager);
+ _compactRoutingManager = newRoutingManager(_compactHelixManager);
+ }
+
+ @AfterClass
+ public void tearDown() {
+ if (_defaultRoutingManager != null) {
+ _defaultRoutingManager.stop();
+ }
+ if (_compactRoutingManager != null) {
+ _compactRoutingManager.stop();
+ }
+ if (_defaultHelixManager != null) {
+ _defaultHelixManager.disconnect();
+ }
+ if (_compactHelixManager != null) {
+ _compactHelixManager.disconnect();
+ }
+ if (_zkClient != null) {
+ _zkClient.close();
+ }
+ ZkStarter.stopLocalZkServer(_zookeeperInstance);
+ }
+
+ private static String server(int i) {
+ return "Server_host" + i + "_" + (8000 + i);
+ }
+
+ private static String offlineSegment(int i) {
+ return OFFLINE_TABLE_NAME + "_" + i;
+ }
+
+ private static Map<String, String> replicas(int i, String state) {
+ Map<String, String> instanceStateMap = new TreeMap<>();
+ instanceStateMap.put(server(i % NUM_SERVERS), state);
+ instanceStateMap.put(server((i + 1) % NUM_SERVERS), state);
+ return instanceStateMap;
+ }
+
+ private static TableConfig tableConfig(TableType tableType) {
+ return new TableConfigBuilder(tableType).setTableName(RAW_TABLE_NAME)
+ .setTimeColumnName(TIME_COLUMN)
+ .setSegmentPartitionConfig(new SegmentPartitionConfig(
+ Map.of(PARTITION_COLUMN, new ColumnPartitionConfig("Modulo",
NUM_PARTITIONS))))
+ .setRoutingConfig(new RoutingConfig(null,
+ List.of(RoutingConfig.PARTITION_SEGMENT_PRUNER_TYPE,
RoutingConfig.TIME_SEGMENT_PRUNER_TYPE), null, false))
+ .build();
+ }
+
+ private void setSegmentZkMetadata(String tableNameWithType, String segment,
int partition, long endTimeMs,
+ boolean consuming) {
+ SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segment);
+ // Old segments, so that the instance selectors do not treat them as new
+ segmentZKMetadata.setCreationTime(BASE_TIME_MS);
+ segmentZKMetadata.setPushTime(BASE_TIME_MS);
+ segmentZKMetadata.setPartitionMetadata(new
SegmentPartitionMetadata(Map.of(PARTITION_COLUMN,
+ new ColumnPartitionMetadata("Modulo", NUM_PARTITIONS,
Set.of(partition), null))));
+ if (consuming) {
+
segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS);
+ } else {
+ segmentZKMetadata.setStartTime(endTimeMs - HOUR_MS);
+ segmentZKMetadata.setEndTime(endTimeMs);
+ segmentZKMetadata.setTimeUnit(TimeUnit.MILLISECONDS);
+ segmentZKMetadata.setTotalDocs(1000);
+ }
+ ZKMetadataProvider.setSegmentZKMetadata(_propertyStore, tableNameWithType,
segmentZKMetadata);
+ }
+
+ private void writeAssignment(String tableNameWithType, Map<String,
Map<String, String>> idealState,
+ Map<String, Map<String, String>> externalView, boolean create) {
+ PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
+ ZNRecord idealStateRecord = new ZNRecord(tableNameWithType);
+ idealStateRecord.setSimpleField("IDEAL_STATE_MODE", "CUSTOMIZED");
+ idealStateRecord.setSimpleField("REBALANCE_MODE", "CUSTOMIZED");
+ idealStateRecord.setSimpleField("STATE_MODEL_DEF_REF",
"SegmentOnlineOfflineStateModel");
+ idealStateRecord.setSimpleField("NUM_PARTITIONS",
Integer.toString(idealState.size()));
+ idealStateRecord.setSimpleField("HELIX_ENABLED", "true");
+ idealStateRecord.setMapFields(idealState);
+ ZNRecord externalViewRecord = new ZNRecord(tableNameWithType);
+ externalViewRecord.setMapFields(externalView);
+ // Compressed, because the compact serializer only parses gzip-compressed
znodes itself
+
idealStateRecord.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD,
true);
+
externalViewRecord.setBooleanField(ZNRecord.ENABLE_COMPRESSION_BOOLEAN_FIELD,
true);
+ String idealStatePath =
keyBuilder.idealStates(tableNameWithType).getPath();
+ String externalViewPath =
keyBuilder.externalView(tableNameWithType).getPath();
+ if (create) {
+ _zkClient.createPersistent(idealStatePath, idealStateRecord);
+ _zkClient.createPersistent(externalViewPath, externalViewRecord);
+ } else {
+ _zkClient.writeData(idealStatePath, idealStateRecord);
+ _zkClient.writeData(externalViewPath, externalViewRecord);
+ }
+ }
+
+ private BrokerRoutingManager newRoutingManager(HelixManager helixManager) {
+ PinotConfiguration config = new PinotConfiguration();
+
config.setProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_PARTITION_METADATA_MANAGER,
true);
+ BrokerRoutingManager routingManager =
+ new BrokerRoutingManager(mock(BrokerMetrics.class),
mock(ServerRoutingStatsManager.class), config);
+ routingManager.init(helixManager);
+ routingManager.processClusterChange(ChangeType.INSTANCE_CONFIG);
+ routingManager.buildRouting(OFFLINE_TABLE_NAME);
+ routingManager.buildRouting(REALTIME_TABLE_NAME);
+ return routingManager;
+ }
+
+ @Test
+ public void testSameRoutingAsDefaultReader() {
+ assertReadsAreCompact();
+ Map<String, Object> expected = snapshot(_defaultRoutingManager);
+ assertEquals(snapshot(_compactRoutingManager), expected);
+
+ // Sanity checks so that the comparison is not vacuous
+ assertEquals(expected.get(OFFLINE_TABLE_NAME + ".timeBoundary"),
+ TIME_COLUMN + "=" + (BASE_TIME_MS + 19 * HOUR_MS -
TimeUnit.DAYS.toMillis(1)));
+ assertEquals(expected.get(OFFLINE_TABLE_NAME + ".disabled"), false);
+ assertNotNull(expected.get(OFFLINE_TABLE_NAME + ".partitions"));
+ assertNotNull(expected.get(REALTIME_TABLE_NAME + ".replicatedServers"));
+ @SuppressWarnings("unchecked")
+ Map<String, Object> filtered = (Map<String, Object>)
expected.get(OFFLINE_TABLE_NAME + "." + QUERIES.get(1) + ".0");
+ assertEquals(filtered.get("numPruned"), 15);
+ @SuppressWarnings("unchecked")
+ Map<String, Object> all = (Map<String, Object>)
expected.get(OFFLINE_TABLE_NAME + "." + QUERIES.get(0) + ".0");
+ assertEquals(all.get("unavailable"), List.of(offlineSegment(7)));
+
+ // Change the offline assignment: add a segment, recover the ERROR
replica, drop a segment
+ Map<String, Map<String, String>> idealState = new TreeMap<>();
+ Map<String, Map<String, String>> externalView = new TreeMap<>();
+ for (int i = 1; i < 21; i++) {
+ String segment = offlineSegment(i);
+ if (i == 20) {
+ setSegmentZkMetadata(OFFLINE_TABLE_NAME, segment, i % NUM_PARTITIONS,
BASE_TIME_MS + i * HOUR_MS, false);
+ }
+ idealState.put(segment, replicas(i, SegmentStateModel.ONLINE));
+ externalView.put(segment, replicas(i, SegmentStateModel.ONLINE));
+ }
+ writeAssignment(OFFLINE_TABLE_NAME, idealState, externalView, false);
+ _defaultRoutingManager.processClusterChange(ChangeType.IDEAL_STATE);
+ _defaultRoutingManager.processClusterChange(ChangeType.EXTERNAL_VIEW);
+ _compactRoutingManager.processClusterChange(ChangeType.IDEAL_STATE);
+ _compactRoutingManager.processClusterChange(ChangeType.EXTERNAL_VIEW);
+
+ assertReadsAreCompact();
+ Map<String, Object> expectedAfterChange = snapshot(_defaultRoutingManager);
+ assertEquals(snapshot(_compactRoutingManager), expectedAfterChange);
+ assertEquals(expectedAfterChange.get(OFFLINE_TABLE_NAME + ".timeBoundary"),
+ TIME_COLUMN + "=" + (BASE_TIME_MS + 20 * HOUR_MS -
TimeUnit.DAYS.toMillis(1)));
+ @SuppressWarnings("unchecked")
+ Map<String, Object> allAfterChange =
+ (Map<String, Object>) expectedAfterChange.get(OFFLINE_TABLE_NAME + "."
+ QUERIES.get(0) + ".0");
+ assertEquals(allAfterChange.get("unavailable"), List.of());
+ }
+
+ /// Checks the raw reads of the data accessors: the compact manager reads
compact records and the default manager
+ /// reads regular records. The wrapping of compact records without a copy is
checked in
+ /// `HelixHelperTest#testToIdealStateAndExternalViewShareCompactRecord`, not
here.
+ private void assertReadsAreCompact() {
+ PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
+ for (String table : List.of(OFFLINE_TABLE_NAME, REALTIME_TABLE_NAME)) {
+ for (String path : List.of(keyBuilder.idealStates(table).getPath(),
keyBuilder.externalView(table).getPath())) {
+ ZNRecord compact =
+
_compactHelixManager.getHelixDataAccessor().getBaseDataAccessor().get(path,
null, AccessOption.PERSISTENT);
+ assertTrue(HelixHelper.isCompact(compact), path);
+ ZNRecord standard =
+
_defaultHelixManager.getHelixDataAccessor().getBaseDataAccessor().get(path,
null, AccessOption.PERSISTENT);
+ assertFalse(HelixHelper.isCompact(standard), path);
+ }
+ }
+ }
+
+ /// Returns what the routing manager exposes for both tables, in a form that
compares by value.
+ private static Map<String, Object> snapshot(BrokerRoutingManager
routingManager) {
+ Map<String, Object> snapshot = new LinkedHashMap<>();
+ for (String table : List.of(OFFLINE_TABLE_NAME, REALTIME_TABLE_NAME)) {
+ snapshot.put(table + ".exists", routingManager.routingExists(table));
+ snapshot.put(table + ".disabled", routingManager.isTableDisabled(table));
+ snapshot.put(table + ".servingInstances", new
TreeSet<>(routingManager.getServingInstances(table)));
+ TablePartitionInfo partitionInfo =
routingManager.getTablePartitionInfo(table);
+ if (partitionInfo != null) {
+ List<List<String>> segmentsByPartition = new ArrayList<>();
+ for (List<String> segments : partitionInfo.getSegmentsByPartition()) {
+ segmentsByPartition.add(sorted(segments));
+ }
+ snapshot.put(table + ".partitions", partitionInfo.getPartitionColumn()
+ "/"
+ + partitionInfo.getNumPartitions() + "/" + segmentsByPartition);
+ }
+ TablePartitionReplicatedServersInfo replicatedServersInfo =
+ routingManager.getTablePartitionReplicatedServersInfo(table);
+ if (replicatedServersInfo != null) {
+ List<String> partitions = new ArrayList<>();
+ for (TablePartitionReplicatedServersInfo.PartitionInfo info :
replicatedServersInfo.getPartitionInfoMap()) {
+ partitions.add(info == null ? "null"
+ : new TreeSet<>(info._fullyReplicatedServers) + ":" +
sorted(info._segments));
+ }
+ snapshot.put(table + ".replicatedServers",
+ partitions + "/" +
sorted(replicatedServersInfo.getSegmentsWithInvalidPartition()));
+ }
+ for (String query : QUERIES) {
+ BrokerRequest brokerRequest =
CalciteSqlCompiler.compileToBrokerRequest(query);
+ for (int requestId = 0; requestId < 4; requestId++) {
+ snapshot.put(table + "." + query + "." + requestId,
+ routing(routingManager.getRoutingTable(brokerRequest, table,
requestId)));
+ }
+ }
+ }
+ TimeBoundaryInfo timeBoundaryInfo =
routingManager.getTimeBoundaryInfo(OFFLINE_TABLE_NAME);
+ snapshot.put(OFFLINE_TABLE_NAME + ".timeBoundary",
+ timeBoundaryInfo == null ? null : timeBoundaryInfo.getTimeColumn() +
"=" + timeBoundaryInfo.getTimeValue());
+ return snapshot;
+ }
+
+ private static Map<String, Object> routing(RoutingTable routingTable) {
+ Map<String, Object> routing = new LinkedHashMap<>();
+ Map<String, String> servers = new TreeMap<>();
+ for (Map.Entry<ServerInstance, SegmentsToQuery> entry :
routingTable.getServerInstanceToSegmentsMap().entrySet()) {
+ servers.put(entry.getKey().getInstanceId(),
+ sorted(entry.getValue().getSegments()) + "+" +
sorted(entry.getValue().getOptionalSegments()));
+ }
+ routing.put("servers", servers);
+ routing.put("unavailable", sorted(routingTable.getUnavailableSegments()));
+ routing.put("numPruned", routingTable.getNumPrunedSegments());
+ return routing;
+ }
+
+ private static List<String> sorted(List<String> list) {
+ List<String> sorted = new ArrayList<>(list);
+ sorted.sort(null);
+ return sorted;
+ }
+}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializer.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializer.java
new file mode 100644
index 00000000000..a66f5c524c3
--- /dev/null
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializer.java
@@ -0,0 +1,468 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.utils.helix;
+
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParseException;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.zip.GZIPInputStream;
+import javax.annotation.Nullable;
+import org.apache.helix.HelixDefinedState;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.util.GZipCompressionUtil;
+import org.apache.helix.zookeeper.zkclient.serialize.ZkSerializer;
+import
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/// A [ZkSerializer] that reads gzip-compressed ideal state and external view
znodes with less allocation than
+/// [ZNRecordSerializer], and returns an equal [ZNRecord].
+///
+/// Only gzip-compressed znodes use the compact parser. A plain (uncompressed)
znode goes to [ZNRecordSerializer]
+/// unchanged. Pinot compresses large table ideal states, so plain znodes are
the small ones, and there the compact
+/// parser does not pay off: in a size sweep it allocated about 1x to 3.5x as
much and used 1x to 1.5x the CPU of
+/// [ZNRecordSerializer] up to about 200 segments. On gzip-compressed znodes
of any size it allocated about 0.06x to
+/// 0.3x as much and used about 0.25x to 0.75x the CPU.
+///
+/// Compared with [ZNRecordSerializer#deserialize], for a gzip-compressed
znode:
+/// - It streams a gzip-compressed znode through a [GZIPInputStream] straight
into the JSON parser. It does not
+/// inflate the whole payload into a growing byte array first.
+/// - Its [JsonFactory] does not canonicalize or intern field names, so the
parser does not build a symbol table with
+/// one entry per segment.
+/// - Instance names (the keys of the instance-state maps and the values of
the list fields) are interned once per
+/// parse. Routing keys its server maps by interned instance ids, so lookups
keep their identity fast path.
+/// - State values come from a fixed table of the known segment states. An
unknown state is interned.
+/// - Segment names (the keys of the map fields) are interned, so they keep
their identity with the names that the
+/// routing components already hold.
+/// - Equal instance-state maps within one znode share one
[ImmutableSortedArrayMap]. The outer segment map is an
+/// [ImmutableSortedArrayMap] too.
+///
+/// The map fields of a record from the compact parser are immutable: any
mutation throws
+/// [UnsupportedOperationException]. Simple fields and list fields are mutable
[LinkedHashMap]s, as with
+/// [ZNRecordSerializer]. Scalar setters of the record (version, creation
time, etc.) work as usual.
+///
+/// When the compact parser fails, the serializer logs a warning and falls
back to [ZNRecordSerializer]. [#serialize]
+/// always delegates to [ZNRecordSerializer].
+///
+/// [HelixHelper#toIdealState] and [HelixHelper#toExternalView] wrap a record
from the compact parser without the copy
+/// of the outer segment map that `new IdealState(record)` and `new
ExternalView(record)` make. Those constructors
+/// share the inner instance-state maps, so an
[org.apache.helix.model.IdealState] or
+/// [org.apache.helix.model.ExternalView] built from a compact record has
immutable instance-state maps either way.
+/// A record from [ZNRecordSerializer] (a plain znode or the fallback) still
goes through those constructors.
+///
+/// Thread-safety: the serializer is stateless apart from the shared
[JsonFactory], which is thread-safe. Each call
+/// uses its own parser and dictionaries, so concurrent calls are safe.
+public class CompactZNRecordSerializer implements ZkSerializer {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(CompactZNRecordSerializer.class);
+
+ private static final JsonFactory JSON_FACTORY = JsonFactory.builder()
+ .disable(JsonFactory.Feature.CANONICALIZE_FIELD_NAMES)
+ .disable(JsonFactory.Feature.INTERN_FIELD_NAMES)
+ .build();
+
+ /// Jackson has its own input buffer, so the inflater only needs a small one.
+ private static final int GZIP_BUFFER_SIZE = 2 * 1024;
+ /// Small, so that the fixed cost per parse stays low; the arrays grow
geometrically.
+ private static final int INITIAL_SEGMENT_CAPACITY = 16;
+ private static final int INITIAL_STATE_MAP_TABLE_CAPACITY = 16;
+ private static final int INITIAL_INSTANCE_CAPACITY = 8;
+
+ /// The states a segment replica can have. The constants are string literals
or enum constant names, so they are
+ /// already interned.
+ private static final String[] KNOWN_STATES = {
+ SegmentStateModel.ONLINE, SegmentStateModel.CONSUMING,
SegmentStateModel.OFFLINE, SegmentStateModel.ERROR,
+ HelixDefinedState.DROPPED.name()
+ };
+
+ private final ZNRecordSerializer _delegate;
+
+ public CompactZNRecordSerializer() {
+ this(new ZNRecordSerializer());
+ }
+
+ public CompactZNRecordSerializer(ZNRecordSerializer delegate) {
+ _delegate = delegate;
+ }
+
+ @Override
+ public byte[] serialize(Object data) {
+ return _delegate.serialize(data);
+ }
+
+ @Nullable
+ @Override
+ public Object deserialize(byte[] bytes) {
+ if (bytes == null || bytes.length == 0) {
+ // Reading a parent or empty node, same as ZNRecordSerializer
+ return null;
+ }
+ if (!GZipCompressionUtil.isCompressed(bytes)) {
+ // Plain znodes are small; the default deserializer is cheaper there
(see the class documentation)
+ return _delegate.deserialize(bytes);
+ }
+ try {
+ return deserializeCompact(bytes);
+ } catch (Exception e) {
+ LOGGER.warn("Caught exception while deserializing ZNRecord with the
compact reader, falling back to the "
+ + "default deserializer", e);
+ return _delegate.deserialize(bytes);
+ }
+ }
+
+ /// Parses the given gzip-compressed znode content with the compact parser.
The content must be gzip-compressed;
+ /// [#deserialize] sends plain content to [ZNRecordSerializer] before
calling this method. Throws on malformed input;
+ /// [#deserialize] then falls back.
+ static ZNRecord deserializeCompact(byte[] bytes)
+ throws IOException {
+ assert GZipCompressionUtil.isCompressed(bytes);
+ GZIPInputStream gzipInputStream = new GZIPInputStream(new
ByteArrayInputStream(bytes), GZIP_BUFFER_SIZE);
+ try (JsonParser parser = JSON_FACTORY.createParser(gzipInputStream)) {
+ ZNRecord record = new Reader(parser).readRecord();
+ // The parser stops at the end of the record. GZIPInputStream verifies
the CRC-32 and the size in the gzip
+ // trailer only when it reaches the end of the stream, so read the rest
of the stream (as
+ // GZipCompressionUtil.uncompress does). A bad trailer throws, and
deserialize falls back.
+ gzipInputStream.transferTo(OutputStream.nullOutputStream());
+ return record;
+ }
+ }
+
+ /// The state of one parse: the parser, the dictionaries, and scratch
buffers. Not thread-safe; one per call.
+ private static final class Reader {
+ private final JsonParser _parser;
+ private final HashMap<String, String> _instanceNames = new HashMap<>();
+ private final StateMapTable _stateMaps = new StateMapTable();
+ private String[] _instanceKeys = new String[INITIAL_INSTANCE_CAPACITY];
+ private Object[] _instanceStates = new Object[INITIAL_INSTANCE_CAPACITY];
+
+ Reader(JsonParser parser) {
+ _parser = parser;
+ }
+
+ ZNRecord readRecord()
+ throws IOException {
+ expect(_parser.nextToken(), JsonToken.START_OBJECT);
+ String id = null;
+ Map<String, String> simpleFields = null;
+ Map<String, List<String>> listFields = null;
+ Map<String, Map<String, String>> mapFields = null;
+ byte[] rawPayload = null;
+ boolean hasSimpleFields = false;
+ boolean hasListFields = false;
+ boolean hasMapFields = false;
+ JsonToken token;
+ while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) {
+ String name = _parser.currentName();
+ token = _parser.nextToken();
+ switch (name) {
+ case "id":
+ id = token == JsonToken.VALUE_NULL ? null : scalarText(token);
+ break;
+ case "simpleFields":
+ hasSimpleFields = true;
+ simpleFields = token == JsonToken.VALUE_NULL ? null :
readSimpleFields(token);
+ break;
+ case "listFields":
+ hasListFields = true;
+ listFields = token == JsonToken.VALUE_NULL ? null :
readListFields(token);
+ break;
+ case "mapFields":
+ hasMapFields = true;
+ mapFields = token == JsonToken.VALUE_NULL ? null :
readMapFields(token);
+ break;
+ case "rawPayload":
+ rawPayload = token == JsonToken.VALUE_NULL ? null :
_parser.getBinaryValue();
+ break;
+ default:
+ // ZNRecord ignores unknown properties
+ _parser.skipChildren();
+ break;
+ }
+ }
+ // Like ObjectMapper.readValue (FAIL_ON_TRAILING_TOKENS is off), content
after the record is ignored
+ expect(token, JsonToken.END_OBJECT);
+
+ ZNRecord record = new ZNRecord(id);
+ if (hasSimpleFields) {
+ record.setSimpleFields(simpleFields);
+ }
+ if (hasListFields) {
+ record.setListFields(listFields);
+ }
+ if (hasMapFields) {
+ record.setMapFields(mapFields);
+ }
+ if (rawPayload != null) {
+ record.setRawPayload(rawPayload);
+ }
+ return record;
+ }
+
+ private Map<String, String> readSimpleFields(JsonToken token)
+ throws IOException {
+ expect(token, JsonToken.START_OBJECT);
+ Map<String, String> simpleFields = new LinkedHashMap<>();
+ while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) {
+ // ZNRecordSerializer interns all field names
+ String key = _parser.currentName().intern();
+ token = _parser.nextToken();
+ simpleFields.put(key, token == JsonToken.VALUE_NULL ? null :
scalarText(token));
+ }
+ expect(token, JsonToken.END_OBJECT);
+ return simpleFields;
+ }
+
+ private Map<String, List<String>> readListFields(JsonToken token)
+ throws IOException {
+ expect(token, JsonToken.START_OBJECT);
+ Map<String, List<String>> listFields = new LinkedHashMap<>();
+ while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) {
+ String key = _parser.currentName().intern();
+ token = _parser.nextToken();
+ if (token == JsonToken.VALUE_NULL) {
+ listFields.put(key, null);
+ continue;
+ }
+ expect(token, JsonToken.START_ARRAY);
+ List<String> list = new ArrayList<>();
+ while ((token = _parser.nextToken()) != JsonToken.END_ARRAY) {
+ // List fields hold instance names (e.g. the preference lists of
SEMI_AUTO resources)
+ list.add(token == JsonToken.VALUE_NULL ? null :
instanceName(scalarText(token)));
+ }
+ listFields.put(key, list);
+ }
+ expect(token, JsonToken.END_OBJECT);
+ return listFields;
+ }
+
+ private Map<String, Map<String, String>> readMapFields(JsonToken token)
+ throws IOException {
+ expect(token, JsonToken.START_OBJECT);
+ String[] segments = new String[INITIAL_SEGMENT_CAPACITY];
+ Object[] stateMaps = new Object[INITIAL_SEGMENT_CAPACITY];
+ int numSegments = 0;
+ while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) {
+ String segment = _parser.currentName().intern();
+ token = _parser.nextToken();
+ Map<String, String> stateMap = token == JsonToken.VALUE_NULL ? null :
readStateMap(token);
+ if (numSegments == segments.length) {
+ int newCapacity = segments.length * 2;
+ segments = Arrays.copyOf(segments, newCapacity);
+ stateMaps = Arrays.copyOf(stateMaps, newCapacity);
+ }
+ segments[numSegments] = segment;
+ stateMaps[numSegments] = stateMap;
+ numSegments++;
+ }
+ expect(token, JsonToken.END_OBJECT);
+ return ImmutableSortedArrayMap.sortAndWrap(segments, stateMaps,
numSegments);
+ }
+
+ private Map<String, String> readStateMap(JsonToken token)
+ throws IOException {
+ expect(token, JsonToken.START_OBJECT);
+ int size = 0;
+ while ((token = _parser.nextToken()) == JsonToken.FIELD_NAME) {
+ String instance = instanceName(_parser.currentName());
+ token = _parser.nextToken();
+ String state = token == JsonToken.VALUE_NULL ? null :
stateValue(token);
+ if (size == _instanceKeys.length) {
+ _instanceKeys = Arrays.copyOf(_instanceKeys, size * 2);
+ _instanceStates = Arrays.copyOf(_instanceStates, size * 2);
+ }
+ _instanceKeys[size] = instance;
+ _instanceStates[size] = state;
+ size++;
+ }
+ expect(token, JsonToken.END_OBJECT);
+ size = sortInPlace(_instanceKeys, _instanceStates, size);
+ return _stateMaps.intern(_instanceKeys, _instanceStates, size);
+ }
+
+ /// Returns the canonical (interned) instance name equal to the given name.
+ private String instanceName(String name) {
+ String canonical = _instanceNames.get(name);
+ if (canonical == null) {
+ canonical = name.intern();
+ _instanceNames.put(canonical, canonical);
+ }
+ return canonical;
+ }
+
+ /// Returns the canonical state for the current value token without
allocating for the known states.
+ private String stateValue(JsonToken token)
+ throws IOException {
+ if (token != JsonToken.VALUE_STRING) {
+ return scalarText(token).intern();
+ }
+ char[] chars = _parser.getTextCharacters();
+ int offset = _parser.getTextOffset();
+ int length = _parser.getTextLength();
+ for (String knownState : KNOWN_STATES) {
+ if (matches(knownState, chars, offset, length)) {
+ return knownState;
+ }
+ }
+ return new String(chars, offset, length).intern();
+ }
+
+ /// Returns the text of a scalar token. Like Jackson's `Map<String,
String>` binding, numbers and booleans become
+ /// their text, and objects or arrays fail.
+ private String scalarText(JsonToken token)
+ throws IOException {
+ if (!token.isScalarValue()) {
+ throw new JsonParseException(_parser, "Expected a scalar value but
got: " + token);
+ }
+ return _parser.getText();
+ }
+
+ private void expect(@Nullable JsonToken actual, JsonToken expected)
+ throws JsonParseException {
+ if (actual != expected) {
+ throw new JsonParseException(_parser, "Expected " + expected + " but
got: " + actual);
+ }
+ }
+ }
+
+ private static boolean matches(String s, char[] chars, int offset, int
length) {
+ if (s.length() != length) {
+ return false;
+ }
+ for (int i = 0; i < length; i++) {
+ if (s.charAt(i) != chars[offset + i]) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ /// Sorts the first `size` entries of the parallel arrays by key with an
insertion sort (instance-state maps hold a
+ /// handful of replicas). When a key appears more than once, the last value
wins. Returns the number of distinct
+ /// keys, which are compacted at the start of the arrays.
+ static int sortInPlace(String[] keys, Object[] values, int size) {
+ int numUnique = 0;
+ for (int i = 0; i < size; i++) {
+ String key = keys[i];
+ Object value = values[i];
+ int low = 0;
+ int high = numUnique - 1;
+ int found = -1;
+ while (low <= high) {
+ int mid = (low + high) >>> 1;
+ int cmp = keys[mid].compareTo(key);
+ if (cmp < 0) {
+ low = mid + 1;
+ } else if (cmp > 0) {
+ high = mid - 1;
+ } else {
+ found = mid;
+ break;
+ }
+ }
+ if (found >= 0) {
+ values[found] = value;
+ continue;
+ }
+ System.arraycopy(keys, low, keys, low + 1, numUnique - low);
+ System.arraycopy(values, low, values, low + 1, numUnique - low);
+ keys[low] = key;
+ values[low] = value;
+ numUnique++;
+ }
+ return numUnique;
+ }
+
+ /// Deduplicates instance-state maps by content within one parse. Keys and
values passed in are canonical (interned
+ /// or from the known state table), so content equality is reference
equality element by element.
+ private static final class StateMapTable {
+ private ImmutableSortedArrayMap<?>[] _maps = new
ImmutableSortedArrayMap<?>[INITIAL_STATE_MAP_TABLE_CAPACITY];
+ private int[] _hashes = new int[INITIAL_STATE_MAP_TABLE_CAPACITY];
+ private int _size;
+
+ @SuppressWarnings("unchecked")
+ Map<String, String> intern(String[] keys, Object[] values, int size) {
+ int hash = 1;
+ for (int i = 0; i < size; i++) {
+ hash = 31 * hash + keys[i].hashCode();
+ hash = 31 * hash + (values[i] == null ? 0 : values[i].hashCode());
+ }
+ int mask = _maps.length - 1;
+ int slot = mix(hash) & mask;
+ while (_maps[slot] != null) {
+ if (_hashes[slot] == hash && sameContent(_maps[slot], keys, values,
size)) {
+ return (Map<String, String>) _maps[slot];
+ }
+ slot = (slot + 1) & mask;
+ }
+ ImmutableSortedArrayMap<String> map = size == 0 ?
ImmutableSortedArrayMap.empty()
+ : new ImmutableSortedArrayMap<>(Arrays.copyOf(keys, size),
Arrays.copyOf(values, size));
+ _maps[slot] = map;
+ _hashes[slot] = hash;
+ if (++_size * 2 > _maps.length) {
+ resize();
+ }
+ return map;
+ }
+
+ private static boolean sameContent(ImmutableSortedArrayMap<?> map,
String[] keys, Object[] values, int size) {
+ if (map.size() != size) {
+ return false;
+ }
+ for (int i = 0; i < size; i++) {
+ if (map.keyAt(i) != keys[i] || map.valueAt(i) != values[i]) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ private void resize() {
+ ImmutableSortedArrayMap<?>[] oldMaps = _maps;
+ int[] oldHashes = _hashes;
+ _maps = new ImmutableSortedArrayMap<?>[oldMaps.length * 2];
+ _hashes = new int[oldMaps.length * 2];
+ int mask = _maps.length - 1;
+ for (int i = 0; i < oldMaps.length; i++) {
+ if (oldMaps[i] != null) {
+ int slot = mix(oldHashes[i]) & mask;
+ while (_maps[slot] != null) {
+ slot = (slot + 1) & mask;
+ }
+ _maps[slot] = oldMaps[i];
+ _hashes[slot] = oldHashes[i];
+ }
+ }
+ }
+
+ private static int mix(int hash) {
+ return hash ^ (hash >>> 16);
+ }
+ }
+}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
index e8636bea991..ac102349d0d 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
@@ -34,6 +34,7 @@ import org.apache.helix.HelixAdmin;
import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixManager;
import org.apache.helix.HelixManagerFactory;
+import org.apache.helix.HelixProperty;
import org.apache.helix.InstanceType;
import org.apache.helix.PropertyKey.Builder;
import org.apache.helix.model.ExternalView;
@@ -94,6 +95,46 @@ public class HelixHelper {
(ZNRecord)
ZN_RECORD_SERIALIZER.deserialize(ZN_RECORD_SERIALIZER.serialize(idealState.getRecord())));
}
+ /// Returns whether the given record came from the compact parser of
[CompactZNRecordSerializer], that is, its map
+ /// fields are an [ImmutableSortedArrayMap]. A record from
[ZNRecordSerializer] returns false, including the records
+ /// that the compact serializer delegates to it (plain znodes and the
fallback).
+ public static boolean isCompact(ZNRecord record) {
+ return record.getMapFields() instanceof ImmutableSortedArrayMap;
+ }
+
+ /// Returns an [IdealState] for the given record. When the record [is
compact][#isCompact], the ideal state takes
+ /// over the maps of the record instead of copying them (see `adoptRecord`),
so the caller must own the record and
+ /// must not use it afterwards. Otherwise this is `new IdealState(record)`.
+ public static IdealState toIdealState(ZNRecord record) {
+ return isCompact(record) ? adoptRecord(record, new
IdealState(record.getId())) : new IdealState(record);
+ }
+
+ /// Returns an [ExternalView] for the given record. Same contract as
[#toIdealState].
+ public static ExternalView toExternalView(ZNRecord record) {
+ return isCompact(record) ? adoptRecord(record, new
ExternalView(record.getId())) : new ExternalView(record);
+ }
+
+ /// Makes the empty record of the given property hold the content of the
source record, as
+ /// `ZNRecord(ZNRecord, String)` (called by `HelixProperty(ZNRecord)`) does,
but by reference instead of by copy:
+ /// the simple, list and map fields and the raw payload are shared. The
scalar fields that the copy constructor
+ /// copies (version, creation time, modified time, ephemeral owner) are
copied, and the property [HelixProperty.Stat]
+ /// is set from them, as `HelixProperty(ZNRecord)` does. Like the copy
constructor, the delta list is not copied.
+ /// Returns the given property.
+ private static <T extends HelixProperty> T adoptRecord(ZNRecord source, T
property) {
+ ZNRecord target = property.getRecord();
+ target.setSimpleFields(source.getSimpleFields());
+ target.setListFields(source.getListFields());
+ target.setMapFields(source.getMapFields());
+ target.setRawPayload(source.getRawPayload());
+ target.setVersion(source.getVersion());
+ target.setCreationTime(source.getCreationTime());
+ target.setModifiedTime(source.getModifiedTime());
+ target.setEphemeralOwner(source.getEphemeralOwner());
+ property.setStat(new HelixProperty.Stat(source.getVersion(),
source.getCreationTime(), source.getModifiedTime(),
+ source.getEphemeralOwner()));
+ return property;
+ }
+
public static IdealState updateIdealState(HelixManager helixManager, String
resourceName,
Function<IdealState, IdealState> updater) {
return IDEAL_STATE_GROUP_COMMIT.commit(helixManager, resourceName, updater,
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMap.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMap.java
new file mode 100644
index 00000000000..13056a8a851
--- /dev/null
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMap.java
@@ -0,0 +1,562 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.utils.helix;
+
+import java.util.AbstractCollection;
+import java.util.AbstractMap;
+import java.util.AbstractSet;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Objects;
+import java.util.Set;
+import java.util.SortedMap;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+
+
+/// An immutable [SortedMap] with [String] keys in natural order, backed by a
sorted key array and a parallel value
+/// array.
+///
+/// [CompactZNRecordSerializer] uses it for the map fields of ideal states and
external views: one instance per
+/// distinct instance-state map, and one for the outer segment map. Lookups
use binary search with an identity fast
+/// path, so interned keys compare cheaply.
+///
+/// The map follows the [Map] and [SortedMap] contracts: [#equals],
[#hashCode] and [#toString] match those of a
+/// [java.util.TreeMap] with the same content, and [#keySet], [#entrySet] and
[#values] are [AbstractSet] and
+/// [AbstractCollection] views with the standard `equals` and `hashCode`.
[#comparator] returns `null`, so
+/// `new TreeMap<>(map)` and `TreeMap.putAll` copy it in linear time. Keys are
never `null`; values can be `null`.
+///
+/// The values must be immutable too (in practice they are strings or other
instances of this class). [#hashCode]
+/// is cached, so a value that changes after the map is built makes the cached
hash code stale.
+///
+/// The map cannot change, but it differs from an unmodifiable
[java.util.TreeMap] in two ways:
+/// - [#subMap], [#headMap] and [#tailMap] return copies without the range
bounds of a view. A key outside the range
+/// of the original call is accepted by a further range call: `subMap("b",
"d").headMap("z")` returns the entries in
+/// `["b", "d")`, where a [java.util.TreeMap] view throws
[IllegalArgumentException].
+/// - Every mutator that would change the map throws
[UnsupportedOperationException]. Some mutators inherited by the
+/// views from [AbstractCollection] and [AbstractSet] return normally when
they have nothing to change, for example
+/// `keySet().addAll(List.of())`, `keySet().removeAll(List.of())` or
`values().removeIf(v -> false)`.
+///
+/// Thread-safety: the map is immutable and safe to share between threads once
published. The cached hash codes and
+/// views are computed lazily with benign races: every thread computes the
same value.
+public final class ImmutableSortedArrayMap<V> extends AbstractMap<String, V>
implements SortedMap<String, V> {
+ private static final ImmutableSortedArrayMap<?> EMPTY = new
ImmutableSortedArrayMap<>(new String[0], new Object[0]);
+
+ private final String[] _keys;
+ private final Object[] _values;
+
+ // Lazily computed, benign races. 0 means "not computed yet" (a real hash
code of 0 is recomputed each time).
+ private int _hashCode;
+ private int _keySetHashCode;
+ private KeySet _keySet;
+ private EntrySet _entrySet;
+ private Values _valuesView;
+
+ /// Wraps the given arrays without copying them. The keys must be non-null,
distinct and sorted in natural order,
+ /// and the caller must not modify either array afterwards.
+ ImmutableSortedArrayMap(String[] keys, Object[] values) {
+ assert keys.length == values.length;
+ _keys = keys;
+ _values = values;
+ }
+
+ /// Returns the empty map.
+ @SuppressWarnings("unchecked")
+ public static <V> ImmutableSortedArrayMap<V> empty() {
+ return (ImmutableSortedArrayMap<V>) EMPTY;
+ }
+
+ /// Returns an immutable sorted copy of the given map. Keys must be
non-null, and values must be immutable (see the
+ /// class documentation).
+ static <V> ImmutableSortedArrayMap<V> copyOf(Map<String, ? extends V> map) {
+ if (map instanceof ImmutableSortedArrayMap) {
+ @SuppressWarnings("unchecked")
+ ImmutableSortedArrayMap<V> immutableMap = (ImmutableSortedArrayMap<V>)
map;
+ return immutableMap;
+ }
+ int size = map.size();
+ if (size == 0) {
+ return empty();
+ }
+ String[] keys = new String[size];
+ Object[] values = new Object[size];
+ int i = 0;
+ for (Map.Entry<String, ? extends V> entry : map.entrySet()) {
+ keys[i] = Objects.requireNonNull(entry.getKey(), "Null key");
+ values[i] = entry.getValue();
+ i++;
+ }
+ return sortAndWrap(keys, values, size);
+ }
+
+ /// Sorts the first `size` entries of the given parallel arrays by key and
wraps them. When a key appears more than
+ /// once, the last value wins, like repeated [Map#put] calls. The arrays may
be modified and must not be used by the
+ /// caller afterwards.
+ static <V> ImmutableSortedArrayMap<V> sortAndWrap(String[] keys, Object[]
values, int size) {
+ if (size == 0) {
+ return empty();
+ }
+ boolean strictlySorted = true;
+ for (int i = 1; i < size; i++) {
+ if (keys[i - 1].compareTo(keys[i]) >= 0) {
+ strictlySorted = false;
+ break;
+ }
+ }
+ if (strictlySorted) {
+ if (keys.length == size) {
+ return new ImmutableSortedArrayMap<>(keys, values);
+ }
+ return new ImmutableSortedArrayMap<>(Arrays.copyOf(keys, size),
Arrays.copyOf(values, size));
+ }
+ // Stable sort, so the last occurrence of a duplicated key stays the last
in its run.
+ Integer[] order = new Integer[size];
+ for (int i = 0; i < size; i++) {
+ order[i] = i;
+ }
+ Arrays.sort(order, Comparator.comparing(i -> keys[i]));
+ String[] sortedKeys = new String[size];
+ Object[] sortedValues = new Object[size];
+ int numUnique = 0;
+ for (int i = 0; i < size; i++) {
+ int index = order[i];
+ String key = keys[index];
+ if (numUnique > 0 && sortedKeys[numUnique - 1].equals(key)) {
+ sortedValues[numUnique - 1] = values[index];
+ } else {
+ sortedKeys[numUnique] = key;
+ sortedValues[numUnique] = values[index];
+ numUnique++;
+ }
+ }
+ if (numUnique != size) {
+ sortedKeys = Arrays.copyOf(sortedKeys, numUnique);
+ sortedValues = Arrays.copyOf(sortedValues, numUnique);
+ }
+ return new ImmutableSortedArrayMap<>(sortedKeys, sortedValues);
+ }
+
+ /// Returns the index of the key, or `-(insertionPoint + 1)` when absent.
+ private int indexOf(String key) {
+ int low = 0;
+ int high = _keys.length - 1;
+ while (low <= high) {
+ int mid = (low + high) >>> 1;
+ String midKey = _keys[mid];
+ int cmp = midKey == key ? 0 : midKey.compareTo(key);
+ if (cmp < 0) {
+ low = mid + 1;
+ } else if (cmp > 0) {
+ high = mid - 1;
+ } else {
+ return mid;
+ }
+ }
+ return -(low + 1);
+ }
+
+ /// Returns the index of the first key that is greater than or equal to the
given key.
+ private int lowerBound(String key) {
+ int index = indexOf(key);
+ return index >= 0 ? index : -(index + 1);
+ }
+
+ /// Returns the key at the given position in sorted order.
+ String keyAt(int index) {
+ return _keys[index];
+ }
+
+ /// Returns the value at the given position in sorted order.
+ @SuppressWarnings("unchecked")
+ V valueAt(int index) {
+ return (V) _values[index];
+ }
+
+ @Override
+ public int size() {
+ return _keys.length;
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return _keys.length == 0;
+ }
+
+ @Override
+ public boolean containsKey(Object key) {
+ return key instanceof String && indexOf((String) key) >= 0;
+ }
+
+ @Override
+ public boolean containsValue(Object value) {
+ for (Object v : _values) {
+ if (v == value || (v != null && v.equals(value))) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ @Nullable
+ @Override
+ public V get(Object key) {
+ if (!(key instanceof String)) {
+ return null;
+ }
+ int index = indexOf((String) key);
+ return index >= 0 ? valueAt(index) : null;
+ }
+
+ @Override
+ public V getOrDefault(Object key, V defaultValue) {
+ if (!(key instanceof String)) {
+ return defaultValue;
+ }
+ int index = indexOf((String) key);
+ return index >= 0 ? valueAt(index) : defaultValue;
+ }
+
+ @Nullable
+ @Override
+ public Comparator<? super String> comparator() {
+ return null;
+ }
+
+ @Override
+ public String firstKey() {
+ if (_keys.length == 0) {
+ throw new NoSuchElementException();
+ }
+ return _keys[0];
+ }
+
+ @Override
+ public String lastKey() {
+ if (_keys.length == 0) {
+ throw new NoSuchElementException();
+ }
+ return _keys[_keys.length - 1];
+ }
+
+ @Override
+ public SortedMap<String, V> subMap(String fromKey, String toKey) {
+ Objects.requireNonNull(fromKey);
+ Objects.requireNonNull(toKey);
+ if (fromKey.compareTo(toKey) > 0) {
+ throw new IllegalArgumentException("fromKey > toKey");
+ }
+ return range(lowerBound(fromKey), lowerBound(toKey));
+ }
+
+ @Override
+ public SortedMap<String, V> headMap(String toKey) {
+ return range(0, lowerBound(Objects.requireNonNull(toKey)));
+ }
+
+ @Override
+ public SortedMap<String, V> tailMap(String fromKey) {
+ return range(lowerBound(Objects.requireNonNull(fromKey)), _keys.length);
+ }
+
+ /// Returns the entries in `[from, to)`. The result is a copy, not a view:
it has the same content as a view
+ /// because neither map can change, but it does not keep the range bounds
(see the class documentation).
+ private SortedMap<String, V> range(int from, int to) {
+ if (from == 0 && to == _keys.length) {
+ return this;
+ }
+ if (from >= to) {
+ return empty();
+ }
+ return new ImmutableSortedArrayMap<>(Arrays.copyOfRange(_keys, from, to),
Arrays.copyOfRange(_values, from, to));
+ }
+
+ @Override
+ public Set<String> keySet() {
+ KeySet keySet = _keySet;
+ if (keySet == null) {
+ keySet = new KeySet();
+ _keySet = keySet;
+ }
+ return keySet;
+ }
+
+ @Override
+ public Collection<V> values() {
+ Values values = _valuesView;
+ if (values == null) {
+ values = new Values();
+ _valuesView = values;
+ }
+ return values;
+ }
+
+ @Override
+ public Set<Map.Entry<String, V>> entrySet() {
+ EntrySet entrySet = _entrySet;
+ if (entrySet == null) {
+ entrySet = new EntrySet();
+ _entrySet = entrySet;
+ }
+ return entrySet;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (o == this) {
+ return true;
+ }
+ if (o instanceof ImmutableSortedArrayMap) {
+ ImmutableSortedArrayMap<?> that = (ImmutableSortedArrayMap<?>) o;
+ return Arrays.equals(_keys, that._keys) && Arrays.equals(_values,
that._values);
+ }
+ return super.equals(o);
+ }
+
+ @Override
+ public int hashCode() {
+ int hashCode = _hashCode;
+ if (hashCode == 0) {
+ for (int i = 0; i < _keys.length; i++) {
+ hashCode += _keys[i].hashCode() ^ Objects.hashCode(_values[i]);
+ }
+ _hashCode = hashCode;
+ }
+ return hashCode;
+ }
+
+ // Mutators. AbstractMap already throws for put; the remaining ones are
overridden so that they throw even when
+ // they have nothing to change, and so that no default method runs a
callback or reaches a mutable code path.
+
+ @Override
+ public V remove(Object key) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void putAll(Map<? extends String, ? extends V> m) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void clear() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void replaceAll(BiFunction<? super String, ? super V, ? extends V>
function) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public V putIfAbsent(String key, V value) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public boolean remove(Object key, Object value) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public boolean replace(String key, V oldValue, V newValue) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public V replace(String key, V value) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public V computeIfAbsent(String key, Function<? super String, ? extends V>
mappingFunction) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public V computeIfPresent(String key, BiFunction<? super String, ? super V,
? extends V> remappingFunction) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public V compute(String key, BiFunction<? super String, ? super V, ? extends
V> remappingFunction) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public V merge(String key, V value, BiFunction<? super V, ? super V, ?
extends V> remappingFunction) {
+ throw new UnsupportedOperationException();
+ }
+
+ /// Iterator over the positions of the map. `remove` throws
[UnsupportedOperationException] (the default).
+ private abstract class ArrayIterator<T> implements Iterator<T> {
+ private int _next;
+
+ @Override
+ public boolean hasNext() {
+ return _next < _keys.length;
+ }
+
+ @Override
+ public T next() {
+ if (_next >= _keys.length) {
+ throw new NoSuchElementException();
+ }
+ return elementAt(_next++);
+ }
+
+ abstract T elementAt(int index);
+ }
+
+ private final class KeySet extends AbstractSet<String> {
+ @Override
+ public Iterator<String> iterator() {
+ return new ArrayIterator<String>() {
+ @Override
+ String elementAt(int index) {
+ return _keys[index];
+ }
+ };
+ }
+
+ @Override
+ public int size() {
+ return _keys.length;
+ }
+
+ @Override
+ public boolean contains(Object o) {
+ return containsKey(o);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (o == this) {
+ return true;
+ }
+ if (o instanceof ImmutableSortedArrayMap<?>.KeySet) {
+ return Arrays.equals(_keys, ((ImmutableSortedArrayMap<?>.KeySet)
o).keys());
+ }
+ return super.equals(o);
+ }
+
+ @Override
+ public int hashCode() {
+ int hashCode = _keySetHashCode;
+ if (hashCode == 0) {
+ for (String key : _keys) {
+ hashCode += key.hashCode();
+ }
+ _keySetHashCode = hashCode;
+ }
+ return hashCode;
+ }
+
+ private String[] keys() {
+ return _keys;
+ }
+
+ @Override
+ public boolean remove(Object o) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void clear() {
+ throw new UnsupportedOperationException();
+ }
+ }
+
+ private final class Values extends AbstractCollection<V> {
+ @Override
+ public Iterator<V> iterator() {
+ return new ArrayIterator<V>() {
+ @Override
+ V elementAt(int index) {
+ return valueAt(index);
+ }
+ };
+ }
+
+ @Override
+ public int size() {
+ return _keys.length;
+ }
+
+ @Override
+ public boolean contains(Object o) {
+ return containsValue(o);
+ }
+
+ @Override
+ public boolean remove(Object o) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void clear() {
+ throw new UnsupportedOperationException();
+ }
+ }
+
+ private final class EntrySet extends AbstractSet<Map.Entry<String, V>> {
+ @Override
+ public Iterator<Map.Entry<String, V>> iterator() {
+ return new ArrayIterator<Map.Entry<String, V>>() {
+ @Override
+ Map.Entry<String, V> elementAt(int index) {
+ return new AbstractMap.SimpleImmutableEntry<>(_keys[index],
valueAt(index));
+ }
+ };
+ }
+
+ @Override
+ public int size() {
+ return _keys.length;
+ }
+
+ @Override
+ public boolean contains(Object o) {
+ if (!(o instanceof Map.Entry)) {
+ return false;
+ }
+ Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
+ Object key = entry.getKey();
+ if (!(key instanceof String)) {
+ return false;
+ }
+ int index = indexOf((String) key);
+ return index >= 0 && Objects.equals(_values[index], entry.getValue());
+ }
+
+ @Override
+ public boolean remove(Object o) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void clear() {
+ throw new UnsupportedOperationException();
+ }
+ }
+}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializerTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializerTest.java
new file mode 100644
index 00000000000..fd747185808
--- /dev/null
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/CompactZNRecordSerializerTest.java
@@ -0,0 +1,360 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.utils.helix;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.SortedMap;
+import java.util.TreeMap;
+import org.apache.helix.model.ExternalView;
+import org.apache.helix.model.IdealState;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.util.GZipCompressionUtil;
+import org.apache.pinot.spi.utils.CommonConstants;
+import
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that [CompactZNRecordSerializer] returns records equal to those of
[ZNRecordSerializer], with interned
+/// names, shared instance-state maps and immutable map fields.
+public class CompactZNRecordSerializerTest {
+ private static final ZNRecordSerializer DEFAULT_SERIALIZER = new
ZNRecordSerializer();
+ private static final CompactZNRecordSerializer COMPACT_SERIALIZER = new
CompactZNRecordSerializer();
+
+ @DataProvider
+ public Object[][] compression() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ private static byte[] bytes(ZNRecord record, boolean compressed)
+ throws IOException {
+ byte[] bytes = DEFAULT_SERIALIZER.serialize(record);
+ if (compressed && !GZipCompressionUtil.isCompressed(bytes)) {
+ bytes = GZipCompressionUtil.compress(bytes);
+ }
+ assertEquals(GZipCompressionUtil.isCompressed(bytes), compressed);
+ return bytes;
+ }
+
+ private static ZNRecord assertSameAsDefault(byte[] bytes)
+ throws IOException {
+ ZNRecord expected = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes);
+ // Call the compact parser directly so that a parse failure cannot hide
behind the fallback. It only takes
+ // gzip-compressed content.
+ byte[] compressed = GZipCompressionUtil.isCompressed(bytes) ? bytes :
GZipCompressionUtil.compress(bytes);
+ ZNRecord actual = CompactZNRecordSerializer.deserializeCompact(compressed);
+ // The public deserializer returns the same record for both, plain content
through the default deserializer
+ assertEquals(COMPACT_SERIALIZER.deserialize(bytes), actual);
+ assertEquals(COMPACT_SERIALIZER.deserialize(compressed), actual);
+ assertNotNull(expected);
+ assertNotNull(actual);
+ assertEquals(actual.getId(), expected.getId());
+ assertEquals(actual, expected);
+ assertEquals(actual.getSimpleFields(), expected.getSimpleFields());
+ assertEquals(actual.getListFields(), expected.getListFields());
+ assertEquals(actual.getMapFields(), expected.getMapFields());
+ assertEquals(actual.getRawPayload(), expected.getRawPayload());
+ // Same iteration order for simple and list fields; map fields are sorted
+ assertEquals(new ArrayList<>(actual.getSimpleFields().keySet()),
+ new ArrayList<>(expected.getSimpleFields().keySet()));
+ assertEquals(new ArrayList<>(actual.getListFields().keySet()), new
ArrayList<>(expected.getListFields().keySet()));
+ assertEquals(new ArrayList<>(actual.getMapFields().keySet()),
+ new ArrayList<>(new TreeMap<>(expected.getMapFields()).keySet()));
+ return actual;
+ }
+
+ static ZNRecord sampleIdealState() {
+ ZNRecord record = new ZNRecord("myTable_REALTIME");
+ record.setSimpleField("IDEAL_STATE_MODE", "CUSTOMIZED");
+ record.setSimpleField("REBALANCE_MODE", "CUSTOMIZED");
+ record.setSimpleField("HELIX_ENABLED", "true");
+
record.setSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY,
"1700000000000");
+ record.setListField("emptyList", new ArrayList<>());
+ record.setListField("myTable__0__1__20240101T0000Z", new
ArrayList<>(List.of("Server_b_1", "Server_a_1")));
+ for (int i = 0; i < 50; i++) {
+ Map<String, String> stateMap = new LinkedHashMap<>();
+ // Unsorted on purpose
+ stateMap.put("Server_" + (i % 3 + 5) + "_8098",
SegmentStateModel.ONLINE);
+ stateMap.put("Server_" + (i % 3 + 1) + "_8098", i % 7 == 0 ?
SegmentStateModel.CONSUMING : "ONLINE");
+ record.setMapField("myTable__" + (49 - i) + "__0__20240101T0000Z",
stateMap);
+ }
+ record.setMapField("myTable__unknown", Map.of("Server_1_8098",
"SOME_NEW_STATE", "Server_2_8098", "ERROR"));
+ record.setMapField("myTable__empty", Map.of());
+ return record;
+ }
+
+ @Test(dataProvider = "compression")
+ public void testRoundTripIdealState(boolean compressed)
+ throws IOException {
+ ZNRecord actual = assertSameAsDefault(bytes(sampleIdealState(),
compressed));
+
+ // IdealState and ExternalView wrappers work on the compact record
+ IdealState idealState = new IdealState(actual);
+ assertTrue(idealState.isEnabled());
+
assertEquals(idealState.getRecord().getSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY),
+ "1700000000000");
+ assertEquals(idealState.getPartitionSet().size(), 52);
+
assertEquals(idealState.getInstanceStateMap("myTable__unknown").get("Server_1_8098"),
"SOME_NEW_STATE");
+ ExternalView externalView = new ExternalView(actual);
+ assertEquals(externalView.getStateMap("myTable__empty"), Map.of());
+ assertTrue(externalView.getStateMap("myTable__0__0__20240101T0000Z")
instanceof SortedMap);
+ }
+
+ @Test(dataProvider = "compression")
+ public void testRoundTripEmptyAndMissingFields(boolean compressed)
+ throws IOException {
+ assertSameAsDefault(bytes(new ZNRecord("empty"), compressed));
+ ZNRecord withPayload = new ZNRecord("payload");
+ withPayload.setRawPayload(new byte[]{1, 2, 3, 4, 5});
+ assertSameAsDefault(bytes(withPayload, compressed));
+
+ // Hand-written JSON: missing fields, unknown properties, null and
non-string scalar values, duplicate keys
+ String json = "{\"id\":\"t_OFFLINE\",\"unknown\":{\"a\":[1,{\"b\":null}]},"
+ +
"\"simpleFields\":{\"n\":12,\"b\":true,\"nul\":null,\"dup\":\"1\",\"dup\":\"2\"},"
+ +
"\"mapFields\":{\"s2\":{\"z\":\"ONLINE\",\"a\":\"OFFLINE\",\"a\":\"ONLINE\"},\"s1\":{\"x\":null,\"y\":7},"
+ +
"\"s0\":null,\"s2\":{\"q\":\"DROPPED\"}},\"listFields\":{\"l\":[\"x\",null,\"y\"],\"nl\":null}}";
+ byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
+ assertSameAsDefault(bytes);
+ assertSameAsDefault(GZipCompressionUtil.compress(bytes));
+ // Only the id
+ assertSameAsDefault("{\"id\":\"only\"}".getBytes(StandardCharsets.UTF_8));
+ }
+
+ @Test(dataProvider = "compression")
+ public void testRoundTripEmptyRawPayload(boolean compressed)
+ throws IOException {
+ // A present but empty payload (base64 of zero bytes) reads back as an
empty array, not as null
+ ZNRecord record = new ZNRecord("emptyPayload");
+ record.setRawPayload(new byte[0]);
+ byte[] bytes = bytes(record, compressed);
+ String json = new String(DEFAULT_SERIALIZER.serialize(record),
StandardCharsets.UTF_8);
+ assertTrue(json.replace(" ", "").contains("\"rawPayload\":\"\""), json);
+ ZNRecord actual = assertSameAsDefault(bytes);
+ assertNotNull(actual.getRawPayload());
+ assertEquals(actual.getRawPayload().length, 0);
+ }
+
+ @Test
+ public void testEscapesAndTrailingContent()
+ throws IOException {
+ // JSON escapes in field names, map keys and values (quote, backslash,
control character, unicode)
+ String json =
"{\"id\":\"t\\\"1\",\"simpleFields\":{\"k\\\\\":\"v\\n\\u00e9\"},"
+ +
"\"mapFields\":{\"seg\\u0041\":{\"Server\\\"1\":\"ONLINE\",\"Server_2\":\"ON\\u004cINE\"}},"
+ + "\"listFields\":{\"l\":[\"Server\\\"1\",\"\\u2603\"]}}";
+ byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
+ ZNRecord record = assertSameAsDefault(bytes);
+ assertSameAsDefault(GZipCompressionUtil.compress(bytes));
+ assertEquals(record.getId(), "t\"1");
+ assertEquals(record.getMapField("segA"), Map.of("Server\"1", "ONLINE",
"Server_2", "ONLINE"));
+ // An escaped known state still maps to the shared constant
+ assertSame(record.getMapField("segA").get("Server_2"),
SegmentStateModel.ONLINE);
+
+ // Content after the record is ignored, as ObjectMapper.readValue does
+ byte[] trailing = "{\"id\":\"x\",\"mapFields\":{\"s\":{\"a\":\"ONLINE\"}}}
{\"id\":\"y\"}"
+ .getBytes(StandardCharsets.UTF_8);
+ assertEquals(assertSameAsDefault(trailing).getId(), "x");
+
assertEquals(assertSameAsDefault(GZipCompressionUtil.compress(trailing)).getId(),
"x");
+ }
+
+ @Test
+ public void testRoundTripRandomized()
+ throws IOException {
+ Random random = new Random(42);
+ for (int iteration = 0; iteration < 20; iteration++) {
+ ZNRecord record = new ZNRecord("table_" + iteration);
+ int numSegments = random.nextInt(200);
+ for (int i = 0; i < numSegments; i++) {
+ Map<String, String> stateMap = new LinkedHashMap<>();
+ int numReplicas = random.nextInt(5);
+ for (int r = 0; r < numReplicas; r++) {
+ String[] states = {"ONLINE", "OFFLINE", "CONSUMING", "ERROR",
"OTHER"};
+ stateMap.put("Server_" + random.nextInt(10),
states[random.nextInt(states.length)]);
+ }
+ record.setMapField("segment_" + random.nextInt(1000), stateMap);
+ }
+ assertSameAsDefault(bytes(record, random.nextBoolean()));
+ }
+ }
+
+ @Test
+ public void testInterningAndSharing()
+ throws IOException {
+ ZNRecord actual = (ZNRecord)
COMPACT_SERIALIZER.deserialize(bytes(sampleIdealState(), true));
+ Map<String, Map<String, String>> mapFields = actual.getMapFields();
+ assertTrue(mapFields instanceof ImmutableSortedArrayMap);
+ for (Map.Entry<String, Map<String, String>> entry : mapFields.entrySet()) {
+ assertSame(entry.getKey(), entry.getKey().intern());
+ Map<String, String> stateMap = entry.getValue();
+ assertTrue(stateMap instanceof ImmutableSortedArrayMap);
+ for (Map.Entry<String, String> stateEntry : stateMap.entrySet()) {
+ assertSame(stateEntry.getKey(), stateEntry.getKey().intern());
+ assertSame(stateEntry.getValue(), stateEntry.getValue().intern());
+ }
+ }
+ // Segment 49 - i: i = 48 is ONLINE on Server_5, i = 42 is CONSUMING on
Server_1
+
assertSame(mapFields.get("myTable__1__0__20240101T0000Z").get("Server_5_8098"),
SegmentStateModel.ONLINE);
+
assertSame(mapFields.get("myTable__7__0__20240101T0000Z").get("Server_1_8098"),
SegmentStateModel.CONSUMING);
+ for (String instance :
actual.getListField("myTable__0__1__20240101T0000Z")) {
+ assertSame(instance, instance.intern());
+ }
+
+ // Equal instance-state maps share one instance; different ones do not
+ // Segment 49 - i: i = 1 and i = 4 have the same servers and both ONLINE
+ Map<String, String> stateMap48 =
mapFields.get("myTable__48__0__20240101T0000Z");
+ Map<String, String> stateMap45 =
mapFields.get("myTable__45__0__20240101T0000Z");
+ assertEquals(stateMap48, stateMap45);
+ assertSame(stateMap48, stateMap45);
+ // i = 0 is CONSUMING on Server_1, i = 3 is ONLINE on Server_1
+ assertNotSame(mapFields.get("myTable__49__0__20240101T0000Z"),
mapFields.get("myTable__46__0__20240101T0000Z"));
+
+ // Map fields are immutable, scalar fields are not
+ assertThrows(UnsupportedOperationException.class, () ->
actual.setMapField("x", Map.of()));
+ assertThrows(UnsupportedOperationException.class, () ->
stateMap48.put("x", "ONLINE"));
+ actual.setVersion(7);
+ actual.setCreationTime(1L);
+ actual.setModifiedTime(2L);
+ actual.setEphemeralOwner(3L);
+ actual.setSimpleField("new", "value");
+ assertEquals(actual.getVersion(), 7);
+
+ // HelixProperty copies the outer map into a TreeMap and keeps the inner
maps
+ IdealState idealState = new IdealState(actual);
+
assertSame(idealState.getInstanceStateMap("myTable__48__0__20240101T0000Z"),
stateMap48);
+ assertEquals(idealState.getRecord().getMapFields(), mapFields);
+ }
+
+ @Test
+ public void testSortInPlace() {
+ String[] keys = {"c", "a", "b", "a", "d", null};
+ Object[] values = {"1", "2", "3", "4", "5", null};
+ int size = CompactZNRecordSerializer.sortInPlace(keys, values, 5);
+ assertEquals(size, 4);
+ assertEquals(Arrays.copyOf(keys, size), new String[]{"a", "b", "c", "d"});
+ assertEquals(Arrays.copyOf(values, size), new Object[]{"4", "3", "1",
"5"});
+ }
+
+ @Test
+ public void testMalformedInputFallsBack()
+ throws IOException {
+ // Both deserializers return null for malformed JSON
+ byte[] malformed =
"{\"id\":\"x\",\"mapFields\":{\"s\":[1]}}".getBytes(StandardCharsets.UTF_8);
+ assertNull(DEFAULT_SERIALIZER.deserialize(malformed));
+ assertNull(COMPACT_SERIALIZER.deserialize(malformed));
+ // Gzip-compressed: the compact parser fails and falls back to the default
deserializer
+ byte[] compressedMalformed = GZipCompressionUtil.compress(malformed);
+ assertThrows(IOException.class, () ->
CompactZNRecordSerializer.deserializeCompact(compressedMalformed));
+ assertNull(DEFAULT_SERIALIZER.deserialize(compressedMalformed));
+ assertNull(COMPACT_SERIALIZER.deserialize(compressedMalformed));
+ assertNull(COMPACT_SERIALIZER.deserialize("not
json".getBytes(StandardCharsets.UTF_8)));
+ assertNull(COMPACT_SERIALIZER.deserialize(new byte[0]));
+ assertNull(COMPACT_SERIALIZER.deserialize(null));
+ }
+
+ @Test
+ public void testUnsupportedInputFallsBackToValidRecord()
+ throws IOException {
+ // The default deserializer reads a byte array written as a JSON array of
numbers; the compact parser only reads
+ // base64 text. The compact parser fails, and the fallback returns the
record of the default deserializer.
+ byte[] bytes = GZipCompressionUtil.compress(
+
"{\"id\":\"x\",\"mapFields\":{\"s\":{\"a\":\"ONLINE\"}},\"rawPayload\":[1,2,3]}"
+ .getBytes(StandardCharsets.UTF_8));
+ assertThrows(IOException.class, () ->
CompactZNRecordSerializer.deserializeCompact(bytes));
+ ZNRecord expected = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes);
+ assertNotNull(expected);
+ assertEquals(expected.getRawPayload(), new byte[]{1, 2, 3});
+ ZNRecord actual = (ZNRecord) COMPACT_SERIALIZER.deserialize(bytes);
+ assertNotNull(actual);
+ assertEquals(actual, expected);
+ assertEquals(actual.getMapFields(), expected.getMapFields());
+ assertEquals(actual.getRawPayload(), expected.getRawPayload());
+ assertFalse(HelixHelper.isCompact(actual));
+ }
+
+ @Test
+ public void testCorruptGzipTrailerFallsBack()
+ throws IOException {
+ byte[] bytes = bytes(sampleIdealState(), true);
+ assertNotNull(CompactZNRecordSerializer.deserializeCompact(bytes));
+ // The gzip trailer is the CRC-32 (4 bytes) and the uncompressed size (4
bytes). The deflate data stays valid.
+ byte[] badCrc = bytes.clone();
+ badCrc[badCrc.length - 8] ^= 0x01;
+ assertThrows(IOException.class, () ->
CompactZNRecordSerializer.deserializeCompact(badCrc));
+ assertNull(DEFAULT_SERIALIZER.deserialize(badCrc));
+ assertNull(COMPACT_SERIALIZER.deserialize(badCrc));
+ byte[] badSize = bytes.clone();
+ badSize[badSize.length - 1] ^= 0x01;
+ assertThrows(IOException.class, () ->
CompactZNRecordSerializer.deserializeCompact(badSize));
+ assertNull(COMPACT_SERIALIZER.deserialize(badSize));
+ }
+
+ @Test(dataProvider = "compression")
+ public void testOnlyGzipUsesCompactParser(boolean compressed)
+ throws IOException {
+ byte[] bytes = bytes(sampleIdealState(), compressed);
+ ZNRecord expected = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes);
+ ZNRecord actual = (ZNRecord) COMPACT_SERIALIZER.deserialize(bytes);
+ assertNotNull(actual);
+ assertEquals(actual, expected);
+ assertEquals(actual.getMapFields(), expected.getMapFields());
+ assertEquals(HelixHelper.isCompact(actual), compressed);
+ Map<String, String> stateMap =
actual.getMapField("myTable__0__0__20240101T0000Z");
+ if (compressed) {
+ assertTrue(actual.getMapFields() instanceof ImmutableSortedArrayMap);
+ assertTrue(stateMap instanceof ImmutableSortedArrayMap);
+ } else {
+ // A plain znode goes to the default deserializer: same map types,
mutable maps
+ assertEquals(actual.getMapFields().getClass(),
expected.getMapFields().getClass());
+ assertEquals(stateMap.getClass(),
expected.getMapField("myTable__0__0__20240101T0000Z").getClass());
+ assertFalse(stateMap instanceof ImmutableSortedArrayMap);
+ stateMap.put("Server_new_8098", SegmentStateModel.ONLINE);
+ actual.setMapField("newSegment", new TreeMap<>());
+ assertEquals(actual.getMapFields().size(),
expected.getMapFields().size() + 1);
+ }
+ }
+
+ @Test
+ public void testSerializeDelegates()
+ throws IOException {
+ ZNRecord record = sampleIdealState();
+ byte[] expected = DEFAULT_SERIALIZER.serialize(record);
+ assertEquals(COMPACT_SERIALIZER.serialize(record), expected);
+ // A compact record serializes to the same bytes as an equal record from
the default deserializer
+ ZNRecord compact =
CompactZNRecordSerializer.deserializeCompact(GZipCompressionUtil.compress(expected));
+ ZNRecord standard = (ZNRecord) DEFAULT_SERIALIZER.deserialize(expected);
+ assertTrue(HelixHelper.isCompact(compact));
+
assertEquals(DEFAULT_SERIALIZER.deserialize(COMPACT_SERIALIZER.serialize(compact)),
standard);
+ }
+}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
index cef05ccfb1b..dd84cea5143 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/HelixHelperTest.java
@@ -18,18 +18,32 @@
*/
package org.apache.pinot.common.utils.helix;
+import java.io.IOException;
import java.util.Arrays;
import java.util.List;
+import java.util.TreeMap;
+import org.apache.helix.HelixProperty;
+import org.apache.helix.model.ExternalView;
+import org.apache.helix.model.IdealState;
import org.apache.helix.model.InstanceConfig;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.util.GZipCompressionUtil;
+import org.apache.pinot.spi.utils.CommonConstants;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
import static org.testng.Assert.assertTrue;
public class HelixHelperTest {
+ private static final ZNRecordSerializer DEFAULT_SERIALIZER = new
ZNRecordSerializer();
@Test
public void testUpdateHostName() {
@@ -88,4 +102,93 @@ public class HelixHelperTest {
assertTrue(HelixHelper.removeDisabledPartitions(instanceConfig));
assertTrue(instanceConfig.getDisabledPartitionsMap().isEmpty());
}
+
+ @Test
+ public void testToIdealStateAndExternalViewShareCompactRecord()
+ throws IOException {
+ ZNRecord written = CompactZNRecordSerializerTest.sampleIdealState();
+ written.setSimpleField("HELIX_ENABLED", "false");
+ written.setRawPayload(new byte[]{1, 2, 3});
+ // The compact parser only takes gzip-compressed content
+ byte[] bytes =
GZipCompressionUtil.compress(DEFAULT_SERIALIZER.serialize(written));
+
+ for (boolean externalView : new boolean[]{false, true}) {
+ ZNRecord compact = CompactZNRecordSerializer.deserializeCompact(bytes);
+ setScalars(compact);
+ assertTrue(HelixHelper.isCompact(compact));
+ ZNRecord copied = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes);
+ setScalars(copied);
+ assertFalse(HelixHelper.isCompact(copied));
+ HelixProperty expected = externalView ? new ExternalView(copied) : new
IdealState(copied);
+
+ HelixProperty actual = externalView ? HelixHelper.toExternalView(compact)
+ : HelixHelper.toIdealState(compact);
+ assertTrue(externalView ? actual instanceof ExternalView : actual
instanceof IdealState);
+ ZNRecord record = actual.getRecord();
+ // No copy: the maps and the payload are the ones of the parsed record
+ assertSame(record.getMapFields(), compact.getMapFields());
+ assertSame(record.getSimpleFields(), compact.getSimpleFields());
+ assertSame(record.getListFields(), compact.getListFields());
+ assertSame(record.getRawPayload(), compact.getRawPayload());
+ // Same content, scalar fields and stat as the copying constructor
+ assertEquals(actual.getId(), expected.getId());
+ assertEquals(record, expected.getRecord());
+ assertEquals(record.getSimpleFields(),
expected.getRecord().getSimpleFields());
+ assertEquals(record.getListFields(),
expected.getRecord().getListFields());
+ assertEquals(record.getMapFields(), expected.getRecord().getMapFields());
+ assertEquals(record.getRawPayload(),
expected.getRecord().getRawPayload());
+ assertEquals(record.getVersion(), 17);
+ assertEquals(record.getCreationTime(), 1000L);
+ assertEquals(record.getModifiedTime(), 2000L);
+ assertEquals(record.getEphemeralOwner(), 3000L);
+ assertEquals(actual.getStat(), expected.getStat());
+ assertEquals(actual.getStat().getVersion(), 17);
+ assertEquals(actual.getBucketSize(), expected.getBucketSize());
+ if (externalView) {
+ assertEquals(((ExternalView) actual).getPartitionSet(),
((ExternalView) expected).getPartitionSet());
+ assertSame(((ExternalView)
actual).getStateMap("myTable__0__0__20240101T0000Z"),
+ compact.getMapField("myTable__0__0__20240101T0000Z"));
+ } else {
+ IdealState idealState = (IdealState) actual;
+ assertFalse(idealState.isEnabled());
+ assertEquals(idealState.isEnabled(), ((IdealState)
expected).isEnabled());
+ assertEquals(idealState.getRebalanceMode(),
IdealState.RebalanceMode.CUSTOMIZED);
+ assertEquals(idealState.getPartitionSet(), ((IdealState)
expected).getPartitionSet());
+
assertEquals(idealState.getRecord().getSimpleField(CommonConstants.IdealState.HYBRID_TABLE_TIME_BOUNDARY),
+ "1700000000000");
+ }
+ // The shared outer map stays immutable
+ assertThrows(UnsupportedOperationException.class, () ->
record.getMapFields().remove("myTable__empty"));
+ }
+ }
+
+ @Test
+ public void testToIdealStateAndExternalViewCopyDefaultRecord() {
+ byte[] bytes =
DEFAULT_SERIALIZER.serialize(CompactZNRecordSerializerTest.sampleIdealState());
+ for (boolean externalView : new boolean[]{false, true}) {
+ // A record from the default deserializer (or from the fallback) goes
through the copying constructor
+ ZNRecord standard = (ZNRecord) DEFAULT_SERIALIZER.deserialize(bytes);
+ assertNotNull(standard);
+ setScalars(standard);
+ assertFalse(HelixHelper.isCompact(standard));
+ HelixProperty actual = externalView ?
HelixHelper.toExternalView(standard)
+ : HelixHelper.toIdealState(standard);
+ HelixProperty expected = externalView ? new ExternalView(standard) : new
IdealState(standard);
+ assertNotSame(actual.getRecord(), standard);
+ assertNotSame(actual.getRecord().getMapFields(),
standard.getMapFields());
+ assertTrue(actual.getRecord().getMapFields() instanceof TreeMap);
+ assertEquals(actual.getRecord(), expected.getRecord());
+ assertEquals(actual.getStat(), expected.getStat());
+ // The copy is mutable, as today
+ actual.getRecord().getMapFields().remove("myTable__empty");
+ assertTrue(standard.getMapFields().containsKey("myTable__empty"));
+ }
+ }
+
+ private static void setScalars(ZNRecord record) {
+ record.setVersion(17);
+ record.setCreationTime(1000L);
+ record.setModifiedTime(2000L);
+ record.setEphemeralOwner(3000L);
+ }
}
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMapTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMapTest.java
new file mode 100644
index 00000000000..83bc7a6b12c
--- /dev/null
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/helix/ImmutableSortedArrayMapTest.java
@@ -0,0 +1,207 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.utils.helix;
+
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Set;
+import java.util.SortedMap;
+import java.util.TreeMap;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that [ImmutableSortedArrayMap] behaves like a [TreeMap] for reads
and rejects every mutation.
+public class ImmutableSortedArrayMapTest {
+
+ private static Map<String, String> unsortedSample() {
+ Map<String, String> map = new LinkedHashMap<>();
+ map.put("Server_c_8098", "ONLINE");
+ map.put("Server_a_8098", "CONSUMING");
+ map.put("Server_b_8098", "OFFLINE");
+ map.put("Server_d_8098", null);
+ return map;
+ }
+
+ @Test
+ public void testReadsMatchTreeMap() {
+ Map<String, String> source = unsortedSample();
+ TreeMap<String, String> expected = new TreeMap<>(source);
+ ImmutableSortedArrayMap<String> map =
ImmutableSortedArrayMap.copyOf(source);
+
+ assertEquals(map.size(), expected.size());
+ assertFalse(map.isEmpty());
+ assertNull(map.comparator());
+ assertEquals(new ArrayList<>(map.keySet()), new
ArrayList<>(expected.keySet()));
+ assertEquals(new ArrayList<>(map.values()), new
ArrayList<>(expected.values()));
+ assertEquals(new ArrayList<>(map.entrySet()), new
ArrayList<>(expected.entrySet()));
+ assertEquals(map.toString(), expected.toString());
+ assertEquals(map.firstKey(), expected.firstKey());
+ assertEquals(map.lastKey(), expected.lastKey());
+ for (String key : List.of("Server_a_8098", "Server_d_8098",
"Server_e_8098", "")) {
+ assertEquals(map.get(key), expected.get(key), key);
+ assertEquals(map.containsKey(key), expected.containsKey(key), key);
+ assertEquals(map.getOrDefault(key, "x"), expected.getOrDefault(key,
"x"), key);
+ }
+ // Lookups with a non-interned equal key and with keys of another type
+ assertEquals(map.get(new String("Server_b_8098")), "OFFLINE");
+ assertNull(map.get(42));
+ assertFalse(map.containsKey(null));
+ assertNull(map.get(null));
+ assertTrue(map.containsValue("ONLINE"));
+ assertTrue(map.containsValue(null));
+ assertFalse(map.containsValue("ERROR"));
+ assertFalse(map.containsValue(42));
+ assertTrue(map.entrySet().contains(new
AbstractMap.SimpleEntry<>("Server_a_8098", "CONSUMING")));
+ assertFalse(map.entrySet().contains(new
AbstractMap.SimpleEntry<>("Server_a_8098", "ONLINE")));
+ assertTrue(map.values().contains(null));
+ }
+
+ @Test
+ public void testEqualsAndHashCodeParity() {
+ Map<String, String> source = unsortedSample();
+ TreeMap<String, String> treeMap = new TreeMap<>(source);
+ HashMap<String, String> hashMap = new HashMap<>(source);
+ ImmutableSortedArrayMap<String> map =
ImmutableSortedArrayMap.copyOf(source);
+
+ assertEquals(map, treeMap);
+ assertEquals(treeMap, map);
+ assertEquals(map, hashMap);
+ assertEquals(hashMap, map);
+ assertEquals(map.hashCode(), treeMap.hashCode());
+ assertEquals(map.hashCode(), hashMap.hashCode());
+ assertEquals(map, ImmutableSortedArrayMap.copyOf(hashMap));
+
+ Set<String> hashSet = new HashSet<>(source.keySet());
+ assertEquals(map.keySet(), hashSet);
+ assertEquals(hashSet, map.keySet());
+ assertEquals(map.keySet(), treeMap.keySet());
+ assertEquals(map.keySet().hashCode(), hashSet.hashCode());
+ assertEquals(map.entrySet(), treeMap.entrySet());
+ assertEquals(treeMap.entrySet(), map.entrySet());
+ assertEquals(map.entrySet().hashCode(), treeMap.entrySet().hashCode());
+
+ // A HashMap keyed by key sets works across implementations, as in
ReplicaGroupInstanceSelector
+ Map<Set<String>, String> byKeySet = new HashMap<>();
+ byKeySet.put(treeMap.keySet(), "value");
+ assertEquals(byKeySet.get(map.keySet()), "value");
+ byKeySet.clear();
+ byKeySet.put(map.keySet(), "value");
+ assertEquals(byKeySet.get(hashSet), "value");
+
+ Map<String, String> different = new TreeMap<>(source);
+ different.put("Server_a_8098", "ONLINE");
+ assertFalse(map.equals(different));
+ assertFalse(different.equals(map));
+ assertFalse(map.equals(ImmutableSortedArrayMap.copyOf(different)));
+ }
+
+ @Test
+ public void testSortedViews() {
+ TreeMap<String, String> expected = new TreeMap<>(unsortedSample());
+ ImmutableSortedArrayMap<String> map =
ImmutableSortedArrayMap.copyOf(expected);
+ for (String from : List.of("", "Server_a_8098", "Server_b",
"Server_c_8098", "Z")) {
+ assertEquals(map.tailMap(from), expected.tailMap(from), from);
+ assertEquals(map.headMap(from), expected.headMap(from), from);
+ for (String to : List.of("Server_b", "Server_c_8098", "Z")) {
+ if (from.compareTo(to) <= 0) {
+ SortedMap<String, String> subMap = map.subMap(from, to);
+ assertEquals(subMap, expected.subMap(from, to), from + "-" + to);
+ assertEquals(new ArrayList<>(subMap.keySet()), new
ArrayList<>(expected.subMap(from, to).keySet()));
+ }
+ }
+ }
+ assertThrows(IllegalArgumentException.class, () -> map.subMap("b", "a"));
+ assertThrows(NullPointerException.class, () -> map.headMap(null));
+ assertThrows(NoSuchElementException.class, () ->
ImmutableSortedArrayMap.empty().firstKey());
+ assertThrows(NoSuchElementException.class, () ->
ImmutableSortedArrayMap.empty().lastKey());
+ }
+
+ @Test
+ public void testDuplicateKeysLastWins() {
+ String[] keys = {"b", "a", "b", "c", "a"};
+ Object[] values = {"1", "2", "3", "4", "5"};
+ ImmutableSortedArrayMap<String> map =
ImmutableSortedArrayMap.sortAndWrap(keys, values, keys.length);
+ assertEquals(map, Map.of("a", "5", "b", "3", "c", "4"));
+ assertEquals(new ArrayList<>(map.keySet()), List.of("a", "b", "c"));
+
+ // Sorted input with a trailing capacity is trimmed
+ String[] sortedKeys = {"a", "b", null, null};
+ Object[] sortedValues = {"1", "2", null, null};
+ assertEquals(ImmutableSortedArrayMap.sortAndWrap(sortedKeys, sortedValues,
2), Map.of("a", "1", "b", "2"));
+ assertSame(ImmutableSortedArrayMap.sortAndWrap(new String[0], new
Object[0], 0), ImmutableSortedArrayMap.empty());
+ }
+
+ @Test
+ public void testMutatorsThrow() {
+ ImmutableSortedArrayMap<String> map =
ImmutableSortedArrayMap.copyOf(unsortedSample());
+ assertThrows(UnsupportedOperationException.class, () -> map.put("x", "y"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.remove("Server_a_8098"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.remove("Server_a_8098", "CONSUMING"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.putAll(Map.of("x", "y")));
+ assertThrows(UnsupportedOperationException.class, map::clear);
+ assertThrows(UnsupportedOperationException.class, () ->
map.putIfAbsent("x", "y"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.replace("Server_a_8098", "y"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.replace("Server_a_8098", "CONSUMING", "y"));
+ assertThrows(UnsupportedOperationException.class, () -> map.replaceAll((k,
v) -> v));
+ assertThrows(UnsupportedOperationException.class, () ->
map.computeIfAbsent("x", k -> "y"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.computeIfPresent("Server_a_8098", (k, v) -> v));
+ assertThrows(UnsupportedOperationException.class, () -> map.compute("x",
(k, v) -> v));
+ assertThrows(UnsupportedOperationException.class, () -> map.merge("x",
"y", (a, b) -> a));
+ assertThrows(UnsupportedOperationException.class, () ->
map.keySet().remove("Server_a_8098"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.keySet().add("x"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.keySet().clear());
+ assertThrows(UnsupportedOperationException.class, () ->
map.keySet().retainAll(Set.of()));
+ assertThrows(UnsupportedOperationException.class, () ->
map.values().remove("ONLINE"));
+ assertThrows(UnsupportedOperationException.class, () ->
map.values().clear());
+ assertThrows(UnsupportedOperationException.class, () ->
map.entrySet().clear());
+ assertThrows(UnsupportedOperationException.class,
+ () -> map.entrySet().remove(new
AbstractMap.SimpleEntry<>("Server_a_8098", "CONSUMING")));
+ assertThrows(UnsupportedOperationException.class, () ->
map.entrySet().iterator().next().setValue("x"));
+ Iterator<String> iterator = map.keySet().iterator();
+ iterator.next();
+ assertThrows(UnsupportedOperationException.class, iterator::remove);
+ assertEquals(map, new TreeMap<>(unsortedSample()));
+ }
+
+ @Test
+ public void testTreeMapCopyCompatibility() {
+ // TreeMap takes the buildFromSorted path for a SortedMap with the same
(natural) comparator, so the copy made by
+ // HelixProperty does not compare keys; comparator() returning null
enables that path. This test checks only the
+ // content and the order of the copy.
+ ImmutableSortedArrayMap<String> map =
ImmutableSortedArrayMap.copyOf(unsortedSample());
+ TreeMap<String, String> copy = new TreeMap<>();
+ copy.putAll(map);
+ assertEquals(copy, map);
+ assertEquals(new ArrayList<>(copy.keySet()), new
ArrayList<>(map.keySet()));
+ }
+}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 77a089e0908..967ab1b8dac 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -611,6 +611,14 @@ public class CommonConstants {
public static final int
DEFAULT_ROUTING_ASSIGNMENT_CHANGE_PROCESS_PARALLELISM =
Runtime.getRuntime().availableProcessors();
+ /// When enabled, the broker's spectator Helix managers read table ideal
states and external views with a
+ /// compact streaming deserializer. It decompresses and parses the znode
in one pass, interns instance and
+ /// segment names, and shares identical instance-state maps. The broker
resource keeps the default Helix
+ /// deserializer. Takes effect on broker start.
+ public static final String
CONFIG_OF_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED =
+ "pinot.broker.routing.compact.assignment.reader.enabled";
+ public static final boolean
DEFAULT_ROUTING_COMPACT_ASSIGNMENT_READER_ENABLED = false;
+
// When enabled, the broker will set a query option to ignore
SERVER_SEGMENT_MISSING errors from servers.
// This is useful to tolerate short windows where routing has not yet
reflected recently deleted segments.
public static final String CONFIG_OF_IGNORE_MISSING_SEGMENTS =
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
