This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang 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 46da5d2f908 Fix bytes pk for dim tables (#19501)
46da5d2f908 is described below
commit 46da5d2f908e7b06ead88d25ae60fbc2640e0f6e
Author: RajaVinoth96 <[email protected]>
AuthorDate: Fri Oct 16 05:49:05 2026 +0530
Fix bytes pk for dim tables (#19501)
---
.../org/apache/pinot/common/utils/DataSchema.java | 4 +-
.../core/common/datablock/DataBlockBuilder.java | 7 +-
.../manager/offline/DimensionTableDataManager.java | 12 +-
.../offline/MemoryOptimizedDimensionTable.java | 25 +-
.../function/LookupTransformFunction.java | 4 +-
.../DimensionTableDataBytesPrimaryKeyTest.java | 235 +++++++++++++++++
...ransformFunctionBytesPrimaryKeyQueriesTest.java | 287 +++++++++++++++++++++
pinot-core/src/test/resources/data/dimAsset.csv | 2 +
.../src/test/resources/data/dimAsset_config.json | 18 ++
.../src/test/resources/data/dimAsset_schema.json | 9 +
.../query/runtime/queries/QueryRunnerTestBase.java | 4 +
.../runtime/queries/ResourceBasedQueriesTest.java | 94 +++++--
.../src/test/resources/queries/LookupJoin.json | 154 ++++++++++-
.../segment/readers/PinotSegmentColumnReader.java | 9 +
.../segment/readers/PinotSegmentRecordReader.java | 83 +++++-
.../readers/PinotSegmentRecordReaderTest.java | 41 ++-
.../local/segment/readers/PinotSegmentUtil.java | 4 +
17 files changed, 947 insertions(+), 45 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java
index 70458405884..11d83a713e1 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/DataSchema.java
@@ -631,7 +631,9 @@ public class DataSchema {
case TIMESTAMP:
return new Timestamp((long) value);
case BYTES:
- return ((ByteArray) value).getBytes();
+ return value instanceof ByteArray
+ ? ((ByteArray) value).getBytes()
+ : (byte[]) value;
case UUID:
return UuidUtils.toUUID((ByteArray) value);
case BOOLEAN_ARRAY:
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
b/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
index 6f0c0c40d4a..e93ab968d03 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/common/datablock/DataBlockBuilder.java
@@ -149,6 +149,7 @@ public class DataBlockBuilder {
fixedSize.putInt(dictId);
break;
case BYTES:
+ value = value instanceof ByteArray ? value : new
ByteArray((byte[]) value);
setColumn(fixedSize, varSize, (ByteArray) value);
break;
case MAP:
@@ -721,7 +722,8 @@ public class DataBlockBuilder {
/// Iterate using two loops.
/// The outer loop will iterate over a maximum of configurable step rows and
check for the interruption flag,
/// calling the inner loop to process the rows without checking the
interruption flag.
- static void interruptableLoop(int start, int max, int step, InnerLoop loop)
throws IOException {
+ static void interruptableLoop(int start, int max, int step, InnerLoop loop)
+ throws IOException {
for (int i = start; i < max; i += step) {
if (Thread.currentThread().isInterrupted()) {
throw new RuntimeException("Thread interrupted while processing rows.
Rows processed so far: " + i);
@@ -732,6 +734,7 @@ public class DataBlockBuilder {
}
private interface InnerLoop {
- void run(int from, int to) throws IOException;
+ void run(int from, int to)
+ throws IOException;
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java
index 19ee7943329..91b30512cd7 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataManager.java
@@ -48,7 +48,6 @@ import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
import org.apache.pinot.spi.data.readers.PrimaryKey;
-
/// Dimension Table is a special type of OFFLINE table which is assigned to
all servers
/// in a tenant and is used to execute a LOOKUP Transform Function.
DimensionTableDataManager
/// loads the contents into a HashMap for faster access thus the size should
be small
@@ -174,7 +173,6 @@ public class DimensionTableDataManager extends
OfflineTableDataManager {
INSTANCES.remove(_tableNameWithType);
}
-
private void closeDimensionTable(DimensionTable dimensionTable) {
try {
dimensionTable.close();
@@ -228,9 +226,8 @@ public class DimensionTableDataManager extends
OfflineTableDataManager {
int numTotalDocs = indexSegment.getSegmentMetadata().getTotalDocs();
if (numTotalDocs > 0) {
try (PinotSegmentRecordReader recordReader = new
PinotSegmentRecordReader()) {
- recordReader.init(indexSegment);
+ recordReader.init(indexSegment, primaryKeyColumns);
- int[] pkIndexes =
recordReader.getIndexesForColumns(primaryKeyColumns);
int[] valIndexes = recordReader.getIndexesForColumns(valueColumns);
for (int i = 0; i < numTotalDocs; i++) {
@@ -239,7 +236,7 @@ public class DimensionTableDataManager extends
OfflineTableDataManager {
return null;
}
- Object[] primaryKey = recordReader.getRecordValues(i, pkIndexes);
+ Object[] primaryKey = recordReader.getPrimaryKeys(i);
Object[] values = recordReader.getRecordValues(i, valIndexes);
Object[] previousValue = lookupTable.put(primaryKey, values);
@@ -306,9 +303,8 @@ public class DimensionTableDataManager extends
OfflineTableDataManager {
try {
int readerIdx = recordReaders.size();
PinotSegmentRecordReader recordReader = new
PinotSegmentRecordReader();
- recordReader.init(indexSegment);
+ recordReader.init(indexSegment, primaryKeyColumns);
recordReaders.add(recordReader);
- int[] pkIndexes =
recordReader.getIndexesForColumns(primaryKeyColumns);
for (int i = 0; i < numTotalDocs; i++) {
if (_loadToken.get() != token) {
@@ -317,7 +313,7 @@ public class DimensionTableDataManager extends
OfflineTableDataManager {
return null;
}
- Object[] primaryKey = recordReader.getRecordValues(i, pkIndexes);
+ Object[] primaryKey = recordReader.getPrimaryKeys(i);
long readerIdxAndDocId = (((long) readerIdx) << 32) | (i &
0xffffffffL);
long previousValue = lookupTable.put(primaryKey,
readerIdxAndDocId);
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java
index b04d9a78051..8c02025748e 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/offline/MemoryOptimizedDimensionTable.java
@@ -19,7 +19,10 @@
package org.apache.pinot.core.data.manager.offline;
import it.unimi.dsi.fastutil.objects.Object2LongOpenCustomHashMap;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
import java.util.List;
+import java.util.Set;
import javax.annotation.Nullable;
import org.apache.pinot.segment.local.data.manager.SegmentDataManager;
import org.apache.pinot.segment.local.data.manager.TableDataManager;
@@ -28,6 +31,7 @@ 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;
+import org.apache.pinot.spi.utils.ByteArray;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -37,7 +41,7 @@ public class MemoryOptimizedDimensionTable implements
DimensionTable {
private final Object2LongOpenCustomHashMap<Object[]> _lookupTable;
private final Schema _tableSchema;
- private final List<String> _primaryKeyColumns;
+ private final Set<String> _primaryKeyColumns;
private final ThreadLocal<GenericRow> _reuse =
ThreadLocal.withInitial(GenericRow::new);
private final List<SegmentDataManager> _segmentDataManagers;
private final List<PinotSegmentRecordReader> _recordReaders;
@@ -50,7 +54,7 @@ public class MemoryOptimizedDimensionTable implements
DimensionTable {
List<PinotSegmentRecordReader> recordReaders,
TableDataManager tableDataManager) {
_tableSchema = tableSchema;
- _primaryKeyColumns = primaryKeyColumns;
+ _primaryKeyColumns = new LinkedHashSet<>(primaryKeyColumns);
_lookupTable = lookupTable;
_segmentDataManagers = segmentDataManagers;
_recordReaders = recordReaders;
@@ -59,7 +63,7 @@ public class MemoryOptimizedDimensionTable implements
DimensionTable {
@Override
public List<String> getPrimaryKeyColumns() {
- return _primaryKeyColumns;
+ return new ArrayList<>(_primaryKeyColumns);
}
@Override
@@ -93,6 +97,12 @@ public class MemoryOptimizedDimensionTable implements
DimensionTable {
PinotSegmentRecordReader recordReader = _recordReaders.get(readerIdx);
recordReader.getRecord(docId, reuse);
+ for (String column : _primaryKeyColumns) {
+ Object value = reuse.getValue(column);
+ if (value instanceof byte[]) {
+ reuse.putValue(column, new ByteArray((byte[]) value));
+ }
+ }
return reuse;
}
@@ -108,7 +118,11 @@ public class MemoryOptimizedDimensionTable implements
DimensionTable {
int readerIdx = (int) (readerIdxAndDocId >> 32);
PinotSegmentRecordReader recordReader = _recordReaders.get(readerIdx);
- return recordReader.getValue(docId, columnName);
+ Object value = recordReader.getValue(docId, columnName);
+ if (_primaryKeyColumns.contains(columnName) && value instanceof byte[]) {
+ value = new ByteArray((byte[]) value);
+ }
+ return value;
}
@Nullable
@@ -127,6 +141,9 @@ public class MemoryOptimizedDimensionTable implements
DimensionTable {
Object[] values = new Object[numColumns];
for (int i = 0; i < numColumns; i++) {
values[i] = recordReader.getValue(docId, columnNames[i]);
+ if (_primaryKeyColumns.contains(columnNames[i]) && values[i] instanceof
byte[]) {
+ values[i] = new ByteArray((byte[]) values[i]);
+ }
}
return values;
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/LookupTransformFunction.java
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/LookupTransformFunction.java
index fa7ac171116..9e41b42a401 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/LookupTransformFunction.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/LookupTransformFunction.java
@@ -377,7 +377,9 @@ public class LookupTransformFunction extends
BaseTransformFunction {
}
private void setBytesSV(int index, Object value) {
- if (value instanceof byte[]) {
+ if (value instanceof ByteArray) {
+ _bytesValuesSV[index] = ((ByteArray) value).getBytes();
+ } else if (value instanceof byte[]) {
_bytesValuesSV[index] = (byte[]) value;
} else {
_bytesValuesSV[index] = EMPTY_BYTES;
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataBytesPrimaryKeyTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataBytesPrimaryKeyTest.java
new file mode 100644
index 00000000000..3bba3da4096
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/offline/DimensionTableDataBytesPrimaryKeyTest.java
@@ -0,0 +1,235 @@
+/**
+ * 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 com.fasterxml.jackson.core.JsonProcessingException;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
+import java.util.List;
+import java.util.concurrent.Executors;
+import org.apache.commons.io.FileUtils;
+import org.apache.helix.AccessOption;
+import org.apache.helix.HelixManager;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.common.utils.config.SchemaSerDeUtils;
+import org.apache.pinot.common.utils.config.TableConfigSerDeUtils;
+import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils;
+import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
+import org.apache.pinot.segment.local.segment.index.loader.LoaderTest;
+import org.apache.pinot.segment.local.utils.SegmentLocks;
+import org.apache.pinot.segment.local.utils.SegmentOperationsThrottler;
+import org.apache.pinot.segment.local.utils.SegmentOperationsThrottlerSet;
+import org.apache.pinot.segment.local.utils.SegmentReloadSemaphore;
+import org.apache.pinot.segment.local.utils.ServerReloadJobStatusCache;
+import org.apache.pinot.segment.spi.SegmentMetadata;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.segment.spi.creator.SegmentIndexCreationDriver;
+import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
+import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
+import org.apache.pinot.spi.config.table.DimensionTableConfig;
+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.data.readers.FileFormat;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.data.readers.PrimaryKey;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.BytesUtils;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.*;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNull;
+
+
+// Tests primary key lookups for BYTES work
+public class DimensionTableDataBytesPrimaryKeyTest {
+ private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(),
LoaderTest.class.getName());
+ private static final String RAW_TABLE_NAME = "dimAsset";
+ private static final String OFFLINE_TABLE_NAME =
TableNameBuilder.OFFLINE.tableNameWithType(RAW_TABLE_NAME);
+ private static final String CSV_DATA_PATH = "data/dimAsset.csv";
+ private static final String SCHEMA_PATH = "data/dimAsset_schema.json";
+ private static final String TABLE_CONFIG_PATH = "data/dimAsset_config.json";
+ private static final SegmentOperationsThrottlerSet
SEGMENT_OPERATIONS_THROTTLER = new SegmentOperationsThrottlerSet(
+ new SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true));
+
+ private File _indexDir;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ ServerMetrics.register(mock(ServerMetrics.class));
+
+ // prepare segment data
+ URL dataPathUrl = getClass().getClassLoader().getResource(CSV_DATA_PATH);
+ URL schemaPathUrl = getClass().getClassLoader().getResource(SCHEMA_PATH);
+ URL configPathUrl =
getClass().getClassLoader().getResource(TABLE_CONFIG_PATH);
+ assertNotNull(dataPathUrl);
+ assertNotNull(schemaPathUrl);
+ assertNotNull(configPathUrl);
+ File csvFile = new File(dataPathUrl.getFile());
+ TableConfig tableConfig = createTableConfig(new
File(configPathUrl.getFile()));
+ Schema schema = createSchema(new File(schemaPathUrl.getFile()));
+
+ // create segment
+ File tableDataDir = new File(TEMP_DIR, OFFLINE_TABLE_NAME);
+
+ SegmentGeneratorConfig segmentGeneratorConfig =
+ SegmentTestUtils.getSegmentGeneratorConfig(csvFile, FileFormat.CSV,
tableDataDir, RAW_TABLE_NAME, tableConfig,
+ schema);
+ SegmentIndexCreationDriver driver = new SegmentIndexCreationDriverImpl();
+ driver.init(segmentGeneratorConfig);
+ driver.build();
+
+ String segmentName = driver.getSegmentName();
+ _indexDir = new File(tableDataDir, segmentName);
+ SegmentMetadata segmentMetadata = new SegmentMetadataImpl(_indexDir);
+ SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName);
+ segmentZKMetadata.setCrc(segmentMetadata.getCrc());
+ }
+
+ @AfterClass
+ public void tearDown() {
+ FileUtils.deleteQuietly(TEMP_DIR);
+ }
+
+ private TableConfig getTableConfig(boolean disablePreload) {
+ DimensionTableConfig dimensionTableConfig = new
DimensionTableConfig(disablePreload, false);
+ return new TableConfigBuilder(TableType.OFFLINE)
+ .setTableName("dimAsset")
+ .setDimensionTableConfig(dimensionTableConfig)
+ .build();
+ }
+
+ private Schema getSchema() {
+ return new Schema.SchemaBuilder()
+ .setSchemaName("dimAsset")
+ .addSingleValueDimension("assetId", FieldSpec.DataType.BYTES)
+ .addSingleValueDimension("assetName", FieldSpec.DataType.STRING)
+ .addSingleValueDimension("contentHash", FieldSpec.DataType.BYTES)
+ .setPrimaryKeyColumns(List.of("assetId"))
+ .build();
+ }
+
+ private DimensionTableDataManager makeTableDataManager(TableConfig
tableConfig, Schema schema)
+ throws IOException {
+ return makeTableDataManager(tableConfig, schema,
mock(ZkHelixPropertyStore.class));
+ }
+
+ private DimensionTableDataManager makeTableDataManager(TableConfig
tableConfig, Schema schema,
+ ZkHelixPropertyStore<ZNRecord> propertyStoreMock)
+ throws JsonProcessingException {
+ HelixManager helixManager = mock(HelixManager.class);
+ when(propertyStoreMock.get("/CONFIGS/TABLE/dimAsset_OFFLINE", null,
AccessOption.PERSISTENT)).thenReturn(
+ TableConfigSerDeUtils.toZNRecord(tableConfig));
+ when(propertyStoreMock.get("/SCHEMAS/dimAsset", null,
AccessOption.PERSISTENT)).thenReturn(
+ SchemaSerDeUtils.toZNRecord(schema));
+ when(helixManager.getHelixPropertyStore()).thenReturn(propertyStoreMock);
+ InstanceDataManagerConfig instanceDataManagerConfig =
mock(InstanceDataManagerConfig.class);
+
when(instanceDataManagerConfig.getInstanceDataDir()).thenReturn(TEMP_DIR.getAbsolutePath());
+ DimensionTableDataManager tableDataManager =
+
DimensionTableDataManager.createInstanceByTableName(OFFLINE_TABLE_NAME);
+ tableDataManager.init(instanceDataManagerConfig, helixManager, new
SegmentLocks(), tableConfig, schema,
+ new SegmentReloadSemaphore(1), Executors.newSingleThreadExecutor(),
null, null, SEGMENT_OPERATIONS_THROTTLER,
+ false, mock(ServerReloadJobStatusCache.class));
+ tableDataManager.start();
+ return tableDataManager;
+ }
+
+ @DataProvider(name = "preloadModes")
+ public Object[][] providePreloadModes() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ @Test(dataProvider = "preloadModes")
+ public void testLookup(boolean disablePreload)
+ throws Exception {
+ TableConfig tableConfig = getTableConfig(disablePreload);
+ Schema schema = getSchema();
+ DimensionTableDataManager tableDataManager =
makeTableDataManager(tableConfig, schema);
+
+ // try fetching data BEFORE loading segment
+ byte[] rawBytes = BytesUtils.toBytes("550e8400e29b41d4a716446655440000");
// from hex string
+ byte[] contentHashBytes =
BytesUtils.toBytes("5eb63bbbe01eeed093cb22bb8f5acdc3");
+ ByteArray keyValue = new ByteArray(rawBytes);
+ PrimaryKey key = new PrimaryKey(new Object[]{keyValue});
+ assertFalse(tableDataManager.containsKey(key));
+ assertNull(tableDataManager.lookupRow(key));
+ assertNull(tableDataManager.lookupValue(key, "assetId"));
+ assertNull(tableDataManager.lookupValue(key, "assetName"));
+ assertNull(tableDataManager.lookupValue(key, "contentHash"));
+ assertNull(tableDataManager.lookupValues(key, new String[]{"assetId",
"assetName", "contentHash"}));
+
+ tableDataManager.addSegment(ImmutableSegmentLoader.load(_indexDir, new
IndexLoadingConfig(tableConfig, schema),
+ SEGMENT_OPERATIONS_THROTTLER));
+
+ // Confirm table is loaded and available for lookup
+ assertTrue(tableDataManager.containsKey(key));
+ GenericRow row = tableDataManager.lookupRow(key);
+ assertNotNull(row, "Should return response after segment load");
+ assertEquals(row.getFieldToValueMap().size(), 3);
+ assertEquals(row.getValue("assetId"), keyValue);
+ assertEquals(row.getValue("assetName"), "logo-final-v2.svg");
+ assertEquals(row.getValue("contentHash"), contentHashBytes);
+ assertEquals(tableDataManager.lookupValue(key, "assetId"), keyValue);
+ assertEquals(tableDataManager.lookupValue(key, "assetName"),
"logo-final-v2.svg");
+ assertEquals(tableDataManager.lookupValue(key, "contentHash"),
contentHashBytes);
+ Object[] values = tableDataManager.lookupValues(key, new
String[]{"assetId", "assetName", "contentHash"});
+ assertNotNull(values);
+ assertEquals(values.length, 3);
+ assertEquals(values[0], keyValue);
+ assertEquals(values[1], "logo-final-v2.svg");
+ assertEquals(values[2], contentHashBytes);
+ }
+
+ protected static Schema createSchema(File schemaFile)
+ throws IOException {
+ InputStream inputStream = new FileInputStream(schemaFile);
+ Assert.assertNotNull(inputStream);
+ return JsonUtils.inputStreamToObject(inputStream, Schema.class);
+ }
+
+ protected static TableConfig createTableConfig(File tableConfigFile)
+ throws IOException {
+ InputStream inputStream = new FileInputStream(tableConfigFile);
+ Assert.assertNotNull(inputStream);
+ return JsonUtils.inputStreamToObject(inputStream, TableConfig.class);
+ }
+}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/queries/LookupTransformFunctionBytesPrimaryKeyQueriesTest.java
b/pinot-core/src/test/java/org/apache/pinot/queries/LookupTransformFunctionBytesPrimaryKeyQueriesTest.java
new file mode 100644
index 00000000000..d1d6d8e9a98
--- /dev/null
+++
b/pinot-core/src/test/java/org/apache/pinot/queries/LookupTransformFunctionBytesPrimaryKeyQueriesTest.java
@@ -0,0 +1,287 @@
+/**
+ * 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.queries;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.Executors;
+import org.apache.commons.io.FileUtils;
+import org.apache.helix.AccessOption;
+import org.apache.helix.HelixManager;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.common.utils.config.SchemaSerDeUtils;
+import org.apache.pinot.common.utils.config.TableConfigSerDeUtils;
+import org.apache.pinot.core.data.manager.offline.DimensionTableDataManager;
+import
org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.local.utils.SegmentLocks;
+import org.apache.pinot.segment.local.utils.SegmentOperationsThrottler;
+import org.apache.pinot.segment.local.utils.SegmentOperationsThrottlerSet;
+import org.apache.pinot.segment.local.utils.SegmentReloadSemaphore;
+import org.apache.pinot.segment.local.utils.ServerReloadJobStatusCache;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
+import org.apache.pinot.spi.config.table.DimensionTableConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.BytesUtils;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+
+
+/// Query-level regression test for the single-stage `lookUp(...)` UDF against
a dimension table with a BYTES
+/// primary key, using a *real* [DimensionTableDataManager] in preload mode
(i.e. a real [FastLookupDimensionTable]
+/// built by [DimensionTableDataManager#createFastLookupDimensionTable])
rather than a mock.
+///
+/// This exists because two separate things needed a BYTES primary key value
to be represented consistently as
+/// [org.apache.pinot.spi.utils.ByteArray] rather than a raw `byte[]`:
+/// 1. The value stored as the
[org.apache.pinot.spi.data.readers.PrimaryKey] component when the dimension
table
+/// is loaded
(`DimensionTableDataManager#createFastLookupDimensionTable`, via
+/// `PinotSegmentRecordReader#getPrimaryKeys`).
+/// 2. The value returned by `FastLookupDimensionTable#getValue` when the
*primary key column itself* is the
+/// column being looked up (it returns the raw `PrimaryKey` component,
unlike a non-key value column, which
+/// is read from the separate `values` array).
+///
+public class LookupTransformFunctionBytesPrimaryKeyQueriesTest extends
BaseQueriesTest {
+ private static final File INDEX_DIR =
+ new File(FileUtils.getTempDirectory(),
"LookupTransformFunctionBytesPrimaryKeyQueriesTest");
+
+ private static final String FACT_RAW_TABLE_NAME = "factTable";
+ private static final String FACT_SEGMENT_NAME = "factTableSegment";
+ private static final String FACT_ID_COLUMN = "id";
+ private static final String FACT_BYTES_KEY_COLUMN = "factBytesKey";
+
+ private static final String DIM_RAW_TABLE_NAME = "dimTable";
+ private static final String DIM_SEGMENT_NAME = "dimTableSegment";
+ private static final String DIM_BYTES_PK_COLUMN = "bytesPk";
+ private static final String DIM_LABEL_COLUMN = "label";
+ private static final String DIM_OFFLINE_TABLE_NAME =
TableNameBuilder.OFFLINE.tableNameWithType(DIM_RAW_TABLE_NAME);
+
+ // Rows that have a matching dimension row, plus one (id = 3) that
legitimately does not match anything, so a
+ // "miss" can be told apart from the bug: both a real miss and the bug under
test return empty bytes, but only
+ // the bug does so for ids 1 and 2, which do have a matching row.
+ private static final byte[] BYTES_ALPHA = {1, 2, 3};
+ private static final byte[] BYTES_BETA = {4, 5, 6};
+ private static final byte[] BYTES_NO_MATCH = {9, 9, 9};
+
+ private static final SegmentOperationsThrottlerSet
SEGMENT_OPERATIONS_THROTTLER = new SegmentOperationsThrottlerSet(
+ new SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true));
+
+ private IndexSegment _indexSegment;
+ private List<IndexSegment> _indexSegments;
+ private DimensionTableDataManager _dimensionTableDataManager;
+
+ @Override
+ protected String getFilter() {
+ return "";
+ }
+
+ @Override
+ protected IndexSegment getIndexSegment() {
+ return _indexSegment;
+ }
+
+ @Override
+ protected List<IndexSegment> getIndexSegments() {
+ return _indexSegments;
+ }
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ FileUtils.deleteDirectory(INDEX_DIR);
+ ServerMetrics.register(mock(ServerMetrics.class));
+
+ ImmutableSegment factSegment = buildFactSegment();
+ _indexSegment = factSegment;
+ _indexSegments = Arrays.asList(factSegment);
+ }
+
+ @AfterClass
+ public void tearDown() {
+ if (_dimensionTableDataManager != null) {
+ _dimensionTableDataManager.shutDown();
+ }
+ FileUtils.deleteQuietly(INDEX_DIR);
+ }
+
+ private ImmutableSegment buildFactSegment()
+ throws Exception {
+ Schema schema = new
Schema.SchemaBuilder().setSchemaName(FACT_RAW_TABLE_NAME)
+ .addSingleValueDimension(FACT_ID_COLUMN, DataType.INT)
+ .addSingleValueDimension(FACT_BYTES_KEY_COLUMN, DataType.BYTES)
+ .build();
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(FACT_RAW_TABLE_NAME).build();
+
+ List<GenericRow> records = new ArrayList<>();
+ records.add(makeRow(FACT_ID_COLUMN, 1, FACT_BYTES_KEY_COLUMN,
BYTES_ALPHA));
+ records.add(makeRow(FACT_ID_COLUMN, 2, FACT_BYTES_KEY_COLUMN, BYTES_BETA));
+ records.add(makeRow(FACT_ID_COLUMN, 3, FACT_BYTES_KEY_COLUMN,
BYTES_NO_MATCH));
+
+ return buildSegment(tableConfig, schema, FACT_SEGMENT_NAME, records);
+ }
+
+ private ImmutableSegment buildDimSegment()
+ throws Exception {
+ Schema schema = new
Schema.SchemaBuilder().setSchemaName(DIM_RAW_TABLE_NAME)
+ .addSingleValueDimension(DIM_BYTES_PK_COLUMN, DataType.BYTES)
+ .addSingleValueDimension(DIM_LABEL_COLUMN, DataType.STRING)
+ .setPrimaryKeyColumns(List.of(DIM_BYTES_PK_COLUMN))
+ .build();
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(DIM_RAW_TABLE_NAME)
+ .setDimensionTableConfig(new DimensionTableConfig(false /*
disablePreload */, false))
+ .build();
+
+ List<GenericRow> records = new ArrayList<>();
+ records.add(makeRow(DIM_BYTES_PK_COLUMN, BYTES_ALPHA, DIM_LABEL_COLUMN,
"alpha"));
+ records.add(makeRow(DIM_BYTES_PK_COLUMN, BYTES_BETA, DIM_LABEL_COLUMN,
"beta"));
+
+ return buildSegment(tableConfig, schema, DIM_SEGMENT_NAME, records);
+ }
+
+ private static GenericRow makeRow(Object... columnValuePairs) {
+ GenericRow row = new GenericRow();
+ for (int i = 0; i < columnValuePairs.length; i += 2) {
+ row.putValue((String) columnValuePairs[i], columnValuePairs[i + 1]);
+ }
+ return row;
+ }
+
+ private ImmutableSegment buildSegment(TableConfig tableConfig, Schema
schema, String segmentName,
+ List<GenericRow> records)
+ throws Exception {
+ SegmentGeneratorConfig segmentGeneratorConfig = new
SegmentGeneratorConfig(tableConfig, schema);
+ segmentGeneratorConfig.setTableName(schema.getSchemaName());
+ segmentGeneratorConfig.setSegmentName(segmentName);
+ segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath());
+
+ SegmentIndexCreationDriverImpl driver = new
SegmentIndexCreationDriverImpl();
+ driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
+ driver.build();
+
+ return ImmutableSegmentLoader.load(new File(INDEX_DIR, segmentName),
ReadMode.mmap);
+ }
+
+ private void registerRealDimensionTable(ImmutableSegment dimSegment, boolean
disablePreload)
+ throws Exception {
+ Schema schema = new
Schema.SchemaBuilder().setSchemaName(DIM_RAW_TABLE_NAME)
+ .addSingleValueDimension(DIM_BYTES_PK_COLUMN, DataType.BYTES)
+ .addSingleValueDimension(DIM_LABEL_COLUMN, DataType.STRING)
+ .setPrimaryKeyColumns(List.of(DIM_BYTES_PK_COLUMN))
+ .build();
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(DIM_RAW_TABLE_NAME)
+ .setDimensionTableConfig(new DimensionTableConfig(disablePreload,
false))
+ .build();
+
+ ZkHelixPropertyStore<ZNRecord> propertyStoreMock =
mock(ZkHelixPropertyStore.class);
+ HelixManager helixManager = mock(HelixManager.class);
+ when(propertyStoreMock.get("/CONFIGS/TABLE/" + DIM_OFFLINE_TABLE_NAME,
null, AccessOption.PERSISTENT))
+ .thenReturn(TableConfigSerDeUtils.toZNRecord(tableConfig));
+ when(propertyStoreMock.get("/SCHEMAS/" + DIM_RAW_TABLE_NAME, null,
AccessOption.PERSISTENT))
+ .thenReturn(SchemaSerDeUtils.toZNRecord(schema));
+ when(helixManager.getHelixPropertyStore()).thenReturn(propertyStoreMock);
+
+ InstanceDataManagerConfig instanceDataManagerConfig =
mock(InstanceDataManagerConfig.class);
+
when(instanceDataManagerConfig.getInstanceDataDir()).thenReturn(INDEX_DIR.getAbsolutePath());
+
+ _dimensionTableDataManager =
DimensionTableDataManager.createInstanceByTableName(DIM_OFFLINE_TABLE_NAME);
+ _dimensionTableDataManager.init(instanceDataManagerConfig, helixManager,
new SegmentLocks(), tableConfig, schema,
+ new SegmentReloadSemaphore(1), Executors.newSingleThreadExecutor(),
null, null, SEGMENT_OPERATIONS_THROTTLER,
+ false, mock(ServerReloadJobStatusCache.class));
+ _dimensionTableDataManager.start();
+ // This is what actually triggers loadLookupTable() ->
createFastLookupDimensionTable(), i.e. the real
+ // production code path, using real
PinotSegmentRecordReader#getPrimaryKeys against the real segment.
+ _dimensionTableDataManager.addSegment(dimSegment, null);
+ }
+
+ @DataProvider(name = "preloadModes")
+ public Object[][] providePreloadModes() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ @Test(dataProvider = "preloadModes")
+ public void testLookupReturnsPrimaryKeyBytesColumnUnwrapped(boolean
disablePreload)
+ throws Exception {
+ ImmutableSegment dimSegment = buildDimSegment();
+ registerRealDimensionTable(dimSegment, disablePreload);
+ String query = "SELECT id, lookUp('" + DIM_RAW_TABLE_NAME + "', '" +
DIM_BYTES_PK_COLUMN + "', '"
+ + DIM_BYTES_PK_COLUMN + "', " + FACT_BYTES_KEY_COLUMN + ") FROM " +
FACT_RAW_TABLE_NAME + " ORDER BY id";
+ List<Object[]> rows = getBrokerResponse(query).getResultTable().getRows();
+ // BaseQueriesTest#getBrokerResponse always simulates one OFFLINE and one
REALTIME server response from the
+ // same underlying segment(s) (see its dataTableMap construction), so with
a single fact segment each logical
+ // row appears twice -- 3 fact rows -> 6 returned rows. This is expected
harness behavior, not a bug in the
+ // query or the lookup logic, so assert per-id rather than on fixed row
indices.
+ assertEquals(rows.size(), 6);
+ assertAllRowsForId(rows, 1, BYTES_ALPHA, "matching row (id=1) must return
its own primary key bytes");
+ assertAllRowsForId(rows, 2, BYTES_BETA, "matching row (id=2) must return
its own primary key bytes");
+ assertAllRowsForId(rows, 3, new byte[0], "non-matching row (id=3) has no
dimension row to return");
+ }
+
+ @Test(dataProvider = "preloadModes")
+ public void testLookupReturnsNonPrimaryKeyValueColumn(boolean disablePreload)
+ throws Exception {
+ ImmutableSegment dimSegment = buildDimSegment();
+ registerRealDimensionTable(dimSegment, disablePreload);
+ String query = "SELECT id, lookUp('" + DIM_RAW_TABLE_NAME + "', '" +
DIM_LABEL_COLUMN + "', '"
+ + DIM_BYTES_PK_COLUMN + "', " + FACT_BYTES_KEY_COLUMN + ") FROM " +
FACT_RAW_TABLE_NAME + " ORDER BY id";
+ List<Object[]> rows = getBrokerResponse(query).getResultTable().getRows();
+ assertEquals(rows.size(), 6); // see comment in
testLookupReturnsPrimaryKeyBytesColumnUnwrapped
+ for (Object[] row : rows) {
+ int id = (Integer) row[0];
+ if (id == 1) {
+ assertEquals(row[1], "alpha");
+ } else if (id == 2) {
+ assertEquals(row[1], "beta");
+ }
+ }
+ }
+
+ private static void assertAllRowsForId(List<Object[]> rows, int id, byte[]
expectedBytes, String message) {
+ int matchCount = 0;
+ for (Object[] row : rows) {
+ if (((Integer) row[0]) == id) {
+ assertEquals(BytesUtils.toBytes((String) row[1]), expectedBytes,
message);
+ matchCount++;
+ }
+ }
+ assertEquals(matchCount, 2, "expected exactly 2 rows for id=" + id + "
(OFFLINE + REALTIME duplication)");
+ }
+}
diff --git a/pinot-core/src/test/resources/data/dimAsset.csv
b/pinot-core/src/test/resources/data/dimAsset.csv
new file mode 100644
index 00000000000..01a8e70857c
--- /dev/null
+++ b/pinot-core/src/test/resources/data/dimAsset.csv
@@ -0,0 +1,2 @@
+assetId,assetName,contentHash
+"550e8400e29b41d4a716446655440000","logo-final-v2.svg","5eb63bbbe01eeed093cb22bb8f5acdc3"
\ No newline at end of file
diff --git a/pinot-core/src/test/resources/data/dimAsset_config.json
b/pinot-core/src/test/resources/data/dimAsset_config.json
new file mode 100644
index 00000000000..151a2128027
--- /dev/null
+++ b/pinot-core/src/test/resources/data/dimAsset_config.json
@@ -0,0 +1,18 @@
+{
+ "tableName": "dimAsset",
+ "tableType": "OFFLINE",
+ "isDimTable": true,
+ "segmentsConfig": {
+ "segmentPushType": "REFRESH",
+ "replication": "1"
+ },
+ "tenants": {
+ },
+ "tableIndexConfig": {
+ "loadMode": "MMAP"
+ },
+ "metadata": {
+ "customConfigs": {
+ }
+ }
+}
diff --git a/pinot-core/src/test/resources/data/dimAsset_schema.json
b/pinot-core/src/test/resources/data/dimAsset_schema.json
new file mode 100644
index 00000000000..21257c313bf
--- /dev/null
+++ b/pinot-core/src/test/resources/data/dimAsset_schema.json
@@ -0,0 +1,9 @@
+{
+ "schemaName": "dimAsset",
+ "primaryKeyColumns": ["assetId"],
+ "dimensionFieldSpecs": [
+ {"name": "assetId", "dataType": "BYTES"},
+ {"name": "assetName", "dataType": "STRING"},
+ {"name": "contentHash", "dataType": "BYTES"}
+ ]
+}
\ No newline at end of file
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
index fc5db95f8f2..f9a9983f761 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
@@ -591,6 +591,10 @@ public abstract class QueryRunnerTestBase extends
QueryTestSet {
public boolean _replicated;
@JsonProperty("isDimTable")
public boolean _isDimTable;
+ @JsonProperty("isActualDimTableNeeded")
+ public boolean _isActualDimTableNeeded;
+ @JsonProperty("disablePreload")
+ public boolean _disablePreload; // Help diff between FastLookup and
MemOpt DimData table
@JsonProperty("primaryKeyColumns")
public List<String> _primaryKeyColumns;
}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/ResourceBasedQueriesTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/ResourceBasedQueriesTest.java
index cfc7724ba9c..73a75451a5e 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/ResourceBasedQueriesTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/ResourceBasedQueriesTest.java
@@ -36,11 +36,19 @@ import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.TimeZone;
+import java.util.concurrent.Executors;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.Nullable;
+import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.tuple.Pair;
+import org.apache.helix.AccessOption;
+import org.apache.helix.HelixManager;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
import org.apache.pinot.common.response.broker.BrokerResponseNativeV2;
+import org.apache.pinot.common.utils.config.SchemaSerDeUtils;
+import org.apache.pinot.common.utils.config.TableConfigSerDeUtils;
import org.apache.pinot.core.data.manager.offline.DimensionTableDataManager;
import org.apache.pinot.query.QueryEnvironmentTestBase;
import org.apache.pinot.query.QueryServerEnclosure;
@@ -54,8 +62,16 @@ import
org.apache.pinot.query.service.dispatch.QueryDispatcher;
import org.apache.pinot.query.testutils.MockInstanceDataManagerFactory;
import org.apache.pinot.query.testutils.QueryTestUtils;
import org.apache.pinot.segment.local.segment.readers.PinotSegmentRecordReader;
+import org.apache.pinot.segment.local.utils.SegmentLocks;
+import org.apache.pinot.segment.local.utils.SegmentOperationsThrottler;
+import org.apache.pinot.segment.local.utils.SegmentOperationsThrottlerSet;
+import org.apache.pinot.segment.local.utils.SegmentReloadSemaphore;
+import org.apache.pinot.segment.local.utils.ServerReloadJobStatusCache;
import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.spi.config.instance.InstanceDataManagerConfig;
import org.apache.pinot.spi.config.instance.InstanceType;
+import org.apache.pinot.spi.config.table.DimensionTableConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
@@ -63,10 +79,10 @@ import org.apache.pinot.spi.data.readers.PrimaryKey;
import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner;
import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.apache.pinot.spi.utils.builder.TableNameBuilder;
import org.assertj.core.api.Assertions;
import org.mockito.ArgumentMatchers;
-import org.mockito.Mockito;
import org.testng.Assert;
import org.testng.SkipException;
import org.testng.annotations.AfterClass;
@@ -74,6 +90,9 @@ import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
public class ResourceBasedQueriesTest extends QueryRunnerTestBase {
private static final ObjectMapper MAPPER = new ObjectMapper();
@@ -87,6 +106,13 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
private boolean _isRunIgnored;
private TimeZone _currentSystemTimeZone;
+ // For real dimension table uses
+ private static final SegmentOperationsThrottlerSet
SEGMENT_OPERATIONS_THROTTLER =
+ new SegmentOperationsThrottlerSet(new SegmentOperationsThrottler(1, 2,
true),
+ new SegmentOperationsThrottler(1, 2, true), new
SegmentOperationsThrottler(1, 2, true),
+ new SegmentOperationsThrottler(1, 2, true));
+ private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(),
ResourceBasedQueriesTest.class.getName());
+
@BeforeClass
public void setUp()
throws Exception {
@@ -136,7 +162,11 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
if (table._replicated) {
ImmutableSegment segment = addSegmentReplicated(factory1, factory2,
offlineTableName, genericRows);
if (table._isDimTable) {
- registerMockDimensionTable(offlineTableName, schema, table,
segment);
+ if (table._isActualDimTableNeeded) {
+ registerRealDimensionTable(offlineTableName, schema, table,
segment, table._disablePreload);
+ } else {
+ registerMockDimensionTable(offlineTableName, schema, table,
segment);
+ }
}
continue;
}
@@ -300,20 +330,20 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
for (int i = 0; i < columns.size(); i++) {
row.put(columns.get(i), values[i]);
}
- lookupMap.put(new PrimaryKey(recordReader.getRecordValues(docId,
primaryKeyIndexes)), row);
+ Object[] o = recordReader.getRecordValues(docId, primaryKeyIndexes);
+ lookupMap.put(new PrimaryKey(o), row);
}
} catch (Exception e) {
throw new RuntimeException("Failed to build the mock dimension table
for: " + offlineTableName, e);
}
// Create and register a mock DimensionTableDataManager
- DimensionTableDataManager mockDimManager =
Mockito.mock(DimensionTableDataManager.class);
-
Mockito.when(mockDimManager.getPrimaryKeyColumns()).thenReturn(primaryKeyColumns);
-
Mockito.when(mockDimManager.containsKey(ArgumentMatchers.any(PrimaryKey.class)))
- .thenAnswer(invocation -> {
- PrimaryKey pk = invocation.getArgument(0);
- return lookupMap.containsKey(pk);
- });
-
Mockito.when(mockDimManager.lookupValues(ArgumentMatchers.any(PrimaryKey.class),
+ DimensionTableDataManager mockDimManager =
mock(DimensionTableDataManager.class);
+ when(mockDimManager.getPrimaryKeyColumns()).thenReturn(primaryKeyColumns);
+
when(mockDimManager.containsKey(ArgumentMatchers.any(PrimaryKey.class))).thenAnswer(invocation
-> {
+ PrimaryKey pk = invocation.getArgument(0);
+ return lookupMap.containsKey(pk);
+ });
+ when(mockDimManager.lookupValues(ArgumentMatchers.any(PrimaryKey.class),
ArgumentMatchers.any(String[].class))).thenAnswer(invocation -> {
PrimaryKey pk = invocation.getArgument(0);
String[] lookupColumns = invocation.getArgument(1);
@@ -330,6 +360,34 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
DimensionTableDataManager.registerDimensionTable(offlineTableName,
mockDimManager);
}
+ private void registerRealDimensionTable(String offlineTableName, Schema
schema, QueryTestCase.Table table,
+ ImmutableSegment segment, boolean disablePreload)
+ throws Exception {
+ TableConfig tableConfig =
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(TableNameBuilder.extractRawTableName(offlineTableName))
+ .setDimensionTableConfig(new DimensionTableConfig(disablePreload,
false))
+ .build();
+
+ ZkHelixPropertyStore<ZNRecord> propertyStoreMock =
mock(ZkHelixPropertyStore.class);
+ HelixManager helixManager = mock(HelixManager.class);
+ when(propertyStoreMock.get("/CONFIGS/TABLE/" + offlineTableName, null,
AccessOption.PERSISTENT)).thenReturn(
+ TableConfigSerDeUtils.toZNRecord(tableConfig));
+ when(propertyStoreMock.get("/SCHEMAS/" + schema.getSchemaName(), null,
AccessOption.PERSISTENT)).thenReturn(
+ SchemaSerDeUtils.toZNRecord(schema));
+ when(helixManager.getHelixPropertyStore()).thenReturn(propertyStoreMock);
+
+ InstanceDataManagerConfig instanceDataManagerConfig =
mock(InstanceDataManagerConfig.class);
+
when(instanceDataManagerConfig.getInstanceDataDir()).thenReturn(TEMP_DIR.getAbsolutePath());
+
+ DimensionTableDataManager tableDataManager =
DimensionTableDataManager.createInstanceByTableName(offlineTableName);
+ tableDataManager.init(instanceDataManagerConfig, helixManager, new
SegmentLocks(), tableConfig, schema,
+ new SegmentReloadSemaphore(1), Executors.newSingleThreadExecutor(),
null, null, SEGMENT_OPERATIONS_THROTTLER,
+ false, mock(ServerReloadJobStatusCache.class));
+ tableDataManager.start();
+ tableDataManager.addSegment(segment, null); // <-- triggers the real
createFastLookupDimensionTable()
+ DimensionTableDataManager.registerDimensionTable(offlineTableName,
tableDataManager);
+ }
+
@AfterClass
public void tearDown() {
// Restore the original default timezone
@@ -338,6 +396,7 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
server.shutDown();
}
_mailboxService.shutdown();
+ FileUtils.deleteQuietly(TEMP_DIR);
}
// TODO: name the test using testCaseName for testng reports
@@ -400,8 +459,8 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
}
@Test(dataProvider = "testResourceQueryTestCaseProviderBoth")
- public void testQueryTestCasesWithLiteModeWithOutput(String testCaseName,
boolean isIgnored, String sql,
- String h2Sql, List<Object[]> expectedRows, String expect, boolean
keepOutputRowOrder, boolean ignoreV2Optimizer,
+ public void testQueryTestCasesWithLiteModeWithOutput(String testCaseName,
boolean isIgnored, String sql, String h2Sql,
+ List<Object[]> expectedRows, String expect, boolean keepOutputRowOrder,
boolean ignoreV2Optimizer,
boolean ignoreLiteMode)
throws Exception {
if (ignoreV2Optimizer || ignoreLiteMode) {
@@ -487,8 +546,9 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
Assertions.assertThat(queryResult.getProcessingException().getMessage()).matches(pattern);
return Optional.empty();
}
- Assert.assertNull(expectedErrorMsg, "Expected error with message '" +
expectedErrorMsg
- + "'. But instead rows were returned: " +
JsonUtils.objectToPrettyString(queryResult.getResultTable()));
+ Assert.assertNull(expectedErrorMsg,
+ "Expected error with message '" + expectedErrorMsg + "'. But instead
rows were returned: "
+ + JsonUtils.objectToPrettyString(queryResult.getResultTable()));
Assert.assertNotNull(queryResult.getResultTable(),
"Result table is null: " +
JsonUtils.objectToPrettyString(queryResult));
return Optional.of(queryResult);
@@ -611,8 +671,8 @@ public class ResourceBasedQueriesTest extends
QueryRunnerTestBase {
String h2Sql = queryCase._h2Sql != null ?
replaceTableName(testCaseName, queryCase._h2Sql)
: replaceTableName(testCaseName, queryCase._sql);
Object[] testEntry = new Object[]{
- testCaseName, queryCase._ignored, sql, h2Sql,
queryCase._expectedException, queryCase._keepOutputRowOrder,
- queryCase._ignoreV2Optimizer, queryCase._ignoreLiteMode
+ testCaseName, queryCase._ignored, sql, h2Sql,
queryCase._expectedException,
+ queryCase._keepOutputRowOrder, queryCase._ignoreV2Optimizer,
queryCase._ignoreLiteMode
};
providerContent.add(testEntry);
}
diff --git a/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
b/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
index d5ed8dcabd5..ab9810e917f 100644
--- a/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
+++ b/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
@@ -262,5 +262,157 @@
"ignoreLiteMode": true
}
]
+ },
+ "lookup_join_bytes_primary_key_mem_opt_dim_table": {
+ "comment": "Regression test for issue 19228. This fixes the case where a
dimension table contains bytes type primary key.The equals and hashcode of
bytes array is object identity based and so fails on equality check. This fix
wraps bytes primary in the internal ByteArray.",
+ "tables": {
+ "fact_tbl": {
+ "schema": [
+ {"name": "bytes_value", "type": "BYTES"},
+ {"name": "int_value", "type": "INT"}
+ ],
+ "inputs": [
+ ["f47ac10b58cc4372a5670e02b2c3d479", 1],
+ ["9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d", 2],
+ ["3c9b7e128a4d4e90b183f6a89c25e21a", 3],
+ ["a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d", 4]
+ ]
+ },
+ "dim_tbl": {
+ "schema": [
+ {"name": "bytes_value", "type": "BYTES"},
+ {"name": "int_value", "type": "INT"}
+ ],
+ "inputs": [
+ ["f47ac10b58cc4372a5670e02b2c3d479", 1],
+ ["9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d", 2],
+ ["a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d", 4]
+ ],
+ "replicated": true,
+ "isDimTable": true,
+ "isActualDimTableNeeded": true,
+ "disablePreload": true,
+ "primaryKeyColumns": ["bytes_value"]
+ }
+ }, "queries": [
+ {
+ "description": "Checking with a constant value check",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} JOIN {dim_tbl} ON {dim_tbl}.bytes_value =
'a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d' AND {dim_tbl}.bytes_value =
{fact_tbl}.bytes_value",
+ "outputs": [
+ [4]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking full table join",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.bytes_value,{dim_tbl}.int_value FROM {fact_tbl} JOIN {dim_tbl} ON
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value",
+ "outputs": [
+ ["f47ac10b58cc4372a5670e02b2c3d479", 1],
+ ["9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d", 2],
+ ["a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d", 4]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking with a constant value check",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} JOIN {dim_tbl} ON {dim_tbl}.bytes_value in
('a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d','9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d') AND
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value",
+ "outputs": [
+ [4], [2]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking left join with comparison operator to test
case where predicates can possibly fail on a byte[] expecting a ByteArray",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} LEFT JOIN {dim_tbl} ON
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value AND {dim_tbl}.bytes_value <
'a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d'",
+ "outputs": [
+ [null], [2], [null], [null]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking left join with comparison operator to test
case where predicates can possibly fail on a byte[] expecting a ByteArray,
ensuring the comparison value is considered bytes",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} LEFT JOIN {dim_tbl} ON
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value AND {dim_tbl}.bytes_value <
X'a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d'",
+ "outputs": [
+ [null], [2], [null], [null]
+ ],
+ "ignoreLiteMode": true
+ }
+ ]
+ },
+ "lookup_join_bytes_primary_key_fast_lookup_dim_table": {
+ "comment": "Regression test for issue 19228. This fixes the case where a
dimension table contains bytes type primary key.The equals and hashcode of
bytes array is object identity based and so fails on equality check. This fix
wraps bytes primary in the internal ByteArray.",
+ "tables": {
+ "fact_tbl": {
+ "schema": [
+ {"name": "bytes_value", "type": "BYTES"},
+ {"name": "int_value", "type": "INT"}
+ ],
+ "inputs": [
+ ["f47ac10b58cc4372a5670e02b2c3d479", 1],
+ ["9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d", 2],
+ ["3c9b7e128a4d4e90b183f6a89c25e21a", 3],
+ ["a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d", 4]
+ ]
+ },
+ "dim_tbl": {
+ "schema": [
+ {"name": "bytes_value", "type": "BYTES"},
+ {"name": "int_value", "type": "INT"}
+ ],
+ "inputs": [
+ ["f47ac10b58cc4372a5670e02b2c3d479", 1],
+ ["9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d", 2],
+ ["a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d", 4]
+ ],
+ "replicated": true,
+ "isDimTable": true,
+ "isActualDimTableNeeded": true,
+ "disablePreload": false,
+ "primaryKeyColumns": ["bytes_value"]
+ }
+ }, "queries": [
+ {
+ "description": "Checking with a constant value check",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} JOIN {dim_tbl} ON {dim_tbl}.bytes_value =
'a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d' AND {dim_tbl}.bytes_value =
{fact_tbl}.bytes_value",
+ "outputs": [
+ [4]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking full table join",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.bytes_value,{dim_tbl}.int_value FROM {fact_tbl} JOIN {dim_tbl} ON
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value",
+ "outputs": [
+ ["f47ac10b58cc4372a5670e02b2c3d479", 1],
+ ["9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d", 2],
+ ["a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d", 4]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking with a constant value check",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} JOIN {dim_tbl} ON {dim_tbl}.bytes_value in
('a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d','9b1deb4d3b7d4bad9bdd2b0d7b3dcb6d') AND
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value",
+ "outputs": [
+ [4], [2]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking left join with comparison operator to test
case where predicates can possibly fail on a byte[] expecting a ByteArray",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} LEFT JOIN {dim_tbl} ON
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value AND {dim_tbl}.bytes_value <
'a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d'",
+ "outputs": [
+ [null], [2], [null], [null]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Checking left join with comparison operator to test
case where predicates can possibly fail on a byte[] expecting a ByteArray,
ensuring the comparison value is considered bytes",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */
{dim_tbl}.int_value FROM {fact_tbl} LEFT JOIN {dim_tbl} ON
{dim_tbl}.bytes_value = {fact_tbl}.bytes_value AND {dim_tbl}.bytes_value <
X'a1b2c3d4e5f64a7b8c9d0e1f2a3b4c5d'",
+ "outputs": [
+ [null], [2], [null], [null]
+ ],
+ "ignoreLiteMode": true
+ }
+ ]
}
-}
+}
\ No newline at end of file
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
index 12c7ec5a747..a97c998c314 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentColumnReader.java
@@ -650,6 +650,15 @@ public class PinotSegmentColumnReader implements Closeable
{
}
}
+ public Object getInternalValue(int docId) {
+ assert _forwardIndexReader.isSingleValue();
+ Object value = getValue(docId);
+ if (_valueType == DataType.BYTES) {
+ return new ByteArray((byte[]) value);
+ }
+ return value;
+ }
+
@Override
public void close()
throws IOException {
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReader.java
index bb7ab182618..3619ec9397b 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReader.java
@@ -68,6 +68,8 @@ public class PinotSegmentRecordReader implements RecordReader
{
private boolean _skipDefaultNullValues;
private int _nextDocId = 0;
+ private List<String> _primaryKeyColumns;
+ private int[] _primaryKeyIndexes;
public PinotSegmentRecordReader() {
}
@@ -80,7 +82,7 @@ public class PinotSegmentRecordReader implements RecordReader
{
public PinotSegmentRecordReader(File indexDir)
throws Exception {
try {
- init(indexDir, null, null, false);
+ init(indexDir, null, null, false, null);
} catch (Exception e) {
close();
throw e;
@@ -100,8 +102,9 @@ public class PinotSegmentRecordReader implements
RecordReader {
public PinotSegmentRecordReader(File indexDir, @Nullable Schema schema,
@Nullable List<String> sortOrder)
throws Exception {
Set<String> fieldsToRead = schema != null ?
schema.getPhysicalColumnNames() : null;
+ List<String> primaryKeyColumns = schema != null ?
schema.getPrimaryKeyColumns() : null;
try {
- init(indexDir, fieldsToRead, sortOrder, false);
+ init(indexDir, fieldsToRead, sortOrder, false, primaryKeyColumns);
} catch (Exception e) {
close();
throw e;
@@ -110,7 +113,19 @@ public class PinotSegmentRecordReader implements
RecordReader {
@Override
public void init(File indexDir, @Nullable Set<String> fieldsToRead,
@Nullable RecordReaderConfig recordReaderConfig) {
- init(indexDir, fieldsToRead, null, true);
+ init(indexDir, fieldsToRead, null, true, null);
+ }
+
+ /// Initializes the record reader from an index directory.
+ ///
+ /// @param indexDir Index directory
+ /// @param fieldsToRead The fields to read from the segment. If null or
empty, reads all fields
+ /// @param sortOrder List of sorted columns
+ /// @param skipDefaultNullValues Whether to skip putting default null values
into the record
+ /// @param primaryKeyColumns Columns which are marked as primary keys
+ public void init(File indexDir, @Nullable Set<String> fieldsToRead,
@Nullable List<String> sortOrder,
+ boolean skipDefaultNullValues, @Nullable List<String> primaryKeyColumns)
{
+ init(indexDir, fieldsToRead, sortOrder, skipDefaultNullValues, false,
primaryKeyColumns);
}
/// Initializes the record reader from an index directory.
@@ -121,7 +136,7 @@ public class PinotSegmentRecordReader implements
RecordReader {
/// @param skipDefaultNullValues Whether to skip putting default null values
into the record
public void init(File indexDir, @Nullable Set<String> fieldsToRead,
@Nullable List<String> sortOrder,
boolean skipDefaultNullValues) {
- init(indexDir, fieldsToRead, sortOrder, skipDefaultNullValues, false);
+ init(indexDir, fieldsToRead, sortOrder, skipDefaultNullValues, false,
null);
}
/// Initializes the record reader from an index directory with an option to
skip column-level secondary indexes.
@@ -132,22 +147,44 @@ public class PinotSegmentRecordReader implements
RecordReader {
/// @param skipDefaultNullValues Whether to skip putting default null values
into the record
/// @param forwardIndexOnly Whether to load only column-level forward index,
dictionary, and null value vector,
/// skipping other column-level secondary indexes
+ /// @param primaryKeyColumns Columns which are marked as primary keys
public void init(File indexDir, @Nullable Set<String> fieldsToRead,
@Nullable List<String> sortOrder,
- boolean skipDefaultNullValues, boolean forwardIndexOnly) {
+ boolean skipDefaultNullValues, boolean forwardIndexOnly, @Nullable
List<String> primaryKeyColumns) {
IndexSegment indexSegment;
try {
indexSegment = ImmutableSegmentLoader.load(indexDir, ReadMode.mmap,
forwardIndexOnly);
} catch (Exception e) {
throw new RuntimeException("Caught exception while loading the segment
from: " + indexDir, e);
}
- init(indexSegment, true, fieldsToRead, null, sortOrder,
skipDefaultNullValues);
+ init(indexSegment, true, fieldsToRead, null, sortOrder,
skipDefaultNullValues, primaryKeyColumns);
+ }
+
+ /// Initializes the record reader from an index directory with an option to
skip column-level secondary indexes.
+ ///
+ /// @param indexDir Index directory
+ /// @param fieldsToRead The fields to read from the segment. If null or
empty, reads all fields
+ /// @param sortOrder List of sorted columns
+ /// @param skipDefaultNullValues Whether to skip putting default null values
into the record
+ /// @param forwardIndexOnly Whether to load only column-level forward index,
dictionary, and null value vector,
+ /// skipping other column-level secondary indexes
+ public void init(File indexDir, @Nullable Set<String> fieldsToRead,
@Nullable List<String> sortOrder,
+ boolean skipDefaultNullValues, boolean forwardIndexOnly) {
+ init(indexDir, fieldsToRead, sortOrder, skipDefaultNullValues,
forwardIndexOnly, null);
+ }
+
+ /// Initializes the record reader from a segment.
+ ///
+ /// @param indexSegment Index segment to read from
+ /// @param primaryKeyColumns Columns which are marked as primary keys
+ public void init(IndexSegment indexSegment, @Nullable List<String>
primaryKeyColumns) {
+ init(indexSegment, false, null, null, null, false, primaryKeyColumns);
}
/// Initializes the record reader from a segment.
///
/// @param indexSegment Index segment to read from
public void init(IndexSegment indexSegment) {
- init(indexSegment, false, null, null, null, false);
+ init(indexSegment, false, null, null, null, false, null);
}
/// Initializes the record reader from a mutable segment with optional
sorted document ids.
@@ -155,7 +192,7 @@ public class PinotSegmentRecordReader implements
RecordReader {
/// @param mutableSegment Mutable segment
/// @param sortedDocIds Array of sorted document ids
public void init(MutableSegment mutableSegment, @Nullable int[]
sortedDocIds) {
- init(mutableSegment, false, null, sortedDocIds, null, false);
+ init(mutableSegment, false, null, sortedDocIds, null, false, null);
}
/// Initializes the record reader.
@@ -167,7 +204,8 @@ public class PinotSegmentRecordReader implements
RecordReader {
/// @param sortOrder List of sorted columns
/// @param skipDefaultNullValues Whether to skip putting default null values
into the record
private void init(IndexSegment indexSegment, boolean destroySegmentOnClose,
@Nullable Set<String> fieldsToRead,
- @Nullable int[] sortedDocIds, @Nullable List<String> sortOrder, boolean
skipDefaultNullValues) {
+ @Nullable int[] sortedDocIds, @Nullable List<String> sortOrder, boolean
skipDefaultNullValues,
+ @Nullable List<String> primaryKeyColumns) {
_indexSegment = indexSegment;
_destroySegmentOnClose = destroySegmentOnClose;
_numDocs = _indexSegment.getSegmentMetadata().getTotalDocs();
@@ -204,6 +242,7 @@ public class PinotSegmentRecordReader implements
RecordReader {
}
_skipDefaultNullValues = skipDefaultNullValues;
+ _primaryKeyColumns = primaryKeyColumns;
}
}
@@ -357,4 +396,30 @@ public class PinotSegmentRecordReader implements
RecordReader {
throw closeException;
}
}
+
+ public Object[] getPrimaryKeys(int docId) {
+ int[] primaryKeyIndexes = getPrimaryKeyIndexes();
+ Object[] values = new Object[primaryKeyIndexes.length];
+ for (int i = 0, n = primaryKeyIndexes.length; i < n; i++) {
+ int columnIndex = primaryKeyIndexes[i];
+ if (columnIndex > -1) {
+ PinotSegmentColumnReader columnReader =
_columnReaders.get(columnIndex);
+ if (!columnReader.isNull(docId)) {
+ values[i] = columnReader.getInternalValue(docId);
+ } else if (!_skipDefaultNullValues) {
+ values[i] = columnReader.getInternalValue(docId);
+ } // else null value is kept
+ } // else keep null value
+ }
+
+ return values;
+ }
+
+ private int[] getPrimaryKeyIndexes() {
+ if (_primaryKeyIndexes != null) {
+ return _primaryKeyIndexes;
+ }
+ _primaryKeyIndexes = getIndexesForColumns(_primaryKeyColumns);
+ return _primaryKeyIndexes;
+ }
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReaderTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReaderTest.java
index 68601b7aafb..191f6d6c5df 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReaderTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentRecordReaderTest.java
@@ -31,6 +31,7 @@ import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.TimeGranularitySpec;
import org.apache.pinot.spi.data.readers.GenericRow;
import org.apache.pinot.spi.data.readers.RecordReader;
+import org.apache.pinot.spi.utils.ByteArray;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
@@ -47,6 +48,9 @@ public class PinotSegmentRecordReaderTest {
private static final String M1 = "m1";
private static final String M2 = "m2";
private static final String TIME = "t";
+ private static final String D_SV_PK_1 = "d_sv_pk_1";
+ private static final String D_SV_PK_2 = "d_sv_pk_2";
+ private static final String D_SV_PK_3 = "d_sv_pk_3";
private String _segmentOutputDir;
private File _segmentIndexDir;
@@ -74,9 +78,13 @@ public class PinotSegmentRecordReaderTest {
}
private Schema createPinotSchema() {
- return new
Schema.SchemaBuilder().setSchemaName("schema").addSingleValueDimension(D_SV_1,
DataType.STRING)
+ Schema schema = new
Schema.SchemaBuilder().setSchemaName("schema").addSingleValueDimension(D_SV_1,
DataType.STRING)
.addMultiValueDimension(D_MV_1, DataType.STRING).addMetric(M1,
DataType.INT).addMetric(M2, DataType.FLOAT)
- .addTime(new TimeGranularitySpec(DataType.LONG, TimeUnit.HOURS, TIME),
null).build();
+ .addTime(new TimeGranularitySpec(DataType.LONG, TimeUnit.HOURS, TIME),
null)
+ .addSingleValueDimension(D_SV_PK_1,
DataType.STRING).addSingleValueDimension(D_SV_PK_2, DataType.BYTES)
+ .addSingleValueDimension(D_SV_PK_3, DataType.INT).build();
+ schema.setPrimaryKeyColumns(List.of(D_SV_PK_1, D_SV_PK_2, D_SV_PK_3));
+ return schema;
}
private TableConfig createTableConfig() {
@@ -104,6 +112,9 @@ public class PinotSegmentRecordReaderTest {
Assert.assertEquals(outputRow.getValue(M1), row.getValue(M1));
Assert.assertEquals(outputRow.getValue(M2), row.getValue(M2));
Assert.assertEquals(outputRow.getValue(TIME), row.getValue(TIME));
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_1),
row.getValue(D_SV_PK_1));
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_2),
row.getValue(D_SV_PK_2));
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_3),
row.getValue(D_SV_PK_3));
}
}
@@ -181,6 +192,32 @@ public class PinotSegmentRecordReaderTest {
Assert.assertEquals(outputRow.getValue(M1), row.getValue(M1));
Assert.assertEquals(outputRow.getValue(M2), row.getValue(M2));
Assert.assertEquals(outputRow.getValue(TIME), row.getValue(TIME));
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_1),
row.getValue(D_SV_PK_1));
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_2),
row.getValue(D_SV_PK_2));
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_3),
row.getValue(D_SV_PK_3));
+ }
+ }
+
+ @Test
+ public void testPinotSegmentRecordReaderReadPrimaryKeysOnly()
+ throws Exception {
+ GenericRow outputRow = null;
+
+ PinotSegmentRecordReader pinotSegmentRecordReader = new
PinotSegmentRecordReader();
+ pinotSegmentRecordReader.init(_segmentIndexDir, null, null, false, true,
List.of(D_SV_PK_1, D_SV_PK_2, D_SV_PK_3));
+
+ try (pinotSegmentRecordReader) {
+ int docId = 0;
+ while (pinotSegmentRecordReader.hasNext()) {
+ outputRow = pinotSegmentRecordReader.next();
+ Object[] primaryKeys = pinotSegmentRecordReader.getPrimaryKeys(docId);
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_1), primaryKeys[0]);
+ byte[] primaryKeyBytes = (byte[]) outputRow.getValue(D_SV_PK_2);
+ ByteArray byteArray = new ByteArray(primaryKeyBytes);
+ Assert.assertEquals(byteArray, primaryKeys[1]);
+ Assert.assertEquals(outputRow.getValue(D_SV_PK_3), primaryKeys[2]);
+ docId++;
+ }
}
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentUtil.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentUtil.java
index 67b98d6630a..0eb1cfed579 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentUtil.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/readers/PinotSegmentUtil.java
@@ -113,6 +113,10 @@ public class PinotSegmentUtil {
return Math.abs(random.nextDouble());
case STRING:
return
RandomStringUtils.secure().nextAlphabetic(DEFAULT_STRING_VALUE_LENGTH);
+ case BYTES:
+ byte[] randomBytes = new byte[16];
+ random.nextBytes(randomBytes);
+ return randomBytes;
default:
throw new IllegalStateException("Unsupported data type: " +
storedType);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
