Add memory optimized dimension table

pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTable.java

@@ -18,44 +18,21 @@
  */
 package org.apache.pinot.core.data.manager.offline;
 
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import org.apache.pinot.spi.data.FieldSpec;
-import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.GenericRow;
 import org.apache.pinot.spi.data.readers.PrimaryKey;
 
 
-class DimensionTable {
+public interface DimensionTable<T> {
 
-  private final Map<PrimaryKey, GenericRow> _lookupTable;
-  private final Schema _tableSchema;
-  private final List<String> _primaryKeyColumns;
+  List<String> getPrimaryKeyColumns();
 
-  DimensionTable(Schema tableSchema, List<String> primaryKeyColumns) {
-    this(tableSchema, primaryKeyColumns, new HashMap<>());
-  }
+  GenericRow get(PrimaryKey pk);
 
-  DimensionTable(Schema tableSchema, List<String> primaryKeyColumns, Map<PrimaryKey, GenericRow> lookupTable) {
-    _lookupTable = lookupTable;
-    _tableSchema = tableSchema;
-    _primaryKeyColumns = primaryKeyColumns;
-  }
+  boolean put(PrimaryKey pk, T value);
 
-  List<String> getPrimaryKeyColumns() {
-    return _primaryKeyColumns;
-  }
+  boolean isEmpty();
 
-  GenericRow get(PrimaryKey pk) {
-    return _lookupTable.get(pk);
-  }
-
-  boolean isEmpty() {
-    return _lookupTable.isEmpty();
-  }
-
-  FieldSpec getFieldSpecFor(String columnName) {
-    return _tableSchema.getFieldSpecFor(columnName);
-  }
+  FieldSpec getFieldSpecFor(String columnName);
 }

pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java

@@ -32,6 +32,8 @@
 import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader;
 import org.apache.pinot.segment.spi.ImmutableSegment;
 import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.spi.config.table.DimensionTableConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.GenericRow;
@@ -78,16 +80,31 @@ public static DimensionTableDataManager getInstanceByTableName(String tableNameW
       AtomicReferenceFieldUpdater.newUpdater(DimensionTableDataManager.class, DimensionTable.class, "_dimensionTable");
 
   private volatile DimensionTable _dimensionTable;
+  private boolean _isMemoryOptimized = true;
 
   @Override
   protected void doInit() {
     super.doInit();
     Schema schema = ZKMetadataProvider.getTableSchema(_propertyStore, _tableNameWithType);
     Preconditions.checkState(schema != null, "Failed to find schema for dimension table: %s", _tableNameWithType);
+
     List<String> primaryKeyColumns = schema.getPrimaryKeyColumns();
     Preconditions.checkState(CollectionUtils.isNotEmpty(primaryKeyColumns),
         "Primary key columns must be configured for dimension table: %s", _tableNameWithType);
-    _dimensionTable = new DimensionTable(schema, primaryKeyColumns);
+
+    TableConfig tableConfig = ZKMetadataProvider.getTableConfig(_propertyStore, _tableNameWithType);
+    if (tableConfig != null) {
+      DimensionTableConfig dimensionTableConfig = tableConfig.getDimensionTableConfig();
+      if (dimensionTableConfig != null) {
+        _isMemoryOptimized = dimensionTableConfig.isOptimizeMemory();
+      }
+    }
+
+    if (_isMemoryOptimized) {
+      _dimensionTable = new MemoryOptimizedDimensionTable(schema, primaryKeyColumns);
+    } else {
+      _dimensionTable = new FastLookupDimensionTable(schema, primaryKeyColumns);
+    }
   }
 
   @Override
@@ -125,13 +142,18 @@ private void loadLookupTable() {
     DimensionTable replacement;
     do {
       snapshot = _dimensionTable;
-      replacement = createDimensionTable();
+      if (_isMemoryOptimized) {
+        replacement = createMemOptimisedDimensionTable();
+      } else {
+        replacement = createDimensionTable();
+      }
     } while (!UPDATER.compareAndSet(this, snapshot, replacement));
   }
 
   private DimensionTable createDimensionTable() {
     Schema schema = ZKMetadataProvider.getTableSchema(_propertyStore, _tableNameWithType);
     Preconditions.checkState(schema != null, "Failed to find schema for dimension table: %s", _tableNameWithType);
+
     List<String> primaryKeyColumns = schema.getPrimaryKeyColumns();
     Preconditions.checkState(CollectionUtils.isNotEmpty(primaryKeyColumns),
         "Primary key columns must be configured for dimension table: %s", _tableNameWithType);
@@ -156,7 +178,43 @@ private DimensionTable createDimensionTable() {
           }
         }
       }
-      return new DimensionTable(schema, primaryKeyColumns, lookupTable);
+        return new FastLookupDimensionTable(schema, primaryKeyColumns, lookupTable);
+    } finally {
+      for (SegmentDataManager segmentManager : segmentManagers) {
+        releaseSegment(segmentManager);
+      }
+    }
+  }
+
+  private DimensionTable createMemOptimisedDimensionTable() {
+    Schema schema = ZKMetadataProvider.getTableSchema(_propertyStore, _tableNameWithType);
+    Preconditions.checkState(schema != null, "Failed to find schema for dimension table: %s", _tableNameWithType);
+
+    List<String> primaryKeyColumns = schema.getPrimaryKeyColumns();
+    Preconditions.checkState(CollectionUtils.isNotEmpty(primaryKeyColumns),
+        "Primary key columns must be configured for dimension table: %s", _tableNameWithType);
+
+    Map<PrimaryKey, LookupRecordLocation> lookupTable = new HashMap<>();
+    List<SegmentDataManager> segmentManagers = acquireAllSegments();
+    try {
+      for (SegmentDataManager segmentManager : segmentManagers) {
+        IndexSegment indexSegment = segmentManager.getSegment();
+        int numTotalDocs = indexSegment.getSegmentMetadata().getTotalDocs();
+        if (numTotalDocs > 0) {
+          try (PinotSegmentRecordReader recordReader = new PinotSegmentRecordReader()) {
+            recordReader.init(indexSegment);
+            for (int i = 0; i < numTotalDocs; i++) {
+              GenericRow row = new GenericRow();
+              recordReader.getRecord(i, row);
+              lookupTable.put(row.getPrimaryKey(primaryKeyColumns), new LookupRecordLocation(indexSegment, i));
+            }
+          } catch (Exception e) {
+            throw new RuntimeException(
+                "Caught exception while reading records from segment: " + indexSegment.getSegmentName());
+          }
+        }
+      }
+      return new MemoryOptimizedDimensionTable(schema, primaryKeyColumns, lookupTable);
     } finally {
       for (SegmentDataManager segmentManager : segmentManagers) {
         releaseSegment(segmentManager);

pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/FastLookupDimensionTable.java

@@ -0,0 +1,72 @@
+/**
+ * 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.core.data.manager.offline;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.data.readers.PrimaryKey;
+
+
+class FastLookupDimensionTable implements DimensionTable<GenericRow> {
+
+  private Map<PrimaryKey, GenericRow> _lookupTable;
+  private final Schema _tableSchema;
+  private final List<String> _primaryKeyColumns;
+
+  FastLookupDimensionTable(Schema tableSchema, List<String> primaryKeyColumns) {
+    this(tableSchema, primaryKeyColumns, new HashMap<>());
+  }
+
+  FastLookupDimensionTable(Schema tableSchema, List<String> primaryKeyColumns,
+      Map<PrimaryKey, GenericRow> lookupTable) {
+    _lookupTable = lookupTable;
+    _tableSchema = tableSchema;
+    _primaryKeyColumns = primaryKeyColumns;
+  }
+
+  @Override
+  public List<String> getPrimaryKeyColumns() {
+    return _primaryKeyColumns;
+  }
+
+  @Override
+  public GenericRow get(PrimaryKey pk) {
+    return _lookupTable.get(pk);
+  }
+
+  @Override
+  public boolean put(PrimaryKey pk, GenericRow value) {
+    _lookupTable.put(pk, value);
+    return true;
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return _lookupTable.isEmpty();
+  }
+
+  @Override
+  public FieldSpec getFieldSpecFor(String columnName) {
+    return _tableSchema.getFieldSpecFor(columnName);
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/LookupRecordLocation.java

@@ -0,0 +1,40 @@
+/**
+ * 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.core.data.manager.offline;
+
+import org.apache.pinot.segment.spi.IndexSegment;
+
+
+public class LookupRecordLocation {
+  private final IndexSegment _indexSegment;
+  private final int _docId;
+
+  public LookupRecordLocation(IndexSegment indexSegment, int docId) {
+    _indexSegment = indexSegment;
+    _docId = docId;
+  }
+
+  public IndexSegment getIndexSegment() {
+    return _indexSegment;
+  }
+
+  public int getDocId() {
+    return _docId;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java

@@ -0,0 +1,76 @@
+/**
+ * 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.core.data.manager.offline;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.data.readers.PrimaryKey;
+
+
+class MemoryOptimizedDimensionTable implements DimensionTable<LookupRecordLocation> {
+
+  private Map<PrimaryKey, LookupRecordLocation> _lookupTable;
+  private final Schema _tableSchema;
+  private final List<String> _primaryKeyColumns;
+
+  MemoryOptimizedDimensionTable(Schema tableSchema, List<String> primaryKeyColumns) {
+    this(tableSchema, primaryKeyColumns, new HashMap<>());
+  }
+
+  MemoryOptimizedDimensionTable(Schema tableSchema, List<String> primaryKeyColumns,
+      Map<PrimaryKey, LookupRecordLocation> lookupTable) {
+    _lookupTable = lookupTable;
+    _tableSchema = tableSchema;
+    _primaryKeyColumns = primaryKeyColumns;
+  }
+
+  @Override
+  public List<String> getPrimaryKeyColumns() {
+    return _primaryKeyColumns;
+  }
+
+  @Override
+  public GenericRow get(PrimaryKey pk) {
+    LookupRecordLocation lookupRecordLocation = _lookupTable.get(pk);
+    if (lookupRecordLocation == null) {
+      return null;
+    }
+    return lookupRecordLocation.getIndexSegment().getRecord(lookupRecordLocation.getDocId(), new GenericRow());
+  }
+
+  @Override
+  public boolean put(PrimaryKey pk, LookupRecordLocation value) {
+    _lookupTable.put(pk, value);
+    return true;
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return _lookupTable.isEmpty();
+  }
+
+  @Override
+  public FieldSpec getFieldSpecFor(String columnName) {
+    return _tableSchema.getFieldSpecFor(columnName);
+  }
+}

pinot-spi/src/main/java/org/apache/pinot/spi/config/table/DimensionTableConfig.java

@@ -0,0 +1,37 @@
+/**
+ * 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.spi.config.table;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.pinot.spi.config.BaseJsonConfig;
+
+
+public class DimensionTableConfig extends BaseJsonConfig {
+  private final boolean _optimizeMemory;
+
+  @JsonCreator
+  public DimensionTableConfig(@JsonProperty(value = "optimizeMemory", required = true) boolean optimizeMemory) {
+    _optimizeMemory = optimizeMemory;
+  }
+
+  public boolean isOptimizeMemory() {
+    return _optimizeMemory;
+  }
+}