From a7ec3177b4d9864d2444767e33955407291ace7c Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 28 Jul 2026 15:34:23 -0700 Subject: [PATCH 1/8] Add end-to-end Parquet VARIANT support Introduce the VARIANT logical type and envelope, Parquet reconstruction, scalar and transform functions, and segment and query validation across the single-stage and multi-stage engines. Propagate VARIANT through schemas, wire formats, responses, clients, DDL, quickstart, and integration coverage while rejecting raw operations that require ordering, equality, or hashing. --- .../workflows/scripts/.pinot_quickstart.sh | 35 + .../pinot/client/ResultTableResultSet.java | 9 +- .../client/ResultTableResultSetTest.java | 19 + .../apache/pinot/client/PinotResultSet.java | 13 +- .../pinot/client/grpc/PinotGrpcResultSet.java | 14 +- .../pinot/client/utils/DriverUtils.java | 1 + .../pinot/client/PinotResultSetTest.java | 61 + .../client/grpc/PinotGrpcResultSetTest.java | 101 + .../test/resources/variant_result_table.json | 35 + pinot-common/pom.xml | 16 + .../common/datablock/DataBlockEquals.java | 1 + .../evaluator/InbuiltFunctionEvaluator.java | 154 ++ .../common/function/FunctionRegistry.java | 4 +- .../function/TransformFunctionType.java | 87 +- .../function/scalar/VariantFunctions.java | 102 ++ .../encoder/ArrowResponseEncoder.java | 3 + .../response/encoder/JsonResponseEncoder.java | 1 + .../apache/pinot/common/utils/DataSchema.java | 85 +- .../pinot/common/utils/VariantUtils.java | 1624 +++++++++++++++++ pinot-common/src/main/proto/expressions.proto | 1 + .../InbuiltFunctionEvaluatorTest.java | 116 ++ ...ssionsColumnDataTypeCompatibilityTest.java | 87 + .../encoder/ArrowResponseEncoderTest.java | 25 + .../pinot/common/utils/DataSchemaTest.java | 138 +- .../pinot/common/utils/VariantUtilsTest.java | 541 ++++++ .../src/test/proto/legacy_expressions.proto | 58 + .../resources/PinotDdlRestletResource.java | 5 + .../api/PinotDdlRestletResourceTest.java | 41 + .../predicate/PredicateEvaluatorProvider.java | 39 + .../query/FilteredGroupByOperator.java | 9 +- .../core/operator/query/GroupByOperator.java | 9 +- .../BaseVariantTransformFunction.java | 161 ++ .../BinaryOperatorTransformFunction.java | 2 + .../function/InTransformFunction.java | 5 + .../IsVariantNullTransformFunction.java | 82 + .../ParseJsonToVariantTransformFunction.java | 168 ++ .../function/TransformFunctionFactory.java | 21 +- .../VariantExistsTransformFunction.java | 82 + .../function/VariantGetTransformFunction.java | 253 +++ .../VariantTypeOfTransformFunction.java | 82 + .../pinot/core/plan/DistinctPlanNode.java | 13 + .../function/AggregationFunctionUtils.java | 53 + .../function/AnyValueAggregationFunction.java | 7 +- .../distinct/DistinctExecutorFactory.java | 8 + .../query/utils/OrderByComparatorFactory.java | 4 + .../datablock/DataBlockBuilderTest.java | 12 +- .../common/datablock/DataBlockTestUtils.java | 5 + .../common/datatable/DataTableSerDeTest.java | 11 + .../core/data/manager/TableIndexingTest.java | 3 +- .../core/function/FunctionRegistryTest.java | 18 + .../PredicateEvaluatorProviderTest.java | 24 + .../IsVariantNullTransformFunctionTest.java | 162 ++ .../VariantExistsTransformFunctionTest.java | 154 ++ .../VariantGetTransformFunctionTest.java | 519 ++++++ .../VariantTypeOfTransformFunctionTest.java | 164 ++ .../pinot/core/plan/DistinctPlanNodeTest.java | 58 + .../AggregationFunctionUtilsTest.java | 53 + .../AnyValueAggregationFunctionTest.java | 23 + .../utils/OrderByComparatorFactoryTest.java | 19 + .../pinot/integration/tests/ClusterTest.java | 1 + .../tests/custom/VariantTypeTest.java | 404 ++++ .../pinot-input-format/pinot-parquet/pom.xml | 4 + .../parquet/ParquetAvroRecordExtractor.java | 6 +- .../parquet/ParquetAvroRecordReader.java | 59 +- .../parquet/ParquetNativeRecordExtractor.java | 98 +- .../ParquetNativeRecordExtractorConfig.java | 20 + .../parquet/ParquetNativeRecordReader.java | 75 +- .../parquet/ParquetRecordReader.java | 38 +- .../parquet/ParquetRecordReaderConfig.java | 3 +- .../inputformat/parquet/ParquetUtils.java | 26 +- .../parquet/ParquetVariantConverter.java | 305 ++++ .../ParquetVariantRecordReaderTest.java | 876 +++++++++ .../rel/rules/PinotEvaluateLiteralRule.java | 7 + .../calcite/sql/fun/PinotOperatorTable.java | 7 + .../parser/CalciteRexExpressionParser.java | 2 +- .../logical/RelToPlanNodeConverter.java | 5 + .../planner/logical/RexExpressionUtils.java | 4 +- .../physical/PinotDispatchPlanner.java | 4 +- .../physical/v2/PRelToPlanNodeConverter.java | 5 + .../serde/ProtoExpressionToRexExpression.java | 6 + .../serde/RexExpressionToProtoExpression.java | 2 + .../VariantTypeValidationVisitor.java | 201 ++ .../apache/pinot/query/type/TypeFactory.java | 2 + .../pinot/query/QueryCompilationTest.java | 34 + .../pinot/query/QueryEnvironmentTestBase.java | 9 +- .../planner/serde/RexExpressionSerDeTest.java | 22 +- .../VariantTypeValidationVisitorTest.java | 235 +++ .../pinot/query/type/TypeFactoryTest.java | 17 +- .../runtime/operator/AggregateOperator.java | 4 + .../runtime/operator/HashJoinOperator.java | 78 +- .../operator/MultistageGroupByExecutor.java | 7 + .../query/runtime/operator/SortOperator.java | 6 + .../SortedMailboxReceiveOperator.java | 5 + .../operator/WindowAggregateOperator.java | 2 + .../factory/DefaultJoinOperatorFactory.java | 28 +- .../operator/operands/FilterOperand.java | 7 +- .../operands/LiteralParseJsonOperand.java | 95 + .../operands/TransformOperandFactory.java | 9 + .../operator/operands/VariantOperand.java | 221 +++ .../operator/set/BinarySetOperator.java | 1 + .../runtime/operator/set/SetOperator.java | 8 + .../runtime/operator/set/UnionOperator.java | 1 + .../operator/AggregateOperatorTest.java | 38 + .../operator/HashJoinOperatorTest.java | 69 +- .../runtime/operator/SortOperatorTest.java | 12 + .../SortedMailboxReceiveOperatorTest.java | 13 + .../operator/WindowAggregateOperatorTest.java | 39 + .../DefaultJoinOperatorFactoryTest.java | 105 ++ .../operator/operands/FilterOperandTest.java | 50 + .../operator/operands/VariantOperandTest.java | 271 +++ .../operator/set/IntersectOperatorTest.java | 13 + .../operator/set/UnionOperatorTest.java | 15 + .../creator/impl/BaseSegmentCreator.java | 5 +- .../creator/impl/ColumnarValueNormalizer.java | 24 +- .../AbstractColumnStatisticsCollector.java | 4 +- .../ColumnMinMaxValueGenerator.java | 10 +- .../utils/SanitizationTransformerUtils.java | 17 +- .../segment/local/utils/SchemaUtils.java | 19 + .../segment/local/utils/TableConfigUtils.java | 76 +- .../ExpressionTransformerTest.java | 36 + .../impl/ColumnarValueNormalizerTest.java | 43 + .../segment/index/ColumnMetadataTest.java | 12 + .../ColumnMinMaxValueGeneratorTest.java | 69 +- .../utils/VariantSchemaValidationTest.java | 134 ++ .../VariantTableConfigValidationTest.java | 300 +++ .../index/metadata/ColumnMetadataImpl.java | 6 + pinot-spi/VARIANT_DESIGN.md | 330 ++++ .../org/apache/pinot/spi/data/FieldSpec.java | 111 +- .../org/apache/pinot/spi/data/Schema.java | 6 +- .../apache/pinot/spi/utils/PinotDataType.java | 77 + .../pinot/spi/utils/VariantEnvelope.java | 263 +++ .../pinot/spi/data/VariantFieldSpecTest.java | 115 ++ .../pinot/spi/utils/PinotDataTypeTest.java | 15 + .../pinot/spi/utils/VariantEnvelopeTest.java | 263 +++ .../pinot/sql/ddl/compile/DataTypeMapper.java | 1 + .../pinot/sql/ddl/reverse/SchemaEmitter.java | 6 + .../pinot/sql/ddl/reverse/SqlIdentifiers.java | 2 +- .../sql/ddl/compile/DdlCompilerTest.java | 11 + .../ddl/reverse/CanonicalDdlEmitterTest.java | 14 + .../sql/ddl/roundtrip/RoundTripTest.java | 12 + pinot-tools/pom.xml | 10 + .../apache/pinot/tools/VariantQuickStart.java | 115 ++ .../examples/batch/variantEvents/README.md | 88 + .../batch/variantEvents/ingestionJobSpec.yaml | 51 + .../rawdata/variantEvents_data.parquet | Bin 0 -> 1442 bytes .../variantEvents_offline_table_config.json | 39 + .../variantEvents/variantEvents_schema.json | 26 + .../tools/admin/command/QuickStartTest.java | 7 + pom.xml | 5 + 149 files changed, 11351 insertions(+), 145 deletions(-) create mode 100644 pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java create mode 100644 pinot-clients/pinot-jdbc-client/src/test/resources/variant_result_table.json create mode 100644 pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java create mode 100644 pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java create mode 100644 pinot-common/src/test/java/org/apache/pinot/common/proto/ExpressionsColumnDataTypeCompatibilityTest.java create mode 100644 pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java create mode 100644 pinot-common/src/test/proto/legacy_expressions.proto create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunctionTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunctionTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunctionTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/plan/DistinctPlanNodeTest.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java create mode 100644 pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java create mode 100644 pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java create mode 100644 pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/VariantOperandTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizerTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantSchemaValidationTest.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java create mode 100644 pinot-spi/VARIANT_DESIGN.md create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java create mode 100644 pinot-spi/src/test/java/org/apache/pinot/spi/data/VariantFieldSpecTest.java create mode 100644 pinot-spi/src/test/java/org/apache/pinot/spi/utils/VariantEnvelopeTest.java create mode 100644 pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java create mode 100644 pinot-tools/src/main/resources/examples/batch/variantEvents/README.md create mode 100644 pinot-tools/src/main/resources/examples/batch/variantEvents/ingestionJobSpec.yaml create mode 100644 pinot-tools/src/main/resources/examples/batch/variantEvents/rawdata/variantEvents_data.parquet create mode 100644 pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_offline_table_config.json create mode 100644 pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_schema.json diff --git a/.github/workflows/scripts/.pinot_quickstart.sh b/.github/workflows/scripts/.pinot_quickstart.sh index dad7e5fb2587..82b4954d3fdf 100755 --- a/.github/workflows/scripts/.pinot_quickstart.sh +++ b/.github/workflows/scripts/.pinot_quickstart.sh @@ -319,6 +319,41 @@ if [ "${PASS}" -eq 0 ]; then exit 1 fi +# Test the packaged Parquet VARIANT quickstart entrypoint +bin/quick-start-variant-batch.sh & +PID=$! + +PASS=0 + +# Wait at most 5 minutes for the sample segment and a nested VARIANT value to be queryable +for i in $(seq 1 150) +do + QUERY_RES=$(curl -sS -X POST --header 'Accept: application/json' \ + -d "{\"sql\":\"SET enableNullHandling=true; SELECT COUNT(*) FROM variantEvents\",\"trace\":false}" \ + http://localhost:8000/query/sql) + if [ $? -eq 0 ]; then + COUNT_STAR_RES=$(echo "${QUERY_RES}" | jq '.resultTable.rows[0][0]') + VARIANT_QUERY_RES=$(curl -sS -X POST --header 'Accept: application/json' \ + -d "{\"sql\":\"SET enableNullHandling=true; SELECT variant_get(payload, '$.user.id', 'STRING') FROM variantEvents WHERE eventId = 'evt-001'\",\"trace\":false}" \ + http://localhost:8000/query/sql) + if [ $? -eq 0 ]; then + USER_ID_RES=$(echo "${VARIANT_QUERY_RES}" | jq -r '.resultTable.rows[0][0] // empty') + if [[ "${COUNT_STAR_RES}" =~ ^[0-9]+$ ]] && [ "${COUNT_STAR_RES}" -eq 5 ] \ + && [ "${USER_ID_RES}" = "u-1" ]; then + PASS=1 + break + fi + fi + fi + sleep 2 +done + +cleanup "${PID}" +if [ "${PASS}" -eq 0 ]; then + echo 'Parquet VARIANT quickstart failed: Cannot query the five sample rows and nested user ID.' + exit 1 +fi + # Test quick-start-streaming bin/quick-start-streaming.sh & PID=$! diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java index b806670e1ecf..0e15e411f0df 100644 --- a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java @@ -59,6 +59,12 @@ public String getColumnDataType(int columnIndex) { @Override public String getString(int rowIndex, int columnIndex) { JsonNode jsonValue = _rowsArray.get(rowIndex).get(columnIndex); + // Historically getString() exposes a JSON null as the string "null". Preserve that behavior for all established + // types. VARIANT needs to distinguish a SQL null (JSON null) from a Variant null (the canonical JSON string + // "null"), so only the new type maps a JSON null to Java null. + if (jsonValue.isNull() && "VARIANT".equals(getColumnDataType(columnIndex))) { + return null; + } if (jsonValue.isTextual()) { return jsonValue.textValue(); } else { @@ -125,7 +131,8 @@ public String toString() { String[] columnValues = new String[numColumns]; for (int c = 0; c < numColumns; c++) { try { - columnValues[c] = getString(r, c); + String value = getString(r, c); + columnValues[c] = value != null ? value : "null"; } catch (Exception e) { columnNames[c] = "ERROR"; } diff --git a/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/ResultTableResultSetTest.java b/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/ResultTableResultSetTest.java index e57514f71f49..955ee9f78396 100644 --- a/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/ResultTableResultSetTest.java +++ b/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/ResultTableResultSetTest.java @@ -27,6 +27,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotEquals; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; public class ResultTableResultSetTest { @@ -90,6 +91,24 @@ public void testGetString() { assertEquals("r1c1", result); } + @Test + public void testGetStringPreservesLegacyNullBehaviorAndDistinguishesVariantNulls() + throws Exception { + JsonNode resultTable = JsonUtils.stringToJsonNode( + "{\"rows\":[[null,null,\"null\"],[\"null\",7,\"\\\"null\\\"\"],[\"value\",8,null]]," + + "\"dataSchema\":{\"columnNames\":[\"legacyString\",\"legacyInt\",\"payload\"]," + + "\"columnDataTypes\":[\"STRING\",\"INT\",\"VARIANT\"]}}"); + ResultTableResultSet resultSet = new ResultTableResultSet(resultTable); + + assertEquals(resultSet.getString(0, 0), "null", "JSON null must retain the established STRING behavior"); + assertEquals(resultSet.getString(0, 1), "null", "JSON null must retain the established INT behavior"); + assertEquals(resultSet.getString(1, 0), "null", "A textual null must retain the established STRING behavior"); + assertEquals(resultSet.getString(0, 2), "null", "A Variant null is represented by canonical JSON"); + assertEquals(resultSet.getString(1, 2), "\"null\"", "A Variant string containing null remains quoted"); + assertNull(resultSet.getString(2, 2), "Only a SQL null maps to Java null for VARIANT"); + assertNotEquals(resultSet.toString(), ""); + } + @Test public void testGetAllColumns() { // Run the test diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotResultSet.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotResultSet.java index f62b52e497e1..186056f803bc 100644 --- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotResultSet.java +++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotResultSet.java @@ -297,7 +297,7 @@ public String getString(int columnIndex) validateColumn(columnIndex); String val = _resultSet.getString(_currentRow, columnIndex - 1); - if (checkIsNull(val)) { + if (checkIsNull(columnIndex, val)) { return null; } @@ -316,6 +316,7 @@ public Object getObject(int columnIndex) switch (dataType) { case "STRING": + case "VARIANT": return getString(columnIndex); case "INT": return getInt(columnIndex); @@ -352,8 +353,14 @@ public T getObject(String columnLabel, Class type) return super.getObject(columnLabel, type); } - private boolean checkIsNull(String val) { - if (val == null || val.toLowerCase().contentEquals(NULL_STRING)) { + private boolean checkIsNull(int columnIndex, String val) { + if (val == null) { + _wasNull = true; + return true; + } + // A VARIANT null is a non-SQL-null value whose canonical JSON representation is the text "null". + // The underlying ResultTableResultSet preserves a JSON null as Java null, so these two cases remain distinct. + if (!"VARIANT".equals(_columnDataTypes.get(columnIndex)) && val.equalsIgnoreCase(NULL_STRING)) { _wasNull = true; return true; } diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcResultSet.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcResultSet.java index a47b3fcfc44c..c9b6ab9d6567 100644 --- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcResultSet.java +++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcResultSet.java @@ -272,8 +272,13 @@ public short getShort(int columnIndex) public String getString(int columnIndex) throws SQLException { validateColumn(columnIndex); - String val = _currentRowBatch.getRows().get(_currentBatchIndex)[columnIndex - 1].toString(); - if (checkIsNull(val)) { + Object value = _currentRowBatch.getRows().get(_currentBatchIndex)[columnIndex - 1]; + if (value == null) { + _wasNull = true; + return null; + } + String val = value.toString(); + if (checkIsNull(columnIndex, val)) { return null; } return val; @@ -291,6 +296,7 @@ public Object getObject(int columnIndex) switch (dataType) { case "STRING": + case "VARIANT": return getString(columnIndex); case "INT": return getInt(columnIndex); @@ -327,8 +333,8 @@ public T getObject(String columnLabel, Class type) return super.getObject(columnLabel, type); } - private boolean checkIsNull(String val) { - if (val == null || val.toLowerCase().contentEquals(NULL_STRING)) { + private boolean checkIsNull(int columnIndex, String val) { + if (!"VARIANT".equals(_columnDataTypes.get(columnIndex)) && val.equalsIgnoreCase(NULL_STRING)) { _wasNull = true; return true; } diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java index 2e99c58bd17e..7f9cefed3962 100644 --- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java +++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/DriverUtils.java @@ -155,6 +155,7 @@ public static Integer getSQLDataType(String columnDataType) { Integer columnsSQLDataType; switch (columnDataType) { case "STRING": + case "VARIANT": columnsSQLDataType = Types.VARCHAR; break; case "INT": diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotResultSetTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotResultSetTest.java index 144256a7fca9..31a0eec81310 100644 --- a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotResultSetTest.java +++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotResultSetTest.java @@ -23,6 +23,7 @@ import java.nio.charset.StandardCharsets; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.Types; import java.util.Calendar; import java.util.Date; import java.util.List; @@ -33,6 +34,7 @@ import java.util.concurrent.atomic.AtomicReference; import org.apache.commons.io.IOUtils; import org.apache.pinot.client.utils.DateTimeUtils; +import org.apache.pinot.client.utils.DriverUtils; import org.apache.pinot.spi.utils.JsonUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -41,6 +43,7 @@ /// Tests deserialization of a ResultSet given hardcoded Pinot results. public class PinotResultSetTest { public static final String TEST_RESULT_SET_RESOURCE = "result_table.json"; + public static final String VARIANT_RESULT_SET_RESOURCE = "variant_result_table.json"; private DummyJsonTransport _dummyJsonTransport = new DummyJsonTransport(); @Test @@ -187,6 +190,64 @@ public void testGetResultMetadata() } } + @Test + public void testVariantUsesCanonicalJsonStringJdbcType() { + Assert.assertEquals(DriverUtils.getSQLDataType("VARIANT"), Integer.valueOf(Types.VARCHAR)); + } + + @Test + public void testVariantGetStringAndGetObjectPreserveVariantNullAndSqlNull() + throws Exception { + ResultSet resultSet = getResultSet(VARIANT_RESULT_SET_RESOURCE).getResultSet(0); + PinotResultSet pinotResultSet = new PinotResultSet(resultSet); + + Assert.assertTrue(pinotResultSet.next()); + Assert.assertEquals(pinotResultSet.getString(1), "{\"a\":[1,true,null],\"b\":\"text\"}"); + Assert.assertFalse(pinotResultSet.wasNull()); + Assert.assertEquals(pinotResultSet.getObject(1), "{\"a\":[1,true,null],\"b\":\"text\"}"); + Assert.assertFalse(pinotResultSet.wasNull()); + + Assert.assertTrue(pinotResultSet.next()); + Assert.assertEquals(pinotResultSet.getString(1), "null"); + Assert.assertFalse(pinotResultSet.wasNull(), "An encoded Variant null is not SQL null"); + Assert.assertEquals(pinotResultSet.getObject(1), "null"); + Assert.assertFalse(pinotResultSet.wasNull(), "An encoded Variant null is not SQL null"); + + Assert.assertTrue(pinotResultSet.next()); + Assert.assertEquals(pinotResultSet.getString(1), "\"null\""); + Assert.assertFalse(pinotResultSet.wasNull(), "A Variant string containing null is not SQL null"); + Assert.assertEquals(pinotResultSet.getObject(1), "\"null\""); + Assert.assertFalse(pinotResultSet.wasNull(), "A Variant string containing null is not SQL null"); + + Assert.assertTrue(pinotResultSet.next()); + Assert.assertNull(pinotResultSet.getString(1)); + Assert.assertTrue(pinotResultSet.wasNull()); + Assert.assertNull(pinotResultSet.getObject(1)); + Assert.assertTrue(pinotResultSet.wasNull()); + Assert.assertFalse(pinotResultSet.next()); + } + + @Test + public void testEstablishedTypesRetainNullHandling() + throws Exception { + PinotResultSet resultSet = PinotResultSet.fromJson( + "{\"resultTable\":{\"dataSchema\":{\"columnNames\":[\"stringValue\",\"intValue\"]," + + "\"columnDataTypes\":[\"STRING\",\"INT\"]},\"rows\":[[null,null],[\"NuLl\",7]]}}"); + + Assert.assertTrue(resultSet.next()); + Assert.assertNull(resultSet.getString(1)); + Assert.assertTrue(resultSet.wasNull()); + Assert.assertEquals(resultSet.getInt(2), 0); + Assert.assertTrue(resultSet.wasNull()); + + Assert.assertTrue(resultSet.next()); + Assert.assertNull(resultSet.getString(1), "STRING null matching remains case-insensitive"); + Assert.assertTrue(resultSet.wasNull()); + Assert.assertEquals(resultSet.getInt(2), 7); + Assert.assertFalse(resultSet.wasNull()); + Assert.assertFalse(resultSet.next()); + } + @Test public void testGetCalculatedScale() { PinotResultSet pinotResultSet = new PinotResultSet(); diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java new file mode 100644 index 000000000000..92123145ab6a --- /dev/null +++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java @@ -0,0 +1,101 @@ +/** + * 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.client.grpc; + +import com.google.protobuf.ByteString; +import java.io.IOException; +import java.util.Iterator; +import java.util.List; +import org.apache.pinot.common.proto.Broker; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.response.encoder.JsonResponseEncoder; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +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.assertTrue; + + +/** + * Verifies the JDBC result contract over the same metadata, schema, and data block sequence emitted by the gRPC + * broker endpoint. + */ +public class PinotGrpcResultSetTest { + + @Test + public void testVariantGetStringAndGetObjectPreserveVariantNullAndSqlNull() + throws Exception { + PinotGrpcResultSet resultSet = new PinotGrpcResultSet(createResponses()); + + assertTrue(resultSet.next()); + assertEquals(resultSet.getString(1), "{\"a\":[1,true,null],\"b\":\"text\"}"); + assertFalse(resultSet.wasNull()); + assertEquals(resultSet.getObject(1), "{\"a\":[1,true,null],\"b\":\"text\"}"); + assertFalse(resultSet.wasNull()); + + assertTrue(resultSet.next()); + assertEquals(resultSet.getString(1), "null"); + assertFalse(resultSet.wasNull(), "An encoded Variant null is not SQL null"); + assertEquals(resultSet.getObject(1), "null"); + assertFalse(resultSet.wasNull(), "An encoded Variant null is not SQL null"); + + assertTrue(resultSet.next()); + assertEquals(resultSet.getString(1), "\"null\""); + assertFalse(resultSet.wasNull(), "A Variant string containing null is not SQL null"); + assertEquals(resultSet.getObject(1), "\"null\""); + assertFalse(resultSet.wasNull(), "A Variant string containing null is not SQL null"); + + assertTrue(resultSet.next()); + assertNull(resultSet.getString(1)); + assertTrue(resultSet.wasNull()); + assertNull(resultSet.getObject(1)); + assertTrue(resultSet.wasNull()); + assertFalse(resultSet.next()); + } + + private static Iterator createResponses() + throws IOException { + DataSchema schema = new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + List rows = List.of( + new Object[]{"{\"a\":[1,true,null],\"b\":\"text\"}"}, + new Object[]{"null"}, + new Object[]{"\"null\""}, + new Object[]{null} + ); + byte[] payload = new JsonResponseEncoder().encodeResultTable( + new ResultTable(schema, rows), 0, rows.size()); + + Broker.BrokerResponse metadata = Broker.BrokerResponse.newBuilder() + .setPayload(ByteString.copyFromUtf8("{}")) + .build(); + Broker.BrokerResponse schemaBlock = Broker.BrokerResponse.newBuilder() + .setPayload(ByteString.copyFrom(schema.toBytes())) + .build(); + Broker.BrokerResponse dataBlock = Broker.BrokerResponse.newBuilder() + .setPayload(ByteString.copyFrom(payload)) + .putMetadata("rowSize", Integer.toString(rows.size())) + .putMetadata("compression", "NONE") + .putMetadata("encoding", "JSON") + .build(); + return List.of(metadata, schemaBlock, dataBlock).iterator(); + } +} diff --git a/pinot-clients/pinot-jdbc-client/src/test/resources/variant_result_table.json b/pinot-clients/pinot-jdbc-client/src/test/resources/variant_result_table.json new file mode 100644 index 000000000000..00bd123766eb --- /dev/null +++ b/pinot-clients/pinot-jdbc-client/src/test/resources/variant_result_table.json @@ -0,0 +1,35 @@ +{ + "resultTable": { + "dataSchema": { + "columnNames": [ + "payload" + ], + "columnDataTypes": [ + "VARIANT" + ] + }, + "rows": [ + [ + "{\"a\":[1,true,null],\"b\":\"text\"}" + ], + [ + "null" + ], + [ + "\"null\"" + ], + [ + null + ] + ] + }, + "exceptions": [], + "numServersQueried": 1, + "numServersResponded": 1, + "numSegmentsQueried": 1, + "numSegmentsProcessed": 1, + "numSegmentsMatched": 1, + "numDocsScanned": 4, + "totalDocs": 4, + "timeUsedMs": 1 +} diff --git a/pinot-common/pom.xml b/pinot-common/pom.xml index 539625d7f2e2..f9789fc77bfb 100644 --- a/pinot-common/pom.xml +++ b/pinot-common/pom.xml @@ -200,6 +200,22 @@ org.apache.pinot pinot-timeseries-spi + + org.apache.parquet + parquet-variant + + + + org.apache.parquet + parquet-column + + + org.apache.parquet + parquet-common + + + org.apache.arrow diff --git a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java index b2fd5549255f..f1c5be4551bd 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/datablock/DataBlockEquals.java @@ -288,6 +288,7 @@ public boolean equals(DataBlock left, DataBlock right) { } break; case BYTES: + case VARIANT: for (int did = 0; did < numRows; did++) { if (!left.getBytes(did, colId).equals(right.getBytes(did, colId))) { if (_failOnFalse) { diff --git a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java index c72b1bdd4de4..62d0ebf04798 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java @@ -22,16 +22,23 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.common.function.FunctionInfo; import org.apache.pinot.common.function.FunctionInvoker; import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.FunctionUtils; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.FunctionContext; import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.common.utils.VariantUtils.ResultType; +import org.apache.pinot.common.utils.VariantUtils.ReusableResult; +import org.apache.pinot.common.utils.VariantUtils.VariantPath; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.function.FunctionEvaluator; +import org.apache.pinot.spi.utils.PinotDataType; /// Evaluates an expression. @@ -103,6 +110,11 @@ private ExecutableNode planExecution(ExpressionContext expression) { throw new IllegalStateException(String.format("Unsupported function: %s", functionName)); } } + ExecutableNode variantExecutionNode = + VariantExecutionNode.tryCreate(functionName, canonicalName, arguments, childNodes); + if (variantExecutionNode != null) { + return variantExecutionNode; + } return new FunctionExecutionNode(functionInfo, childNodes); } default: @@ -130,6 +142,148 @@ public String toString() { return _functionExpression; } + /** + * Planned ingestion evaluator for Variant scalar functions with literal path and target-type operands. + * + *

Each node compiles its literals once and owns a reusable cursor result. Like the enclosing evaluator, instances + * are intended to be confined to the record-transformer thread and are not thread-safe. + */ + private static class VariantExecutionNode implements ExecutableNode { + private static final String VARIANT_GET = "variantget"; + private static final String TRY_VARIANT_GET = "tryvariantget"; + private static final String VARIANT_EXISTS = "variantexists"; + private static final String IS_VARIANT_NULL = "isvariantnull"; + private static final String VARIANT_TYPE_OF = "varianttypeof"; + + private final String _functionName; + private final ExecutableNode _variantNode; + private final VariantOperation _operation; + private final VariantPath _path; + @Nullable + private final ResultType _targetType; + private final ReusableResult _reusableResult = new ReusableResult(); + + private VariantExecutionNode(String functionName, ExecutableNode variantNode, VariantOperation operation, + VariantPath path, @Nullable ResultType targetType) { + _functionName = functionName; + _variantNode = variantNode; + _operation = operation; + _path = path; + _targetType = targetType; + } + + @Nullable + static ExecutableNode tryCreate(String functionName, String canonicalName, List arguments, + ExecutableNode[] childNodes) { + int numArguments = arguments.size(); + switch (canonicalName) { + case VARIANT_GET: + case TRY_VARIANT_GET: + if ((numArguments != 2 && numArguments != 3) || !hasStringLiteral(arguments, 1) + || (numArguments == 3 && !hasStringLiteral(arguments, 2))) { + return null; + } + return new VariantExecutionNode(functionName, childNodes[0], + canonicalName.equals(VARIANT_GET) ? VariantOperation.GET : VariantOperation.TRY_GET, + VariantUtils.compilePath(stringLiteral(arguments, 1)), + numArguments == 2 ? ResultType.VARIANT + : VariantUtils.parseResultType(stringLiteral(arguments, 2))); + case VARIANT_EXISTS: + if (numArguments != 2 || !hasStringLiteral(arguments, 1)) { + return null; + } + return new VariantExecutionNode(functionName, childNodes[0], VariantOperation.EXISTS, + VariantUtils.compilePath(stringLiteral(arguments, 1)), null); + case IS_VARIANT_NULL: + case VARIANT_TYPE_OF: + if ((numArguments != 1 && numArguments != 2) + || (numArguments == 2 && !hasStringLiteral(arguments, 1))) { + return null; + } + return new VariantExecutionNode(functionName, childNodes[0], + canonicalName.equals(IS_VARIANT_NULL) ? VariantOperation.IS_NULL : VariantOperation.TYPE_OF, + VariantUtils.compilePath(numArguments == 1 ? "$" : stringLiteral(arguments, 1)), null); + default: + return null; + } + } + + private static boolean hasStringLiteral(List arguments, int index) { + ExpressionContext argument = arguments.get(index); + return argument.getType() == ExpressionContext.Type.LITERAL + && argument.getLiteral().getValue() instanceof String; + } + + private static String stringLiteral(List arguments, int index) { + return (String) arguments.get(index).getLiteral().getValue(); + } + + @Override + public Object execute(GenericRow row) { + return executeVariant(_variantNode.execute(row)); + } + + @Override + public Object execute(Object[] values) { + return executeVariant(_variantNode.execute(values)); + } + + @Nullable + private Object executeVariant(@Nullable Object input) { + try { + byte[] variant = toBytes(input); + switch (_operation) { + case GET: + return extract(variant, false); + case TRY_GET: + return extract(variant, true); + case EXISTS: + return VariantUtils.variantExists(variant, _path, _reusableResult); + case IS_NULL: + return VariantUtils.isVariantNull(variant, _path, _reusableResult); + case TYPE_OF: + return VariantUtils.variantTypeOf(variant, _path, _reusableResult); + default: + throw new IllegalStateException("Unhandled Variant operation: " + _operation); + } + } catch (Exception e) { + throw new RuntimeException( + "Caught exception while executing function: " + _functionName + ": " + e.getMessage(), e); + } + } + + @Nullable + private Object extract(@Nullable byte[] variant, boolean tolerant) { + ResultType targetType = Preconditions.checkNotNull(_targetType, "Variant target type must be planned"); + boolean present = tolerant + ? VariantUtils.tryExtractInto(variant, _path, targetType, _reusableResult) + : VariantUtils.extractInto(variant, _path, targetType, _reusableResult); + if (!present) { + return null; + } + return _reusableResult.getExternalValue(targetType); + } + + @Nullable + private static byte[] toBytes(@Nullable Object input) { + if (input == null) { + return null; + } + if (input instanceof byte[]) { + return (byte[]) input; + } + return (byte[]) PinotDataType.BYTES.convert(input, FunctionUtils.getArgumentType(input)); + } + } + + private enum VariantOperation { + GET, + TRY_GET, + EXISTS, + IS_NULL, + TYPE_OF + } + private interface ExecutableNode { Object execute(GenericRow row); diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java b/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java index f88b73cf7404..5d7da4751012 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java @@ -26,6 +26,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; @@ -36,7 +37,6 @@ import org.apache.calcite.sql.type.SqlOperandTypeChecker; import org.apache.calcite.sql.type.SqlReturnTypeInference; import org.apache.calcite.sql.type.SqlTypeFamily; -import org.apache.commons.lang3.StringUtils; import org.apache.pinot.common.function.sql.PinotSqlFunction; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.spi.annotations.ScalarFunction; @@ -216,7 +216,7 @@ public static FunctionInfo lookupFunctionInfo(String canonicalName, int numArgum } public static String canonicalize(String name) { - return StringUtils.remove(name, '_').toLowerCase(); + return name.replace("_", "").toLowerCase(Locale.ROOT); } public static class ArgumentCountBasedScalarFunction implements PinotScalarFunction { diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java index 23c972d26849..5afc6cc7a48d 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java @@ -20,7 +20,10 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.HashSet; import java.util.List; +import java.util.Locale; +import java.util.Set; import javax.annotation.Nullable; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; @@ -34,6 +37,8 @@ import org.apache.calcite.sql.type.SqlTypeFamily; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.sql.type.SqlTypeTransforms; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.common.utils.VariantUtils.ResultType; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DateTimeFormatSpec; @@ -117,6 +122,17 @@ public enum TransformFunctionType { JSON_EXTRACT_KEY("jsonExtractKey", ReturnTypes.TO_ARRAY, OperandTypes.family( List.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), i -> i > 1)), + VARIANT_GET("variantGet", TransformFunctionType::variantGetReturnTypeInference, variantGetOperandTypeChecker()), + TRY_VARIANT_GET("tryVariantGet", TransformFunctionType::variantGetReturnTypeInference, + variantGetOperandTypeChecker()), + VARIANT_EXISTS("variantExists", ReturnTypes.BOOLEAN_NULLABLE, variantPathOperandTypeChecker()), + IS_VARIANT_NULL("isVariantNull", ReturnTypes.BOOLEAN, optionalVariantPathOperandTypeChecker()), + VARIANT_TYPE_OF("variantTypeOf", ReturnTypes.VARCHAR_2000_NULLABLE, optionalVariantPathOperandTypeChecker()), + VARIANT_TO_JSON("variantToJson", ReturnTypes.VARCHAR_2000_NULLABLE, OperandTypes.ANY), + PARSE_JSON_TO_VARIANT("parseJson", TransformFunctionType::nullableVariantReturnTypeInference, + OperandTypes.CHARACTER, "parseJsonToVariant"), + TRY_PARSE_JSON_TO_VARIANT("tryParseJson", TransformFunctionType::nullableVariantReturnTypeInference, + OperandTypes.CHARACTER, "tryParseJsonToVariant"), // Date time functions TIME_CONVERT("timeConvert", ReturnTypes.BIGINT, @@ -259,6 +275,8 @@ public enum TransformFunctionType { // Time series functions TIME_SERIES_BUCKET("timeSeriesBucket"); + private static final Set NULL_HANDLING_REQUIRED_FUNCTION_NAMES = + createNullHandlingRequiredFunctionNames(); private final String _name; private final List _names; private final SqlReturnTypeInference _returnTypeInference; @@ -300,6 +318,30 @@ public SqlOperandTypeChecker getOperandTypeChecker() { return _operandTypeChecker; } + /// Returns whether the function depends on SQL-null semantics and therefore requires query null handling. + /// + /// The check accepts every registered spelling, including underscore-insensitive aliases, so all query engines + /// can enforce the same contract before selecting an implementation. + public static boolean requiresNullHandling(String functionName) { + return NULL_HANDLING_REQUIRED_FUNCTION_NAMES.contains(canonicalize(functionName)); + } + + private static Set createNullHandlingRequiredFunctionNames() { + Set names = new HashSet<>(); + for (TransformFunctionType functionType + : List.of(VARIANT_GET, TRY_VARIANT_GET, VARIANT_EXISTS, IS_VARIANT_NULL, VARIANT_TYPE_OF, VARIANT_TO_JSON, + PARSE_JSON_TO_VARIANT, TRY_PARSE_JSON_TO_VARIANT)) { + for (String name : functionType.getNames()) { + names.add(canonicalize(name)); + } + } + return Set.copyOf(names); + } + + private static String canonicalize(String functionName) { + return FunctionRegistry.canonicalize(functionName); + } + /// Returns the optional explicit returning type specification. private static RelDataType positionalReturnTypeInferenceFromStringLiteral(SqlOperatorBinding opBinding, int pos) { return positionalReturnTypeInferenceFromStringLiteral(opBinding, pos, SqlTypeName.ANY); @@ -308,7 +350,7 @@ private static RelDataType positionalReturnTypeInferenceFromStringLiteral(SqlOpe private static RelDataType positionalReturnTypeInferenceFromStringLiteral(SqlOperatorBinding opBinding, int pos, SqlTypeName defaultSqlType) { if (opBinding.getOperandCount() > pos && opBinding.isOperandLiteral(pos, false)) { - String operandType = opBinding.getOperandLiteralValue(pos, String.class).toUpperCase(); + String operandType = opBinding.getOperandLiteralValue(pos, String.class).toUpperCase(Locale.ROOT); return inferTypeFromStringLiteral(operandType, opBinding.getTypeFactory()); } return opBinding.getTypeFactory().createSqlType(defaultSqlType); @@ -316,7 +358,7 @@ private static RelDataType positionalReturnTypeInferenceFromStringLiteral(SqlOpe private static RelDataType jsonExtractScalarReturnTypeInference(SqlOperatorBinding opBinding) { if (opBinding.getOperandCount() > 2 && opBinding.isOperandLiteral(2, false)) { - String resultsType = opBinding.getOperandLiteralValue(2, String.class).toUpperCase(); + String resultsType = opBinding.getOperandLiteralValue(2, String.class).toUpperCase(Locale.ROOT); RelDataTypeFactory typeFactory = opBinding.getTypeFactory(); switch (resultsType) { case "JSON": @@ -332,6 +374,47 @@ private static RelDataType jsonExtractScalarReturnTypeInference(SqlOperatorBindi return positionalReturnTypeInferenceFromStringLiteral(opBinding, 2, SqlTypeName.VARCHAR); } + private static RelDataType nullableVariantReturnTypeInference(SqlOperatorBinding opBinding) { + RelDataType variantType = opBinding.getTypeFactory().createSqlType(SqlTypeName.VARIANT); + return opBinding.getTypeFactory().createTypeWithNullability(variantType, true); + } + + private static RelDataType variantGetReturnTypeInference(SqlOperatorBinding opBinding) { + if (opBinding.getOperandCount() == 2) { + return nullableVariantReturnTypeInference(opBinding); + } + if (!opBinding.isOperandLiteral(2, false)) { + throw new IllegalArgumentException("variantGet target type must be a string literal"); + } + ResultType targetType = + VariantUtils.parseResultType(opBinding.getOperandLiteralValue(2, String.class)); + RelDataType resultType = opBinding.getTypeFactory().createSqlType(targetType.getSqlTypeName()); + return opBinding.getTypeFactory().createTypeWithNullability(resultType, true); + } + + private static SqlOperandTypeChecker variantGetOperandTypeChecker() { + SqlSingleOperandTypeChecker stringLiteral = OperandTypes.and(OperandTypes.CHARACTER, OperandTypes.LITERAL); + return OperandTypes.or( + OperandTypes.sequence( + (operator, ignored) -> "'" + operator.getName() + "(, )'", + OperandTypes.ANY, stringLiteral), + OperandTypes.sequence( + (operator, ignored) -> "'" + operator.getName() + + "(, , )'", + OperandTypes.ANY, stringLiteral, stringLiteral)); + } + + private static SqlOperandTypeChecker variantPathOperandTypeChecker() { + SqlSingleOperandTypeChecker stringLiteral = OperandTypes.and(OperandTypes.CHARACTER, OperandTypes.LITERAL); + return OperandTypes.sequence( + (operator, ignored) -> "'" + operator.getName() + "(, )'", + OperandTypes.ANY, stringLiteral); + } + + private static SqlOperandTypeChecker optionalVariantPathOperandTypeChecker() { + return OperandTypes.or(OperandTypes.ANY, variantPathOperandTypeChecker()); + } + /// Operand checker shared by `jsonExtractScalar` and its `Fast` / `FirstMatch` variants. /// /// `jsonPath` deliberately does **not** require [OperandTypes#LITERAL]. Operand checking runs on the raw diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java new file mode 100644 index 000000000000..a5d968319270 --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java @@ -0,0 +1,102 @@ +/** + * 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.function.scalar; + +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.spi.annotations.ScalarFunction; + + +/** + * Scalar functions for the Pinot {@code VARIANT} logical type. + */ +public final class VariantFunctions { + private VariantFunctions() { + } + + @Nullable + @ScalarFunction + public static byte[] variantGet(byte[] variant, String path) { + return VariantUtils.variantGet(variant, path); + } + + @Nullable + @ScalarFunction + public static Object variantGet(byte[] variant, String path, String targetType) { + return VariantUtils.variantGet(variant, path, targetType); + } + + @Nullable + @ScalarFunction + public static byte[] tryVariantGet(byte[] variant, String path) { + return VariantUtils.tryVariantGet(variant, path); + } + + @Nullable + @ScalarFunction + public static Object tryVariantGet(byte[] variant, String path, String targetType) { + return VariantUtils.tryVariantGet(variant, path, targetType); + } + + @Nullable + @ScalarFunction + public static Boolean variantExists(byte[] variant, String path) { + return VariantUtils.variantExists(variant, path); + } + + @ScalarFunction(nullableParameters = true) + public static boolean isVariantNull(@Nullable byte[] variant) { + return VariantUtils.isVariantNull(variant); + } + + @ScalarFunction(nullableParameters = true) + public static boolean isVariantNull(@Nullable byte[] variant, String path) { + return VariantUtils.isVariantNull(variant, path); + } + + @Nullable + @ScalarFunction + public static String variantTypeOf(byte[] variant) { + return VariantUtils.variantTypeOf(variant); + } + + @Nullable + @ScalarFunction + public static String variantTypeOf(byte[] variant, String path) { + return VariantUtils.variantTypeOf(variant, path); + } + + @Nullable + @ScalarFunction + public static String variantToJson(byte[] variant) { + return VariantUtils.variantToJson(variant); + } + + @Nullable + @ScalarFunction(names = {"parseJson", "parseJsonToVariant"}) + public static byte[] parseJsonToVariant(String json) { + return VariantUtils.parseJsonToVariant(json); + } + + @Nullable + @ScalarFunction(names = {"tryParseJson", "tryParseJsonToVariant"}) + public static byte[] tryParseJsonToVariant(String json) { + return VariantUtils.tryParseJsonToVariant(json); + } +} diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java index efd9098991e3..463396ea8f0a 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoder.java @@ -107,6 +107,7 @@ private VectorSchemaRoot createVectorSchemaRoot(ResultTable resultTable, DataSch case STRING: case JSON: case BYTES: + case VARIANT: case UUID: case OBJECT: field = new Field(colName, FieldType.nullable(new ArrowType.Utf8()), null); @@ -236,6 +237,7 @@ private VectorSchemaRoot createVectorSchemaRoot(ResultTable resultTable, DataSch case STRING: case JSON: case BYTES: + case VARIANT: case UUID: case OBJECT: byte[] bytes = ((String) value).getBytes(StandardCharsets.UTF_8); @@ -415,6 +417,7 @@ public ResultTable decodeResultTable(byte[] bytes, int rowSize, DataSchema schem case STRING: case JSON: case BYTES: + case VARIANT: case UUID: case OBJECT: row[col] = new String(((VarCharVector) vector).get(i), StandardCharsets.UTF_8); diff --git a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java index aa86400de852..16b8f85d0c95 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/response/encoder/JsonResponseEncoder.java @@ -225,6 +225,7 @@ private static Object extractValue(DataSchema.ColumnDataType columnDataType, Jso case STRING: case JSON: case BYTES: + case VARIANT: case UUID: case OBJECT: return jsonValue.textValue(); 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 52312c2a741b..8464d0329ea7 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 @@ -207,11 +207,11 @@ public static DataSchema fromBytes(PinotInputStream buffer) /// Resolves a [ColumnDataType] token read off the wire, turning the raw [IllegalArgumentException] from /// [ColumnDataType#valueOf] into a message that names the mixed-version cause. /// - /// Rolling-upgrade limitation: once a node on this build emits a `UUID` token, an older peer that does not know - /// the [ColumnDataType#UUID] constant fails here. There is no version-negotiation shim or fallback to `BYTES` - /// today, so brokers and servers must be upgraded atomically (or UUID columns kept out of queries) until the - /// whole cluster is on this build. Rolling back to a pre-UUID build is likewise unsafe while UUID-typed query - /// results are in flight. No existing (non-UUID) column is affected. + /// Rolling-upgrade limitation: once a node emits a token introduced by a newer build, an older peer that does not + /// know that [ColumnDataType] fails here. There is no version-negotiation shim or fallback to the stored type today, + /// so columns using a newly introduced logical type must stay out of queries until every broker and server has been + /// upgraded. Rolling back is likewise unsafe while results of that logical type are in flight. Existing column types + /// are unaffected because their wire names remain unchanged. private static ColumnDataType parseColumnDataType(String name) { try { return ColumnDataType.valueOf(name); @@ -329,6 +329,14 @@ public RelDataType toType(RelDataTypeFactory typeFactory) { return typeFactory.createSqlType(SqlTypeName.VARBINARY); } }, + // VARIANT is a logical type backed by BYTES. Its nonempty PVAR envelope distinguishes a Variant null from the + // empty SQL-null placeholder. + VARIANT(BYTES, NullValuePlaceHolder.INTERNAL_BYTES) { + @Override + public RelDataType toType(RelDataTypeFactory typeFactory) { + return typeFactory.createSqlType(SqlTypeName.VARIANT); + } + }, // UUID is a logical type backed by BYTES; keep it directly after BYTES. ColumnDataType is serialized by name (not // ordinal) via name()/valueOf(), so enum order does not affect wire compatibility. UUID(BYTES, null) { @@ -465,6 +473,56 @@ public ColumnDataType getStoredType() { return _storedColumnDataType; } + public boolean supportsEquality() { + if (isArray() || this == OBJECT) { + return false; + } + return toCapabilityDataType().supportsEquality(); + } + + public boolean supportsHashing() { + if (isArray() || this == OBJECT) { + return false; + } + return toCapabilityDataType().supportsHashing(); + } + + public boolean supportsOrdering() { + if (isArray() || this == OBJECT) { + return false; + } + return toCapabilityDataType().supportsOrdering(); + } + + public boolean supportsMinMax() { + if (isArray() || this == OBJECT) { + return false; + } + return toCapabilityDataType().supportsMinMax(); + } + + public boolean supportsDirectAggregation() { + return toCapabilityDataType().supportsDirectAggregation(); + } + + public boolean supportsPatternMatching() { + if (isArray() || this == MAP || this == OBJECT) { + return false; + } + return toCapabilityDataType().supportsPatternMatching(); + } + + private DataType toCapabilityDataType() { + switch (this) { + case MAP: + return DataType.MAP; + case OBJECT: + return DataType.OPEN_STRUCT; + default: + return toDataType(); + } + } + public boolean isNumber() { return NUMERIC_TYPES.contains(this); } @@ -522,6 +580,8 @@ public DataType toDataType() { case BYTES: case BYTES_ARRAY: return DataType.BYTES; + case VARIANT: + return DataType.VARIANT; case UUID: case UUID_ARRAY: return DataType.UUID; @@ -541,13 +601,14 @@ public DataType toDataType() { /// - Query response /// /// Internal value type is used within the storage and query engine, where value is always of the stored type. For - /// BYTES type, we use a wrapper class [ByteArray] to make it comparable. + /// BYTES and VARIANT types, we use a wrapper class [ByteArray] as their internal representation. /// /// The conversion applies to the following types: /// /// - BOOLEAN: boolean -> int /// - TIMESTAMP: Timestamp -> long /// - BYTES: byte\[\] -> ByteArray + /// - VARIANT: byte\[\] -> ByteArray /// - BOOLEAN_ARRAY: boolean\[\] -> int\[\] /// - TIMESTAMP_ARRAY: Timestamp\[\] -> long\[\] /// - UUID: UUID/String/byte\[\]/ByteArray -> ByteArray @@ -559,6 +620,7 @@ public Object toInternal(Object value) { case TIMESTAMP: return ((Timestamp) value).getTime(); case BYTES: + case VARIANT: return new ByteArray((byte[]) value); case UUID: return new ByteArray(UuidUtils.toBytes(value)); @@ -602,6 +664,7 @@ public Object toInternal(Object value) { /// - BOOLEAN: int -> boolean /// - TIMESTAMP: long -> Timestamp /// - BYTES: ByteArray -> byte\[\] + /// - VARIANT: ByteArray -> byte\[\] /// - BOOLEAN_ARRAY: int\[\] -> boolean\[\] /// - TIMESTAMP_ARRAY: long\[\] -> Timestamp\[\] /// - BYTES_ARRAY: ByteArray\[\] -> byte\[\]\[\] @@ -614,6 +677,7 @@ public Object toExternal(Object value) { case TIMESTAMP: return new Timestamp((long) value); case BYTES: + case VARIANT: return ((ByteArray) value).getBytes(); case UUID: return UuidUtils.toUUID((ByteArray) value); @@ -652,6 +716,7 @@ public Serializable convert(Object value) { case JSON: return value.toString(); case BYTES: + case VARIANT: return ((ByteArray) value).getBytes(); case UUID: return UuidUtils.toUUID((ByteArray) value); @@ -695,6 +760,8 @@ public Serializable format(Object value) { return value.toString(); case BYTES: return BytesUtils.toHexString((byte[]) value); + case VARIANT: + return VariantUtils.variantToJson((byte[]) value); case UUID: return formatUuid(value); case BIG_DECIMAL_ARRAY: @@ -732,6 +799,8 @@ public Serializable convertAndFormat(Object value) { return value.toString(); case BYTES: return ((ByteArray) value).toHexString(); + case VARIANT: + return VariantUtils.variantToJson(((ByteArray) value).getBytes()); case UUID: return UuidUtils.toString((ByteArray) value); case MAP: @@ -1060,6 +1129,8 @@ public static ColumnDataType fromDataTypeSV(DataType dataType) { return JSON; case BYTES: return BYTES; + case VARIANT: + return VARIANT; case UUID: return UUID; case MAP: @@ -1124,6 +1195,8 @@ public PinotDataType toPinotDataType() { return PinotDataType.JSON; case BYTES: return PinotDataType.BYTES; + case VARIANT: + return PinotDataType.VARIANT; case UUID: return PinotDataType.UUID; case MAP: diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java new file mode 100644 index 000000000000..38fb47ac2b13 --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java @@ -0,0 +1,1624 @@ +/** + * 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; + +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.core.JsonToken; +import java.io.IOException; +import java.io.StringWriter; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.sql.Timestamp; +import java.time.Instant; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.LocalTime; +import java.time.ZoneOffset; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import javax.annotation.Nullable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantArrayBuilder; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.parquet.variant.VariantObjectBuilder; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.VariantEnvelope; + + +/** + * Query-side operations for Pinot {@code VARIANT} values. + * + *

The utility navigates the Parquet Variant binary representation directly. It never materializes a JSON tree. + * Instances are not required, and stateless convenience methods are thread-safe. Overloads that accept a + * caller-provided {@link ReusableResult} require that result to be thread-confined and not shared by concurrent calls. + * An empty byte array is Pinot's SQL-null placeholder and is never decoded as an envelope. + */ +public final class VariantUtils { + private static final JsonFactory JSON_FACTORY = new JsonFactory(); + private static final BigDecimal MIN_INT_DECIMAL = BigDecimal.valueOf(Integer.MIN_VALUE); + private static final BigDecimal MAX_INT_DECIMAL = BigDecimal.valueOf(Integer.MAX_VALUE); + private static final BigDecimal MIN_LONG_DECIMAL = BigDecimal.valueOf(Long.MIN_VALUE); + private static final BigDecimal MAX_LONG_DECIMAL = BigDecimal.valueOf(Long.MAX_VALUE); + private static final int MAX_JSON_NESTING_DEPTH = 100; + private static final int MAX_VARIANT_DECIMAL_PRECISION = 38; + private static final int MAX_VARIANT_DECIMAL_SCALE = 38; + private static final int MAX_VARIANT_DECIMAL_BYTES = 16; + private static final long MICROS_PER_SECOND = TimeUnit.SECONDS.toMicros(1); + private static final long NANOS_PER_MICRO = TimeUnit.MICROSECONDS.toNanos(1); + private static final long NANOS_PER_DAY = TimeUnit.DAYS.toNanos(1); + private static final int VARIANT_BASIC_TYPE_MASK = 0x03; + private static final int VARIANT_PRIMITIVE_TYPE_MASK = 0x3F; + private static final int VARIANT_PRIMITIVE = 0; + private static final int VARIANT_SHORT_STRING = 1; + private static final int VARIANT_OBJECT = 2; + private static final int VARIANT_ARRAY = 3; + private static final int VARIANT_NULL = 0; + private static final int VARIANT_TRUE = 1; + private static final int VARIANT_FALSE = 2; + private static final int VARIANT_INT8 = 3; + private static final int VARIANT_INT16 = 4; + private static final int VARIANT_INT32 = 5; + private static final int VARIANT_INT64 = 6; + private static final int VARIANT_DOUBLE = 7; + private static final int VARIANT_DECIMAL4 = 8; + private static final int VARIANT_DECIMAL8 = 9; + private static final int VARIANT_DECIMAL16 = 10; + private static final int VARIANT_DATE = 11; + private static final int VARIANT_TIMESTAMP_TZ = 12; + private static final int VARIANT_TIMESTAMP_NTZ = 13; + private static final int VARIANT_FLOAT = 14; + private static final int VARIANT_BINARY = 15; + private static final int VARIANT_LONG_STRING = 16; + private static final int VARIANT_TIME = 17; + private static final int VARIANT_TIMESTAMP_NANOS_TZ = 18; + private static final int VARIANT_TIMESTAMP_NANOS_NTZ = 19; + private static final int VARIANT_UUID = 20; + private static final int VARIANT_METADATA_VERSION_MASK = 0x0F; + private static final int VARIANT_METADATA_VERSION = 1; + private static final VariantPath ROOT_PATH = new VariantPath(new PathElement[0]); + + private VariantUtils() { + } + + /** + * Statically supported result types for {@code variantGet} and {@code tryVariantGet}. + */ + public enum ResultType { + BOOLEAN(DataType.BOOLEAN, SqlTypeName.BOOLEAN), + INT(DataType.INT, SqlTypeName.INTEGER), + LONG(DataType.LONG, SqlTypeName.BIGINT), + FLOAT(DataType.FLOAT, SqlTypeName.REAL), + DOUBLE(DataType.DOUBLE, SqlTypeName.DOUBLE), + BIG_DECIMAL(DataType.BIG_DECIMAL, SqlTypeName.DECIMAL), + STRING(DataType.STRING, SqlTypeName.VARCHAR), + BYTES(DataType.BYTES, SqlTypeName.VARBINARY), + UUID(DataType.UUID, SqlTypeName.UUID), + TIMESTAMP(DataType.TIMESTAMP, SqlTypeName.TIMESTAMP), + VARIANT(DataType.VARIANT, SqlTypeName.VARIANT), + JSON(DataType.JSON, SqlTypeName.VARCHAR); + + private final DataType _dataType; + private final SqlTypeName _sqlTypeName; + + ResultType(DataType dataType, SqlTypeName sqlTypeName) { + _dataType = dataType; + _sqlTypeName = sqlTypeName; + } + + public DataType getDataType() { + return _dataType; + } + + public SqlTypeName getSqlTypeName() { + return _sqlTypeName; + } + } + + /** + * An immutable, pre-parsed Variant path. The v1 grammar supports {@code $}, dot-separated object fields, and + * non-negative array subscripts. + */ + public static final class VariantPath { + private final PathElement[] _elements; + + private VariantPath(PathElement[] elements) { + _elements = elements; + } + } + + /** + * Reusable, unboxed destination for vectorized Variant extraction. + * + *

Only the getter corresponding to the requested {@link ResultType} is defined after a successful extraction. + * The instance is mutable and not thread-safe; callers should retain one per transform-function instance. + */ + public static final class ReusableResult { + private final Cursor _cursor = new Cursor(); + private int _intValue; + private long _longValue; + private float _floatValue; + private double _doubleValue; + private BigDecimal _bigDecimalValue; + private String _stringValue; + private byte[] _bytesValue; + + public int getIntValue() { + return _intValue; + } + + public long getLongValue() { + return _longValue; + } + + public float getFloatValue() { + return _floatValue; + } + + public double getDoubleValue() { + return _doubleValue; + } + + public BigDecimal getBigDecimalValue() { + return _bigDecimalValue; + } + + public String getStringValue() { + return _stringValue; + } + + /** + * Returns the extracted BYTES, VARIANT, or direct 16-byte UUID representation. + */ + public byte[] getBytesValue() { + return _bytesValue; + } + + public UUID getUuidValue() { + return UuidUtils.toUUID(_bytesValue); + } + + /** + * Materializes the extracted value in the external representation used by scalar functions and ingestion. + */ + public Object getExternalValue(ResultType resultType) { + switch (resultType) { + case BOOLEAN: + return _intValue != 0; + case INT: + return _intValue; + case LONG: + return _longValue; + case FLOAT: + return _floatValue; + case DOUBLE: + return _doubleValue; + case BIG_DECIMAL: + return _bigDecimalValue; + case STRING: + case JSON: + return _stringValue; + case BYTES: + case VARIANT: + return _bytesValue; + case UUID: + return UuidUtils.toUUID(_bytesValue); + case TIMESTAMP: + return new Timestamp(_longValue); + default: + throw new IllegalStateException("Unhandled Variant target type: " + resultType); + } + } + + /** + * Materializes the extracted value in {@link DataSchema}'s internal representation. + * + *

TIMESTAMP remains epoch milliseconds and UUID wraps the directly copied 16-byte value, avoiding an + * external-object round trip in the multi-stage engine. + */ + public Object getInternalValue(ResultType resultType) { + switch (resultType) { + case BOOLEAN: + return _intValue; + case INT: + return _intValue; + case LONG: + case TIMESTAMP: + return _longValue; + case FLOAT: + return _floatValue; + case DOUBLE: + return _doubleValue; + case BIG_DECIMAL: + return _bigDecimalValue; + case STRING: + case JSON: + return _stringValue; + case BYTES: + case UUID: + case VARIANT: + return new ByteArray(_bytesValue); + default: + throw new IllegalStateException("Unhandled Variant target type: " + resultType); + } + } + } + + /** + * Parses a target type literal once for reuse by a transform function. + */ + public static ResultType parseResultType(String targetType) { + if (targetType == null) { + throw new IllegalArgumentException("Variant target type must not be null"); + } + try { + return ResultType.valueOf(targetType.trim().toUpperCase(Locale.ROOT)); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Unsupported Variant target type: " + targetType, e); + } + } + + /** + * Compiles a v1 Variant path. + */ + public static VariantPath compilePath(String path) { + if (path == null || path.isEmpty() || path.charAt(0) != '$') { + throw new IllegalArgumentException("Variant path must start with '$': " + path); + } + List elements = new ArrayList<>(); + int index = 1; + while (index < path.length()) { + char current = path.charAt(index); + if (current == '.') { + int fieldStart = ++index; + while (index < path.length()) { + char next = path.charAt(index); + if (next == '.' || next == '[') { + break; + } + index++; + } + if (fieldStart == index) { + throw new IllegalArgumentException("Variant path contains an empty field: " + path); + } + elements.add(PathElement.forField(path.substring(fieldStart, index))); + } else if (current == '[') { + int subscriptStart = ++index; + while (index < path.length() && Character.isDigit(path.charAt(index))) { + index++; + } + if (subscriptStart == index || index >= path.length() || path.charAt(index) != ']') { + throw new IllegalArgumentException("Invalid Variant array subscript in path: " + path); + } + try { + elements.add(PathElement.forIndex(Integer.parseInt(path.substring(subscriptStart, index)))); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Variant array subscript is too large in path: " + path, e); + } + index++; + } else { + throw new IllegalArgumentException("Unexpected character at offset " + index + " in Variant path: " + path); + } + } + return new VariantPath(elements.toArray(new PathElement[0])); + } + + /** + * Extracts a Variant value. A missing path or SQL null returns Java null; a Variant null remains an encoded Variant + * value. + */ + @Nullable + public static byte[] variantGet(@Nullable byte[] envelope, String path) { + return (byte[]) variantGet(envelope, compilePath(path), ResultType.VARIANT); + } + + /** + * Strictly extracts and converts a value. A missing path or SQL null returns Java null. A Variant null remains + * encoded when the target type is {@link ResultType#VARIANT}, and returns Java null for other target types. An + * incompatible non-null value throws. + */ + @Nullable + public static Object variantGet(@Nullable byte[] envelope, String path, String targetType) { + return variantGet(envelope, compilePath(path), parseResultType(targetType)); + } + + /** + * Strictly extracts using pre-parsed path and type values. + */ + @Nullable + public static Object variantGet(@Nullable byte[] envelope, VariantPath path, ResultType targetType) { + ReusableResult result = new ReusableResult(); + return extractInto(envelope, path, targetType, result) ? result.getExternalValue(targetType) : null; + } + + /** + * Strictly extracts into a reusable, unboxed result. + * + * @return {@code false} for SQL null, a missing path, or Variant null converted to a non-Variant target + */ + public static boolean extractInto(@Nullable byte[] envelope, VariantPath path, ResultType targetType, + ReusableResult result) { + Objects.requireNonNull(result, "result must not be null"); + if (isSqlNull(envelope)) { + return false; + } + Objects.requireNonNull(path, "path must not be null"); + Objects.requireNonNull(targetType, "targetType must not be null"); + Cursor cursor = result._cursor; + if (!cursor.navigate(envelope, path)) { + return false; + } + if (cursor.getType() == Variant.Type.NULL && targetType != ResultType.VARIANT) { + return false; + } + convert(cursor, targetType, result); + return true; + } + + /** + * Tolerant Variant extraction. Malformed input returns Java null. + */ + @Nullable + public static byte[] tryVariantGet(@Nullable byte[] envelope, String path) { + return (byte[]) tryVariantGet(envelope, compilePath(path), ResultType.VARIANT); + } + + /** + * Tolerant typed extraction. Malformed input and incompatible types return Java null. + */ + @Nullable + public static Object tryVariantGet(@Nullable byte[] envelope, String path, String targetType) { + try { + return tryVariantGet(envelope, compilePath(path), parseResultType(targetType)); + } catch (RuntimeException e) { + return null; + } + } + + /** + * Tolerant extraction using pre-parsed path and type values. + */ + @Nullable + public static Object tryVariantGet(@Nullable byte[] envelope, VariantPath path, ResultType targetType) { + try { + ReusableResult result = new ReusableResult(); + return tryExtractInto(envelope, path, targetType, result) ? result.getExternalValue(targetType) : null; + } catch (RuntimeException e) { + return null; + } + } + + /** + * Tolerantly extracts into a reusable, unboxed result. + * + * @return {@code false} for SQL null, missing paths, Variant null converted to a non-Variant target, malformed input, + * or an incompatible conversion + */ + public static boolean tryExtractInto(@Nullable byte[] envelope, VariantPath path, ResultType targetType, + ReusableResult result) { + Objects.requireNonNull(result, "result must not be null"); + if (isSqlNull(envelope)) { + return false; + } + Objects.requireNonNull(path, "path must not be null"); + Objects.requireNonNull(targetType, "targetType must not be null"); + Cursor cursor = result._cursor; + try { + if (!cursor.navigate(envelope, path)) { + return false; + } + if (cursor.getType() == Variant.Type.NULL && targetType != ResultType.VARIANT) { + return false; + } + return tryConvert(cursor, targetType, result); + } catch (IllegalArgumentException | IllegalStateException | UnsupportedOperationException + | IndexOutOfBoundsException e) { + // Cursor operations use these exceptions only for malformed or unsupported Variant encodings. + return false; + } + } + + /** + * Returns whether the path is present. A present Variant null counts as present. + */ + @Nullable + public static Boolean variantExists(@Nullable byte[] envelope, String path) { + return variantExists(envelope, compilePath(path)); + } + + /** + * Returns whether a compiled path is present. A present Variant null counts as present. + */ + @Nullable + public static Boolean variantExists(@Nullable byte[] envelope, VariantPath path) { + return variantExists(envelope, path, new ReusableResult()); + } + + /** + * Allocation-free compiled-path form of {@link #variantExists(byte[], VariantPath)} when the caller retains the + * supplied result between rows. + */ + @Nullable + public static Boolean variantExists(@Nullable byte[] envelope, VariantPath path, ReusableResult result) { + Objects.requireNonNull(result, "result must not be null"); + if (isSqlNull(envelope)) { + return null; + } + return result._cursor.navigate(envelope, Objects.requireNonNull(path, "path must not be null")); + } + + /** + * Returns whether the root value is a Variant null. SQL null is not a Variant null. + */ + public static boolean isVariantNull(@Nullable byte[] envelope) { + return isVariantNull(envelope, ROOT_PATH, new ReusableResult()); + } + + /** + * Returns whether a present value at the path is a Variant null. SQL null and missing paths return false. + */ + public static boolean isVariantNull(@Nullable byte[] envelope, String path) { + return isVariantNull(envelope, compilePath(path)); + } + + /** + * Returns whether a present value at a compiled path is a Variant null. SQL null and missing paths return false. + */ + public static boolean isVariantNull(@Nullable byte[] envelope, VariantPath path) { + return isVariantNull(envelope, path, new ReusableResult()); + } + + /** + * Allocation-free compiled-path form of {@link #isVariantNull(byte[], VariantPath)} when the caller retains the + * supplied result between rows. + */ + public static boolean isVariantNull(@Nullable byte[] envelope, VariantPath path, ReusableResult result) { + Objects.requireNonNull(result, "result must not be null"); + if (isSqlNull(envelope)) { + return false; + } + Cursor cursor = result._cursor; + return cursor.navigate(envelope, Objects.requireNonNull(path, "path must not be null")) + && cursor.getType() == Variant.Type.NULL; + } + + /** + * Returns the Variant type name at the root, or Java null for SQL null. + */ + @Nullable + public static String variantTypeOf(@Nullable byte[] envelope) { + return variantTypeOf(envelope, ROOT_PATH, new ReusableResult()); + } + + /** + * Returns the Variant type name at a path, or Java null for SQL null or a missing path. + */ + @Nullable + public static String variantTypeOf(@Nullable byte[] envelope, String path) { + return variantTypeOf(envelope, compilePath(path)); + } + + /** + * Returns the Variant type name at a compiled path, or Java null for SQL null or a missing path. + */ + @Nullable + public static String variantTypeOf(@Nullable byte[] envelope, VariantPath path) { + return variantTypeOf(envelope, path, new ReusableResult()); + } + + /** + * Allocation-free compiled-path form of {@link #variantTypeOf(byte[], VariantPath)} when the caller retains the + * supplied result between rows. + */ + @Nullable + public static String variantTypeOf(@Nullable byte[] envelope, VariantPath path, ReusableResult result) { + Objects.requireNonNull(result, "result must not be null"); + if (isSqlNull(envelope)) { + return null; + } + Cursor cursor = result._cursor; + return cursor.navigate(envelope, Objects.requireNonNull(path, "path must not be null")) + ? typeName(cursor.getType()) : null; + } + + /** + * Renders the Variant value as canonical JSON text without constructing a JSON tree. + */ + @Nullable + public static String variantToJson(@Nullable byte[] envelope) { + if (isSqlNull(envelope)) { + return null; + } + ReusableResult result = new ReusableResult(); + Cursor cursor = result._cursor; + cursor.navigate(envelope, ROOT_PATH); + return variantToJson(cursor.asVariant()); + } + + /** + * Parses JSON text into a Pinot Variant envelope without constructing a JSON tree. + */ + @Nullable + public static byte[] parseJsonToVariant(@Nullable String json) { + if (json == null) { + return null; + } + try (JsonParser parser = JSON_FACTORY.createParser(json)) { + JsonToken token = parser.nextToken(); + if (token == null) { + throw new IllegalArgumentException("Cannot parse empty text as Variant"); + } + VariantBuilder builder = new VariantBuilder(); + appendJsonValue(parser, token, builder, 0); + if (parser.nextToken() != null) { + throw new IllegalArgumentException("Unexpected trailing token after Variant JSON value"); + } + Variant variant = builder.build(); + return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); + } catch (IOException | RuntimeException e) { + throw new IllegalArgumentException("Cannot parse JSON as Variant", e); + } + } + + /** + * Tolerant JSON parser. Malformed or unsupported input returns Java null. + */ + @Nullable + public static byte[] tryParseJsonToVariant(@Nullable String json) { + try { + return parseJsonToVariant(json); + } catch (RuntimeException e) { + return null; + } + } + + private static boolean isSqlNull(@Nullable byte[] envelope) { + return envelope == null || envelope.length == 0; + } + + private static void convert(Cursor value, ResultType targetType, ReusableResult result) { + switch (targetType) { + case BOOLEAN: + requireType(value, Variant.Type.BOOLEAN, targetType); + result._intValue = value.getBoolean() ? 1 : 0; + break; + case INT: + result._intValue = toInt(value, targetType); + break; + case LONG: + result._longValue = toLong(value, targetType); + break; + case FLOAT: + result._floatValue = toFloat(value, targetType); + break; + case DOUBLE: + result._doubleValue = toDouble(value, targetType); + break; + case BIG_DECIMAL: + result._bigDecimalValue = toBigDecimal(value, targetType); + break; + case STRING: + requireType(value, Variant.Type.STRING, targetType); + result._stringValue = value.getString(); + break; + case BYTES: + requireType(value, Variant.Type.BINARY, targetType); + result._bytesValue = value.getBinary(); + break; + case UUID: + requireType(value, Variant.Type.UUID, targetType); + result._bytesValue = value.getUuidBytes(); + break; + case TIMESTAMP: + result._longValue = toTimestampMillis(value, targetType); + break; + case VARIANT: + result._bytesValue = value.copyEnvelope(); + break; + case JSON: + result._stringValue = variantToJson(value.asVariant()); + break; + default: + throw new IllegalStateException("Unhandled Variant target type: " + targetType); + } + } + + private static boolean tryConvert(Cursor value, ResultType targetType, ReusableResult result) { + Variant.Type valueType = value.getType(); + switch (targetType) { + case BOOLEAN: + if (valueType != Variant.Type.BOOLEAN) { + return false; + } + result._intValue = value.getBoolean() ? 1 : 0; + return true; + case INT: + return tryConvertToInt(value, valueType, result); + case LONG: + return tryConvertToLong(value, valueType, result); + case FLOAT: + switch (valueType) { + case BYTE: + case SHORT: + case INT: + case LONG: + result._floatValue = value.getInteger(); + return true; + case FLOAT: + result._floatValue = value.getFloat(); + return true; + case DOUBLE: + result._floatValue = (float) value.getDouble(); + return true; + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + result._floatValue = value.getDecimal().floatValue(); + return true; + default: + return false; + } + case DOUBLE: + switch (valueType) { + case BYTE: + case SHORT: + case INT: + case LONG: + result._doubleValue = value.getInteger(); + return true; + case FLOAT: + result._doubleValue = value.getFloat(); + return true; + case DOUBLE: + result._doubleValue = value.getDouble(); + return true; + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + result._doubleValue = value.getDecimal().doubleValue(); + return true; + default: + return false; + } + case BIG_DECIMAL: + switch (valueType) { + case BYTE: + case SHORT: + case INT: + case LONG: + result._bigDecimalValue = BigDecimal.valueOf(value.getInteger()); + return true; + case FLOAT: + float floatValue = value.getFloat(); + if (!Float.isFinite(floatValue)) { + return false; + } + result._bigDecimalValue = BigDecimal.valueOf(floatValue); + return true; + case DOUBLE: + double doubleValue = value.getDouble(); + if (!Double.isFinite(doubleValue)) { + return false; + } + result._bigDecimalValue = BigDecimal.valueOf(doubleValue); + return true; + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + result._bigDecimalValue = value.getDecimal(); + return true; + default: + return false; + } + case STRING: + if (valueType != Variant.Type.STRING) { + return false; + } + result._stringValue = value.getString(); + return true; + case BYTES: + if (valueType != Variant.Type.BINARY) { + return false; + } + result._bytesValue = value.getBinary(); + return true; + case UUID: + if (valueType != Variant.Type.UUID) { + return false; + } + result._bytesValue = value.getUuidBytes(); + return true; + case TIMESTAMP: + switch (valueType) { + case DATE: + result._longValue = value.getInteger() * TimeUnit.DAYS.toMillis(1); + return true; + case TIMESTAMP_TZ: + case TIMESTAMP_NTZ: + result._longValue = Math.floorDiv(value.getInteger(), TimeUnit.MILLISECONDS.toMicros(1)); + return true; + case TIMESTAMP_NANOS_TZ: + case TIMESTAMP_NANOS_NTZ: + result._longValue = Math.floorDiv(value.getInteger(), TimeUnit.MILLISECONDS.toNanos(1)); + return true; + default: + return false; + } + case VARIANT: + result._bytesValue = value.copyEnvelope(); + return true; + case JSON: + result._stringValue = variantToJson(value.asVariant()); + return true; + default: + throw new AssertionError("Unhandled Variant target type: " + targetType); + } + } + + private static boolean tryConvertToInt(Cursor value, Variant.Type valueType, ReusableResult result) { + switch (valueType) { + case BYTE: + case SHORT: + case INT: + result._intValue = (int) value.getInteger(); + return true; + case LONG: + long longValue = value.getInteger(); + if (longValue < Integer.MIN_VALUE || longValue > Integer.MAX_VALUE) { + return false; + } + result._intValue = (int) longValue; + return true; + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + BigDecimal decimalValue = value.getDecimal(); + if (!isIntegralInRange(decimalValue, MIN_INT_DECIMAL, MAX_INT_DECIMAL)) { + return false; + } + result._intValue = decimalValue.intValue(); + return true; + default: + return false; + } + } + + private static boolean tryConvertToLong(Cursor value, Variant.Type valueType, ReusableResult result) { + switch (valueType) { + case BYTE: + case SHORT: + case INT: + case LONG: + result._longValue = value.getInteger(); + return true; + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + BigDecimal decimalValue = value.getDecimal(); + if (!isIntegralInRange(decimalValue, MIN_LONG_DECIMAL, MAX_LONG_DECIMAL)) { + return false; + } + result._longValue = decimalValue.longValue(); + return true; + default: + return false; + } + } + + private static boolean isIntegralInRange(BigDecimal value, BigDecimal minimum, BigDecimal maximum) { + return value.compareTo(minimum) >= 0 && value.compareTo(maximum) <= 0 + && (value.scale() <= 0 || value.stripTrailingZeros().scale() <= 0); + } + + private static int toInt(Cursor value, ResultType targetType) { + switch (value.getType()) { + case BYTE: + case SHORT: + case INT: + return (int) value.getInteger(); + case LONG: + return Math.toIntExact(value.getInteger()); + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + return value.getDecimal().intValueExact(); + default: + throw typeMismatch(value, targetType); + } + } + + private static long toLong(Cursor value, ResultType targetType) { + switch (value.getType()) { + case BYTE: + case SHORT: + case INT: + case LONG: + return value.getInteger(); + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + return value.getDecimal().longValueExact(); + default: + throw typeMismatch(value, targetType); + } + } + + private static float toFloat(Cursor value, ResultType targetType) { + switch (value.getType()) { + case BYTE: + case SHORT: + case INT: + case LONG: + return value.getInteger(); + case FLOAT: + return value.getFloat(); + case DOUBLE: + return (float) value.getDouble(); + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + return value.getDecimal().floatValue(); + default: + throw typeMismatch(value, targetType); + } + } + + private static double toDouble(Cursor value, ResultType targetType) { + switch (value.getType()) { + case BYTE: + case SHORT: + case INT: + case LONG: + return value.getInteger(); + case FLOAT: + return value.getFloat(); + case DOUBLE: + return value.getDouble(); + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + return value.getDecimal().doubleValue(); + default: + throw typeMismatch(value, targetType); + } + } + + private static BigDecimal toBigDecimal(Cursor value, ResultType targetType) { + switch (value.getType()) { + case BYTE: + case SHORT: + case INT: + case LONG: + return BigDecimal.valueOf(value.getInteger()); + case FLOAT: + return BigDecimal.valueOf(value.getFloat()); + case DOUBLE: + return BigDecimal.valueOf(value.getDouble()); + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + return value.getDecimal(); + default: + throw typeMismatch(value, targetType); + } + } + + private static long toTimestampMillis(Cursor value, ResultType targetType) { + switch (value.getType()) { + case DATE: + return Math.multiplyExact(value.getInteger(), TimeUnit.DAYS.toMillis(1)); + case TIMESTAMP_TZ: + case TIMESTAMP_NTZ: + return Math.floorDiv(value.getInteger(), TimeUnit.MILLISECONDS.toMicros(1)); + case TIMESTAMP_NANOS_TZ: + case TIMESTAMP_NANOS_NTZ: + return Math.floorDiv(value.getInteger(), TimeUnit.MILLISECONDS.toNanos(1)); + default: + throw typeMismatch(value, targetType); + } + } + + private static void requireType(Cursor value, Variant.Type expected, ResultType targetType) { + if (value.getType() != expected) { + throw typeMismatch(value, targetType); + } + } + + private static IllegalArgumentException typeMismatch(Cursor value, ResultType targetType) { + return new IllegalArgumentException( + "Cannot convert Variant " + typeName(value.getType()) + " to " + targetType.name()); + } + + private static String typeName(Variant.Type type) { + switch (type) { + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + return "DECIMAL"; + default: + return type.name(); + } + } + + private static String variantToJson(Variant variant) { + try { + StringWriter writer = new StringWriter(); + try (JsonGenerator generator = JSON_FACTORY.createGenerator(writer)) { + writeJsonValue(generator, variant); + } + return writer.toString(); + } catch (IOException e) { + throw new IllegalStateException("Cannot render Variant as JSON", e); + } + } + + private static void writeJsonValue(JsonGenerator generator, Variant variant) + throws IOException { + switch (variant.getType()) { + case OBJECT: + generator.writeStartObject(); + for (int i = 0; i < variant.numObjectElements(); i++) { + Variant.ObjectField field = variant.getFieldAtIndex(i); + generator.writeFieldName(field.key); + writeJsonValue(generator, field.value); + } + generator.writeEndObject(); + break; + case ARRAY: + generator.writeStartArray(); + for (int i = 0; i < variant.numArrayElements(); i++) { + writeJsonValue(generator, variant.getElementAtIndex(i)); + } + generator.writeEndArray(); + break; + case NULL: + generator.writeNull(); + break; + case BOOLEAN: + generator.writeBoolean(variant.getBoolean()); + break; + case BYTE: + generator.writeNumber(variant.getByte()); + break; + case SHORT: + generator.writeNumber(variant.getShort()); + break; + case INT: + generator.writeNumber(variant.getInt()); + break; + case LONG: + generator.writeNumber(variant.getLong()); + break; + case FLOAT: + generator.writeNumber(variant.getFloat()); + break; + case DOUBLE: + generator.writeNumber(variant.getDouble()); + break; + case DECIMAL4: + case DECIMAL8: + case DECIMAL16: + generator.writeNumber(variant.getDecimal()); + break; + case STRING: + generator.writeString(variant.getString()); + break; + case BINARY: + generator.writeBinary(toBytes(variant.getBinary())); + break; + case UUID: + generator.writeString(variant.getUUID().toString()); + break; + case DATE: + generator.writeString(LocalDate.ofEpochDay(variant.getInt()).toString()); + break; + case TIMESTAMP_TZ: + generator.writeString(instantFromMicros(variant.getLong()).toString()); + break; + case TIMESTAMP_NTZ: + generator.writeString(LocalDateTime.ofInstant(instantFromMicros(variant.getLong()), ZoneOffset.UTC).toString()); + break; + case TIMESTAMP_NANOS_TZ: + generator.writeString(instantFromNanos(variant.getLong()).toString()); + break; + case TIMESTAMP_NANOS_NTZ: + generator.writeString(LocalDateTime.ofInstant(instantFromNanos(variant.getLong()), ZoneOffset.UTC).toString()); + break; + case TIME: + generator.writeString(LocalTime.ofNanoOfDay(Math.floorMod(variant.getLong() * NANOS_PER_MICRO, NANOS_PER_DAY)) + .toString()); + break; + default: + throw new IllegalStateException("Unsupported Variant type: " + variant.getType()); + } + } + + private static Instant instantFromMicros(long micros) { + long seconds = Math.floorDiv(micros, MICROS_PER_SECOND); + long nanos = Math.floorMod(micros, MICROS_PER_SECOND) * NANOS_PER_MICRO; + return Instant.ofEpochSecond(seconds, nanos); + } + + private static Instant instantFromNanos(long nanos) { + return Instant.ofEpochSecond(Math.floorDiv(nanos, TimeUnit.SECONDS.toNanos(1)), + Math.floorMod(nanos, TimeUnit.SECONDS.toNanos(1))); + } + + private static byte[] toBytes(ByteBuffer buffer) { + ByteBuffer view = buffer.slice(); + byte[] bytes = new byte[view.remaining()]; + view.get(bytes); + return bytes; + } + + private static void appendJsonValue(JsonParser parser, JsonToken token, VariantBuilder builder, int depth) + throws IOException { + if (depth > MAX_JSON_NESTING_DEPTH) { + throw new IllegalArgumentException("Variant JSON exceeds maximum nesting depth " + MAX_JSON_NESTING_DEPTH); + } + switch (token) { + case START_OBJECT: + VariantObjectBuilder objectBuilder = builder.startObject(); + while (parser.nextToken() != JsonToken.END_OBJECT) { + if (parser.currentToken() != JsonToken.FIELD_NAME) { + throw new IllegalArgumentException("Expected a JSON object field name"); + } + objectBuilder.appendKey(parser.currentName()); + JsonToken fieldValue = parser.nextToken(); + if (fieldValue == null) { + throw new IllegalArgumentException("Unexpected end of JSON object"); + } + appendJsonValue(parser, fieldValue, objectBuilder, depth + 1); + } + builder.endObject(); + break; + case START_ARRAY: + VariantArrayBuilder arrayBuilder = builder.startArray(); + while (true) { + JsonToken element = parser.nextToken(); + if (element == JsonToken.END_ARRAY) { + break; + } + if (element == null) { + throw new IllegalArgumentException("Unexpected end of JSON array"); + } + appendJsonValue(parser, element, arrayBuilder, depth + 1); + } + builder.endArray(); + break; + case VALUE_NULL: + builder.appendNull(); + break; + case VALUE_TRUE: + builder.appendBoolean(true); + break; + case VALUE_FALSE: + builder.appendBoolean(false); + break; + case VALUE_STRING: + builder.appendString(parser.getText()); + break; + case VALUE_NUMBER_INT: + appendInteger(parser, builder); + break; + case VALUE_NUMBER_FLOAT: + appendDecimal(parser.getDecimalValue(), builder); + break; + default: + throw new IllegalArgumentException("Unsupported JSON token for Variant: " + token); + } + } + + private static void appendInteger(JsonParser parser, VariantBuilder builder) + throws IOException { + switch (parser.getNumberType()) { + case INT: + builder.appendInt(parser.getIntValue()); + break; + case LONG: + builder.appendLong(parser.getLongValue()); + break; + case BIG_INTEGER: + appendBigInteger(parser.getBigIntegerValue(), builder); + break; + default: + throw new IllegalArgumentException("Unsupported JSON integer representation: " + parser.getNumberType()); + } + } + + private static void appendBigInteger(BigInteger value, VariantBuilder builder) { + if (value.bitLength() < Integer.SIZE) { + builder.appendInt(value.intValue()); + } else if (value.bitLength() < Long.SIZE) { + builder.appendLong(value.longValue()); + } else { + appendDecimal(new BigDecimal(value), builder); + } + } + + private static void appendDecimal(BigDecimal value, VariantBuilder builder) { + BigDecimal normalized = value; + if (normalized.scale() < 0) { + // Parquet Variant stores scale as an unsigned byte. Expand exponent notation exactly instead of allowing a + // negative scale to wrap during encoding. + long expandedPrecision = (long) normalized.precision() - normalized.scale(); + if (normalized.signum() != 0 && expandedPrecision > MAX_VARIANT_DECIMAL_PRECISION) { + throw unsupportedVariantDecimal(value); + } + normalized = normalized.signum() == 0 ? BigDecimal.ZERO : normalized.setScale(0); + } else if (normalized.scale() > MAX_VARIANT_DECIMAL_SCALE) { + // Accept values whose excessive lexical scale consists only of insignificant trailing zeros. + normalized = normalized.stripTrailingZeros(); + if (normalized.scale() < 0) { + normalized = normalized.setScale(0); + } + } + byte[] unscaledBytes = normalized.unscaledValue().toByteArray(); + if (normalized.scale() > MAX_VARIANT_DECIMAL_SCALE + || normalized.precision() > MAX_VARIANT_DECIMAL_PRECISION + || unscaledBytes.length > MAX_VARIANT_DECIMAL_BYTES) { + throw unsupportedVariantDecimal(value); + } + builder.appendDecimal(normalized); + } + + private static IllegalArgumentException unsupportedVariantDecimal(BigDecimal value) { + return new IllegalArgumentException( + "JSON decimal exceeds Parquet Variant decimal(38) encoding: precision=" + value.precision() + + ", scale=" + value.scale()); + } + + /** + * Mutable zero-copy view over one selected value in a Pinot envelope. + * + *

The constants and layouts used here mirror Parquet Variant encoding version 1. Keeping this cursor on + * {@link ReusableResult} avoids allocating envelope views, Variant wrappers, and navigation wrappers for every row. + */ + private static final class Cursor { + private byte[] _envelope; + private int _metadataOffset; + private int _metadataLength; + private boolean _metadataParsed; + private int _metadataOffsetSize; + private int _metadataDictSize; + private int _metadataOffsetListOffset; + private int _metadataDataOffset; + private int _metadataDataLength; + private int _selectedOffset; + private int _selectedLength; + + private boolean navigate(byte[] envelope, VariantPath path) { + reset(envelope); + for (PathElement element : path._elements) { + if (element._field != null) { + if (getType() != Variant.Type.OBJECT || !selectObjectField(element._fieldUtf8)) { + return false; + } + } else if (getType() != Variant.Type.ARRAY || !selectArrayElement(element._index)) { + return false; + } + } + return true; + } + + private void reset(byte[] envelope) { + int metadataLength = VariantEnvelope.validateAndGetMetadataLength(envelope); + int valueLength = envelope.length - VariantEnvelope.HEADER_SIZE - metadataLength; + + _envelope = envelope; + _metadataOffset = VariantEnvelope.HEADER_SIZE; + _metadataLength = metadataLength; + requireRange(_metadataOffset, 1, _metadataOffset + _metadataLength, "Variant metadata"); + int metadataVersion = _envelope[_metadataOffset] & VARIANT_METADATA_VERSION_MASK; + if (metadataVersion != VARIANT_METADATA_VERSION) { + throw new UnsupportedOperationException("Unsupported variant metadata version: " + metadataVersion); + } + _metadataParsed = false; + _selectedOffset = VariantEnvelope.HEADER_SIZE + metadataLength; + _selectedLength = valueLength; + } + + private Variant.Type getType() { + int header = getHeader(); + int basicType = header & VARIANT_BASIC_TYPE_MASK; + int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; + switch (basicType) { + case VARIANT_SHORT_STRING: + return Variant.Type.STRING; + case VARIANT_OBJECT: + return Variant.Type.OBJECT; + case VARIANT_ARRAY: + return Variant.Type.ARRAY; + case VARIANT_PRIMITIVE: + switch (typeInfo) { + case VARIANT_NULL: + return Variant.Type.NULL; + case VARIANT_TRUE: + case VARIANT_FALSE: + return Variant.Type.BOOLEAN; + case VARIANT_INT8: + return Variant.Type.BYTE; + case VARIANT_INT16: + return Variant.Type.SHORT; + case VARIANT_INT32: + return Variant.Type.INT; + case VARIANT_INT64: + return Variant.Type.LONG; + case VARIANT_DOUBLE: + return Variant.Type.DOUBLE; + case VARIANT_DECIMAL4: + return Variant.Type.DECIMAL4; + case VARIANT_DECIMAL8: + return Variant.Type.DECIMAL8; + case VARIANT_DECIMAL16: + return Variant.Type.DECIMAL16; + case VARIANT_DATE: + return Variant.Type.DATE; + case VARIANT_TIMESTAMP_TZ: + return Variant.Type.TIMESTAMP_TZ; + case VARIANT_TIMESTAMP_NTZ: + return Variant.Type.TIMESTAMP_NTZ; + case VARIANT_FLOAT: + return Variant.Type.FLOAT; + case VARIANT_BINARY: + return Variant.Type.BINARY; + case VARIANT_LONG_STRING: + return Variant.Type.STRING; + case VARIANT_TIME: + return Variant.Type.TIME; + case VARIANT_TIMESTAMP_NANOS_TZ: + return Variant.Type.TIMESTAMP_NANOS_TZ; + case VARIANT_TIMESTAMP_NANOS_NTZ: + return Variant.Type.TIMESTAMP_NANOS_NTZ; + case VARIANT_UUID: + return Variant.Type.UUID; + default: + throw new UnsupportedOperationException("Unknown type in Variant. primitive type: " + typeInfo); + } + default: + throw new IllegalStateException("Unhandled Variant basic type: " + basicType); + } + } + + private boolean getBoolean() { + int typeInfo = getPrimitiveTypeInfo(); + if (typeInfo != VARIANT_TRUE && typeInfo != VARIANT_FALSE) { + throw new IllegalArgumentException("Cannot read non-boolean Variant value as BOOLEAN"); + } + return typeInfo == VARIANT_TRUE; + } + + private long getInteger() { + int typeInfo = getPrimitiveTypeInfo(); + switch (typeInfo) { + case VARIANT_INT8: + return readSignedLittleEndian(_selectedOffset + 1, 1); + case VARIANT_INT16: + return readSignedLittleEndian(_selectedOffset + 1, 2); + case VARIANT_INT32: + case VARIANT_DATE: + return readSignedLittleEndian(_selectedOffset + 1, Integer.BYTES); + case VARIANT_INT64: + case VARIANT_TIMESTAMP_TZ: + case VARIANT_TIMESTAMP_NTZ: + case VARIANT_TIME: + case VARIANT_TIMESTAMP_NANOS_TZ: + case VARIANT_TIMESTAMP_NANOS_NTZ: + return readSignedLittleEndian(_selectedOffset + 1, Long.BYTES); + default: + throw new IllegalArgumentException("Cannot read non-integer Variant value as an integer"); + } + } + + private float getFloat() { + if (getPrimitiveTypeInfo() != VARIANT_FLOAT) { + throw new IllegalArgumentException("Cannot read non-float Variant value as FLOAT"); + } + return Float.intBitsToFloat((int) readSignedLittleEndian(_selectedOffset + 1, Float.BYTES)); + } + + private double getDouble() { + if (getPrimitiveTypeInfo() != VARIANT_DOUBLE) { + throw new IllegalArgumentException("Cannot read non-double Variant value as DOUBLE"); + } + return Double.longBitsToDouble(readSignedLittleEndian(_selectedOffset + 1, Double.BYTES)); + } + + private BigDecimal getDecimal() { + int typeInfo = getPrimitiveTypeInfo(); + requireSelectedRange(1, 1); + int scale = Byte.toUnsignedInt(_envelope[_selectedOffset + 1]); + switch (typeInfo) { + case VARIANT_DECIMAL4: + return BigDecimal.valueOf(readSignedLittleEndian(_selectedOffset + 2, Integer.BYTES), scale); + case VARIANT_DECIMAL8: + return BigDecimal.valueOf(readSignedLittleEndian(_selectedOffset + 2, Long.BYTES), scale); + case VARIANT_DECIMAL16: + requireSelectedRange(2, 16); + byte[] unscaled = new byte[16]; + for (int i = 0; i < unscaled.length; i++) { + unscaled[i] = _envelope[_selectedOffset + 17 - i]; + } + return new BigDecimal(new BigInteger(unscaled), scale); + default: + throw new IllegalArgumentException("Cannot read non-decimal Variant value as DECIMAL"); + } + } + + private String getString() { + int header = getHeader(); + int basicType = header & VARIANT_BASIC_TYPE_MASK; + int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; + int contentOffset; + int length; + if (basicType == VARIANT_SHORT_STRING) { + contentOffset = 1; + length = typeInfo; + } else if (basicType == VARIANT_PRIMITIVE && typeInfo == VARIANT_LONG_STRING) { + contentOffset = 1 + Integer.BYTES; + length = readUnsignedLittleEndian(_selectedOffset + 1, Integer.BYTES, selectedLimit()); + } else { + throw new IllegalArgumentException("Cannot read non-string Variant value as STRING"); + } + requireSelectedRange(contentOffset, length); + return new String(_envelope, _selectedOffset + contentOffset, length, StandardCharsets.UTF_8); + } + + private byte[] getBinary() { + if (getPrimitiveTypeInfo() != VARIANT_BINARY) { + throw new IllegalArgumentException("Cannot read non-binary Variant value as BINARY"); + } + int length = readUnsignedLittleEndian(_selectedOffset + 1, Integer.BYTES, selectedLimit()); + int contentOffset = 1 + Integer.BYTES; + requireSelectedRange(contentOffset, length); + byte[] bytes = new byte[length]; + System.arraycopy(_envelope, _selectedOffset + contentOffset, bytes, 0, length); + return bytes; + } + + private byte[] getUuidBytes() { + if (getPrimitiveTypeInfo() != VARIANT_UUID) { + throw new IllegalArgumentException("Cannot read non-UUID Variant value as UUID"); + } + requireSelectedRange(1, UuidUtils.UUID_NUM_BYTES); + byte[] bytes = new byte[UuidUtils.UUID_NUM_BYTES]; + System.arraycopy(_envelope, _selectedOffset + 1, bytes, 0, UuidUtils.UUID_NUM_BYTES); + return bytes; + } + + private byte[] copyEnvelope() { + return VariantEnvelope.encode(_envelope, _metadataOffset, _metadataLength, _envelope, _selectedOffset, + _selectedLength); + } + + private Variant asVariant() { + return new Variant(_envelope, _selectedOffset, _selectedLength, _envelope, _metadataOffset, _metadataLength); + } + + private boolean selectObjectField(byte[] fieldUtf8) { + int header = getHeader(); + int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; + int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(_selectedOffset + 1, sizeBytes, selectedLimit()); + int idSize = ((typeInfo >>> 2) & 3) + 1; + int offsetSize = (typeInfo & 3) + 1; + int idStart = checkedPosition((long) _selectedOffset + 1 + sizeBytes, selectedLimit(), "Variant object ids"); + int offsetStart = checkedPosition((long) idStart + (long) numElements * idSize, selectedLimit(), + "Variant object offsets"); + int dataStart = checkedPosition((long) offsetStart + ((long) numElements + 1) * offsetSize, selectedLimit(), + "Variant object data"); + int finalOffsetPosition = checkedPosition((long) offsetStart + (long) numElements * offsetSize, + selectedLimit(), "Variant object final offset"); + int totalDataLength = readUnsignedLittleEndian(finalOffsetPosition, offsetSize, selectedLimit()); + requireRange(dataStart, totalDataLength, selectedLimit(), "Variant object data"); + + for (int i = 0; i < numElements; i++) { + int id = readUnsignedLittleEndian(idStart + i * idSize, idSize, offsetStart); + if (metadataKeyEquals(id, fieldUtf8)) { + int offset = readUnsignedLittleEndian(offsetStart + i * offsetSize, offsetSize, dataStart); + int nextOffset = readUnsignedLittleEndian(offsetStart + (i + 1) * offsetSize, offsetSize, dataStart); + if (offset > nextOffset || nextOffset > totalDataLength) { + throw new IllegalStateException( + "Invalid Variant object offsets: " + offset + ", " + nextOffset + ", total=" + totalDataLength); + } + _selectedOffset = dataStart + offset; + _selectedLength = nextOffset - offset; + return true; + } + } + return false; + } + + private boolean selectArrayElement(int index) { + int header = getHeader(); + int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; + int sizeBytes = ((typeInfo >>> 2) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(_selectedOffset + 1, sizeBytes, selectedLimit()); + int offsetSize = (typeInfo & 3) + 1; + int offsetStart = checkedPosition((long) _selectedOffset + 1 + sizeBytes, selectedLimit(), + "Variant array offsets"); + int dataStart = checkedPosition((long) offsetStart + ((long) numElements + 1) * offsetSize, selectedLimit(), + "Variant array data"); + int finalOffsetPosition = checkedPosition((long) offsetStart + (long) numElements * offsetSize, + selectedLimit(), "Variant array final offset"); + int totalDataLength = readUnsignedLittleEndian(finalOffsetPosition, offsetSize, selectedLimit()); + requireRange(dataStart, totalDataLength, selectedLimit(), "Variant array data"); + if (index >= numElements) { + return false; + } + int offset = readUnsignedLittleEndian(offsetStart + index * offsetSize, offsetSize, dataStart); + int nextOffset = readUnsignedLittleEndian(offsetStart + (index + 1) * offsetSize, offsetSize, dataStart); + if (offset > nextOffset || nextOffset > totalDataLength) { + throw new IllegalStateException( + "Invalid Variant array offsets: " + offset + ", " + nextOffset + ", total=" + totalDataLength); + } + _selectedOffset = dataStart + offset; + _selectedLength = nextOffset - offset; + return true; + } + + private boolean metadataKeyEquals(int id, byte[] expected) { + ensureMetadataParsed(); + if (id < 0 || id >= _metadataDictSize) { + throw new IllegalArgumentException( + "Invalid dictionary id: " + id + ". dictionary size: " + _metadataDictSize); + } + int offset = readUnsignedLittleEndian(_metadataOffsetListOffset + id * _metadataOffsetSize, + _metadataOffsetSize, _metadataDataOffset); + int nextOffset = readUnsignedLittleEndian(_metadataOffsetListOffset + (id + 1) * _metadataOffsetSize, + _metadataOffsetSize, _metadataDataOffset); + if (offset > nextOffset || nextOffset > _metadataDataLength) { + throw new IllegalStateException( + "Invalid Variant metadata offsets: " + offset + ", " + nextOffset + ", total=" + _metadataDataLength); + } + int length = nextOffset - offset; + if (length != expected.length) { + return false; + } + int start = _metadataDataOffset + offset; + for (int i = 0; i < length; i++) { + if (_envelope[start + i] != expected[i]) { + return false; + } + } + return true; + } + + private void ensureMetadataParsed() { + if (_metadataParsed) { + return; + } + int metadataLimit = _metadataOffset + _metadataLength; + int header = Byte.toUnsignedInt(_envelope[_metadataOffset]); + int offsetSize = ((header >>> 6) & 3) + 1; + int dictSize = readUnsignedLittleEndian(_metadataOffset + 1, offsetSize, metadataLimit); + int offsetListOffset = checkedPosition((long) _metadataOffset + 1 + offsetSize, metadataLimit, + "Variant metadata offsets"); + int dataOffset = checkedPosition((long) offsetListOffset + ((long) dictSize + 1) * offsetSize, metadataLimit, + "Variant metadata data"); + int finalOffsetPosition = checkedPosition((long) offsetListOffset + (long) dictSize * offsetSize, metadataLimit, + "Variant metadata final offset"); + int dataLength = readUnsignedLittleEndian(finalOffsetPosition, offsetSize, metadataLimit); + requireRange(dataOffset, dataLength, metadataLimit, "Variant metadata data"); + _metadataOffsetSize = offsetSize; + _metadataDictSize = dictSize; + _metadataOffsetListOffset = offsetListOffset; + _metadataDataOffset = dataOffset; + _metadataDataLength = dataLength; + _metadataParsed = true; + } + + private int getHeader() { + requireSelectedRange(0, 1); + return Byte.toUnsignedInt(_envelope[_selectedOffset]); + } + + private int getPrimitiveTypeInfo() { + int header = getHeader(); + if ((header & VARIANT_BASIC_TYPE_MASK) != VARIANT_PRIMITIVE) { + throw new IllegalArgumentException("Variant value is not a primitive"); + } + return (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; + } + + private long readSignedLittleEndian(int offset, int numBytes) { + requireRange(offset, numBytes, selectedLimit(), "Variant value"); + long value = 0; + for (int i = 0; i < numBytes - 1; i++) { + value |= (long) Byte.toUnsignedInt(_envelope[offset + i]) << (Byte.SIZE * i); + } + return value | (long) _envelope[offset + numBytes - 1] << (Byte.SIZE * (numBytes - 1)); + } + + private int readUnsignedLittleEndian(int offset, int numBytes, int limit) { + return VariantUtils.readUnsignedLittleEndian(_envelope, offset, numBytes, limit); + } + + private void requireSelectedRange(int relativeOffset, int length) { + if (relativeOffset < 0 || (long) relativeOffset + length > _selectedLength) { + throw new IllegalArgumentException( + "Invalid Variant value range: offset=" + relativeOffset + ", length=" + length + ", valueLength=" + + _selectedLength); + } + } + + private int selectedLimit() { + return _selectedOffset + _selectedLength; + } + } + + private static int readUnsignedLittleEndian(byte[] bytes, int offset, int numBytes, int limit) { + if (numBytes < 1 || numBytes > Integer.BYTES) { + throw new IllegalArgumentException("Invalid unsigned integer width: " + numBytes); + } + requireRange(offset, numBytes, limit, "Variant unsigned integer"); + long value = 0; + for (int i = 0; i < numBytes; i++) { + value |= (long) Byte.toUnsignedInt(bytes[offset + i]) << (Byte.SIZE * i); + } + if (value > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Variant unsigned integer exceeds supported Java range: " + value); + } + return (int) value; + } + + private static int checkedPosition(long position, int limit, String description) { + if (position < 0 || position > limit) { + throw new IllegalArgumentException(description + " exceeds encoded value bounds"); + } + return (int) position; + } + + private static void requireRange(int offset, int length, int limit, String description) { + if (offset < 0 || length < 0 || (long) offset + length > limit) { + throw new IllegalArgumentException( + description + " exceeds encoded bounds: offset=" + offset + ", length=" + length + ", limit=" + limit); + } + } + + private static final class PathElement { + private final String _field; + private final byte[] _fieldUtf8; + private final int _index; + + private PathElement(String field, byte[] fieldUtf8, int index) { + _field = field; + _fieldUtf8 = fieldUtf8; + _index = index; + } + + private static PathElement forField(String field) { + return new PathElement(field, field.getBytes(StandardCharsets.UTF_8), -1); + } + + private static PathElement forIndex(int index) { + return new PathElement(null, null, index); + } + } +} diff --git a/pinot-common/src/main/proto/expressions.proto b/pinot-common/src/main/proto/expressions.proto index cd185eba0843..b7f3a42f915a 100644 --- a/pinot-common/src/main/proto/expressions.proto +++ b/pinot-common/src/main/proto/expressions.proto @@ -46,6 +46,7 @@ enum ColumnDataType { BIG_DECIMAL_ARRAY = 21; UUID = 22; UUID_ARRAY = 23; + VARIANT = 24; } message InputRef { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java b/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java index d8436568bf2f..978a1a47e0a7 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluatorTest.java @@ -18,16 +18,24 @@ */ package org.apache.pinot.common.evaluator; +import java.sql.Timestamp; import java.util.List; +import java.util.UUID; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantBuilder; import org.apache.pinot.common.function.FunctionUtils; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.utils.PinotDataType; +import org.apache.pinot.spi.utils.VariantEnvelope; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; 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.assertThrows; import static org.testng.Assert.assertTrue; @@ -250,4 +258,112 @@ public void testPolymorphicBitwiseFunctions() { longRow.putValue("shift", 40); assertEquals(new InbuiltFunctionEvaluator("bitExtract(value, shift)").evaluate(longRow), 1); } + + @Test + public void testPlannedVariantScalarFunctions() { + byte[] first = VariantUtils.parseJsonToVariant( + "{\"value\":7,\"name\":\"pinot\",\"nested\":{\"flag\":true},\"nullValue\":null}"); + byte[] second = VariantUtils.parseJsonToVariant( + "{\"value\":9,\"name\":\"apache\",\"nested\":{\"flag\":false},\"nullValue\":null}"); + GenericRow row = new GenericRow(); + + InbuiltFunctionEvaluator getInt = + new InbuiltFunctionEvaluator("variantGet(variant, '$.value', 'INT')"); + InbuiltFunctionEvaluator getNested = new InbuiltFunctionEvaluator("variantGet(variant, '$.nested')"); + InbuiltFunctionEvaluator exists = + new InbuiltFunctionEvaluator("variantExists(variant, '$.nullValue')"); + InbuiltFunctionEvaluator isNull = + new InbuiltFunctionEvaluator("isVariantNull(variant, '$.nullValue')"); + InbuiltFunctionEvaluator typeOf = + new InbuiltFunctionEvaluator("variantTypeOf(variant, '$.nested')"); + + assertEquals(getInt.getArguments(), List.of("variant")); + row.putValue("variant", first); + assertEquals(getInt.evaluate(row), 7); + assertEquals(getInt.evaluate(new Object[]{first}), 7); + assertEquals(VariantUtils.variantToJson((byte[]) getNested.evaluate(row)), "{\"flag\":true}"); + assertEquals(exists.evaluate(row), true); + assertEquals(isNull.evaluate(row), true); + assertEquals(typeOf.evaluate(row), "OBJECT"); + + row.putValue("variant", second); + assertEquals(getInt.evaluate(row), 9); + assertEquals(VariantUtils.variantToJson((byte[]) getNested.evaluate(row)), "{\"flag\":false}"); + assertEquals(exists.evaluate(row), true); + assertEquals(isNull.evaluate(row), true); + assertEquals(typeOf.evaluate(row), "OBJECT"); + } + + @Test + public void testPlannedVariantUsesExternalUuidAndTimestampRepresentations() { + UUID uuid = UUID.fromString("00112233-4455-6677-8899-aabbccddeeff"); + VariantBuilder builder = new VariantBuilder(); + builder.appendUUID(uuid); + GenericRow row = new GenericRow(); + row.putValue("variant", encode(builder)); + assertEquals(new InbuiltFunctionEvaluator("variantGet(variant, '$', 'UUID')").evaluate(row), uuid); + assertEquals(new InbuiltFunctionEvaluator("tryVariantGet(variant, '$', 'UUID')").evaluate(row), uuid); + + builder = new VariantBuilder(); + builder.appendTimestampTz(1_700_000_000_123_000L); + row.putValue("variant", encode(builder)); + Timestamp timestamp = new Timestamp(1_700_000_000_123L); + assertEquals(new InbuiltFunctionEvaluator("variantGet(variant, '$', 'TIMESTAMP')").evaluate(row), timestamp); + assertEquals(new InbuiltFunctionEvaluator("tryVariantGet(variant, '$', 'TIMESTAMP')").evaluate(row), timestamp); + } + + @Test + public void testPlannedVariantStrictTryAndSqlNullSemantics() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"score\":\"not-a-number\",\"nullValue\":null}"); + GenericRow row = new GenericRow(); + row.putValue("variant", variant); + + InbuiltFunctionEvaluator strict = + new InbuiltFunctionEvaluator("variantGet(variant, '$.score', 'DOUBLE')"); + InbuiltFunctionEvaluator tolerant = + new InbuiltFunctionEvaluator("tryVariantGet(variant, '$.score', 'DOUBLE')"); + assertThrows(RuntimeException.class, () -> strict.evaluate(row)); + assertNull(tolerant.evaluate(row)); + assertNull(new InbuiltFunctionEvaluator("variantGet(variant, '$.missing', 'STRING')").evaluate(row)); + assertNull(new InbuiltFunctionEvaluator("tryVariantGet(variant, '$.missing', 'STRING')").evaluate(row)); + assertFalse((Boolean) new InbuiltFunctionEvaluator("variantExists(variant, '$.missing')").evaluate(row)); + assertFalse((Boolean) new InbuiltFunctionEvaluator("isVariantNull(variant, '$.missing')").evaluate(row)); + assertNull(new InbuiltFunctionEvaluator("variantTypeOf(variant, '$.missing')").evaluate(row)); + + row.putValue("variant", new byte[]{1}); + assertThrows(RuntimeException.class, () -> strict.evaluate(row)); + assertNull(tolerant.evaluate(row)); + + row.putValue("variant", null); + assertNull(new InbuiltFunctionEvaluator("variantGet(variant, '$.score', 'STRING')").evaluate(row)); + assertNull(new InbuiltFunctionEvaluator("variantExists(variant, '$.score')").evaluate(row)); + assertFalse((Boolean) new InbuiltFunctionEvaluator("isVariantNull(variant)").evaluate(row)); + assertNull(new InbuiltFunctionEvaluator("variantTypeOf(variant)").evaluate(row)); + } + + @Test + public void testPlannedVariantLiteralsAreValidatedAtConstruction() { + assertThrows(IllegalArgumentException.class, + () -> new InbuiltFunctionEvaluator("variantGet(variant, 'not-a-path', 'STRING')")); + assertThrows(IllegalArgumentException.class, + () -> new InbuiltFunctionEvaluator("variantGet(variant, '$.value', 'NOT_A_TYPE')")); + } + + @Test + public void testDynamicVariantOperandsRetainCompatibilityPath() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"value\":11}"); + InbuiltFunctionEvaluator evaluator = new InbuiltFunctionEvaluator("variantGet(variant, path, targetType)"); + assertEquals(evaluator.getArguments(), List.of("variant", "path", "targetType")); + + GenericRow row = new GenericRow(); + row.putValue("variant", variant); + row.putValue("path", "$.value"); + row.putValue("targetType", "INT"); + assertEquals(evaluator.evaluate(row), 11); + } + + private static byte[] encode(VariantBuilder builder) { + Variant variant = builder.build(); + return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/proto/ExpressionsColumnDataTypeCompatibilityTest.java b/pinot-common/src/test/java/org/apache/pinot/common/proto/ExpressionsColumnDataTypeCompatibilityTest.java new file mode 100644 index 000000000000..349691fd87ab --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/proto/ExpressionsColumnDataTypeCompatibilityTest.java @@ -0,0 +1,87 @@ +/** + * 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.proto; + +import java.util.Map; +import org.apache.pinot.common.proto.legacy.LegacyExpressions; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +/// Locks the permanent protobuf numbers used by query-plan expressions. +/// +/// These numbers are read across broker/server version boundaries. Renumbering an existing value, or assigning a +/// new type to an already allocated number, can silently reinterpret a query plan as a different logical type. +public class ExpressionsColumnDataTypeCompatibilityTest { + @Test + public void testStableWireNumbers() { + Map legacyNumbers = Map.ofEntries( + Map.entry("INT", 0), + Map.entry("LONG", 1), + Map.entry("FLOAT", 2), + Map.entry("DOUBLE", 3), + Map.entry("BIG_DECIMAL", 4), + Map.entry("BOOLEAN", 5), + Map.entry("TIMESTAMP", 6), + Map.entry("STRING", 7), + Map.entry("JSON", 8), + Map.entry("BYTES", 9), + Map.entry("OBJECT", 10), + Map.entry("INT_ARRAY", 11), + Map.entry("LONG_ARRAY", 12), + Map.entry("FLOAT_ARRAY", 13), + Map.entry("DOUBLE_ARRAY", 14), + Map.entry("BOOLEAN_ARRAY", 15), + Map.entry("TIMESTAMP_ARRAY", 16), + Map.entry("STRING_ARRAY", 17), + Map.entry("BYTES_ARRAY", 18), + Map.entry("UNKNOWN", 19), + Map.entry("MAP", 20), + Map.entry("BIG_DECIMAL_ARRAY", 21), + Map.entry("UUID", 22), + Map.entry("UUID_ARRAY", 23)); + + for (Map.Entry entry : legacyNumbers.entrySet()) { + assertEquals(Expressions.ColumnDataType.valueOf(entry.getKey()).getNumber(), entry.getValue().intValue(), + "Wire number changed for " + entry.getKey()); + } + assertEquals(Expressions.ColumnDataType.VARIANT.getNumber(), 24); + } + + @Test + public void testPreVariantAndCurrentPeersHaveDeterministicWireBehavior() + throws Exception { + byte[] legacyPayload = LegacyExpressions.FunctionCall.newBuilder() + .setDataType(LegacyExpressions.ColumnDataType.STRING) + .build() + .toByteArray(); + Expressions.FunctionCall parsedByCurrentPeer = Expressions.FunctionCall.parseFrom(legacyPayload); + assertEquals(parsedByCurrentPeer.getDataType(), Expressions.ColumnDataType.STRING); + assertEquals(parsedByCurrentPeer.getDataTypeValue(), 7); + + byte[] variantPayload = Expressions.FunctionCall.newBuilder() + .setDataType(Expressions.ColumnDataType.VARIANT) + .build() + .toByteArray(); + LegacyExpressions.FunctionCall parsedByLegacyPeer = LegacyExpressions.FunctionCall.parseFrom(variantPayload); + assertEquals(parsedByLegacyPeer.getDataType(), LegacyExpressions.ColumnDataType.UNRECOGNIZED); + assertEquals(parsedByLegacyPeer.getDataTypeValue(), 24); + } +} diff --git a/pinot-common/src/test/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoderTest.java b/pinot-common/src/test/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoderTest.java index ea141a293987..4a7e27f1915a 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoderTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/response/encoder/ArrowResponseEncoderTest.java @@ -31,6 +31,7 @@ import org.apache.pinot.common.response.broker.ResultTable; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.spi.utils.UuidUtils; import org.testng.annotations.Test; @@ -184,6 +185,30 @@ public void testEncodeDecodeUuidColumnWithNulls() assertNull(decodedTable.getRows().get(1)[1], "Null UUID array should round-trip as null"); } + @Test + public void testEncodeDecodeVariantColumnPreservesVariantNullAndSqlNull() + throws IOException { + DataSchema schema = new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + String canonicalJson = VariantUtils.variantToJson( + VariantUtils.parseJsonToVariant("{\"a\":[1,true,null],\"b\":\"text\"}")); + String variantNull = VariantUtils.variantToJson(VariantUtils.parseJsonToVariant("null")); + List rows = Arrays.asList( + new Object[]{canonicalJson}, + new Object[]{variantNull}, + new Object[]{null} + ); + + ArrowResponseEncoder encoder = new ArrowResponseEncoder(); + byte[] encodedBytes = encoder.encodeResultTable(new ResultTable(schema, rows), 0, rows.size()); + ResultTable decodedTable = encoder.decodeResultTable(encodedBytes, rows.size(), schema); + + assertEquals(decodedTable.getRows().size(), 3); + assertEquals(decodedTable.getRows().get(0)[0], "{\"a\":[1,true,null],\"b\":\"text\"}"); + assertEquals(decodedTable.getRows().get(1)[0], "null", + "An encoded Variant null must remain canonical JSON text"); + assertNull(decodedTable.getRows().get(2)[0], "SQL null must remain an Arrow null"); + } + @Test public void testEncodeDecodeAllDataTypes() throws IOException { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/DataSchemaTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/DataSchemaTest.java index ba3bdceef4e5..959c3b63d85b 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/DataSchemaTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/DataSchemaTest.java @@ -18,10 +18,19 @@ */ package org.apache.pinot.common.utils; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; import java.math.BigDecimal; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.sql.Timestamp; import java.util.Locale; +import org.apache.pinot.segment.spi.memory.DataBufferPinotInputStream; +import org.apache.pinot.segment.spi.memory.PinotByteBuffer; +import org.apache.pinot.segment.spi.memory.PinotInputStream; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.BytesUtils; @@ -34,13 +43,13 @@ public class DataSchemaTest { private static final String[] COLUMN_NAMES = { - "int", "long", "float", "double", "string", "uuid", "object", "int_array", "long_array", "float_array", + "int", "long", "float", "double", "string", "variant", "uuid", "object", "int_array", "long_array", "float_array", "double_array", "string_array", "boolean_array", "timestamp_array", "bytes_array", "uuid_array" }; private static final int NUM_COLUMNS = COLUMN_NAMES.length; private static final DataSchema.ColumnDataType[] COLUMN_DATA_TYPES = { - INT, LONG, FLOAT, DOUBLE, STRING, UUID, OBJECT, INT_ARRAY, LONG_ARRAY, FLOAT_ARRAY, DOUBLE_ARRAY, STRING_ARRAY, - BOOLEAN_ARRAY, TIMESTAMP_ARRAY, BYTES_ARRAY, UUID_ARRAY + INT, LONG, FLOAT, DOUBLE, STRING, VARIANT, UUID, OBJECT, INT_ARRAY, LONG_ARRAY, FLOAT_ARRAY, DOUBLE_ARRAY, + STRING_ARRAY, BOOLEAN_ARRAY, TIMESTAMP_ARRAY, BYTES_ARRAY, UUID_ARRAY }; private static final String UUID_VALUE = "550e8400-e29b-41d4-a716-446655440000"; private static final String UUID_VALUE_2 = "550e8400-e29b-41d4-a716-446655440001"; @@ -70,16 +79,20 @@ public void testClone() { public void testSerDe() throws Exception { DataSchema dataSchema = new DataSchema(COLUMN_NAMES, COLUMN_DATA_TYPES); - DataSchema dataSchemaAfterSerDe = DataSchema.fromBytes(ByteBuffer.wrap(dataSchema.toBytes())); + byte[] serialized = dataSchema.toBytes(); + DataSchema dataSchemaAfterSerDe = DataSchema.fromBytes(ByteBuffer.wrap(serialized)); Assert.assertEquals(dataSchema, dataSchemaAfterSerDe); Assert.assertEquals(dataSchema.hashCode(), dataSchemaAfterSerDe.hashCode()); + try (PinotInputStream input = new DataBufferPinotInputStream(PinotByteBuffer.wrap(serialized))) { + Assert.assertEquals(DataSchema.fromBytes(input), dataSchema); + } } @Test public void testToString() { DataSchema dataSchema = new DataSchema(COLUMN_NAMES, COLUMN_DATA_TYPES); Assert.assertEquals(dataSchema.toString(), - "[int(INT),long(LONG),float(FLOAT),double(DOUBLE),string(STRING),uuid(UUID),object(OBJECT)," + "[int(INT),long(LONG),float(FLOAT),double(DOUBLE),string(STRING),variant(VARIANT),uuid(UUID),object(OBJECT)," + "int_array(INT_ARRAY),long_array(LONG_ARRAY),float_array(FLOAT_ARRAY),double_array(DOUBLE_ARRAY)," + "string_array(STRING_ARRAY),boolean_array(BOOLEAN_ARRAY),timestamp_array(TIMESTAMP_ARRAY)," + "bytes_array(BYTES_ARRAY),uuid_array(UUID_ARRAY)]"); @@ -134,6 +147,30 @@ public void testColumnDataType() { Assert.assertFalse(UUID.isCompatible(BYTES)); Assert.assertFalse(UUID.isCompatible(STRING)); + Assert.assertFalse(VARIANT.isNumber()); + Assert.assertFalse(VARIANT.isArray()); + Assert.assertTrue(VARIANT.isCompatible(VARIANT)); + Assert.assertFalse(VARIANT.isCompatible(BYTES)); + Assert.assertFalse(VARIANT.supportsEquality()); + Assert.assertFalse(VARIANT.supportsHashing()); + Assert.assertFalse(VARIANT.supportsOrdering()); + Assert.assertFalse(VARIANT.supportsMinMax()); + Assert.assertFalse(VARIANT.supportsDirectAggregation()); + Assert.assertFalse(VARIANT.supportsPatternMatching()); + Assert.assertTrue(BYTES.supportsEquality()); + Assert.assertTrue(BYTES.supportsHashing()); + Assert.assertTrue(BYTES.supportsOrdering()); + Assert.assertTrue(BYTES.supportsMinMax()); + Assert.assertTrue(BYTES.supportsDirectAggregation()); + Assert.assertTrue(BYTES.supportsPatternMatching()); + Assert.assertFalse(STRING_ARRAY.supportsEquality()); + Assert.assertFalse(STRING_ARRAY.supportsHashing()); + Assert.assertFalse(STRING_ARRAY.supportsOrdering()); + Assert.assertFalse(STRING_ARRAY.supportsMinMax()); + Assert.assertFalse(STRING_ARRAY.supportsPatternMatching()); + Assert.assertEquals(fromDataType(FieldSpec.DataType.VARIANT, true), VARIANT); + Assert.expectThrows(IllegalStateException.class, () -> fromDataType(FieldSpec.DataType.VARIANT, false)); + Assert.assertFalse(OBJECT.isNumber()); Assert.assertFalse(OBJECT.isWholeNumber()); Assert.assertFalse(OBJECT.isArray()); @@ -232,6 +269,97 @@ public void testColumnDataType() { Assert.assertEquals(BYTES.format(bytesValue), BytesUtils.toHexString(bytesValue)); } + @Test + public void testUnknownWireTypeFailsWithUpgradeGuidance() + throws Exception { + byte[] columnName = "payload".getBytes(StandardCharsets.UTF_8); + byte[] unknownType = "FUTURE_LOGICAL_TYPE".getBytes(StandardCharsets.UTF_8); + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (DataOutputStream output = new DataOutputStream(bytes)) { + output.writeInt(1); + output.writeInt(columnName.length); + output.write(columnName); + output.writeInt(unknownType.length); + output.write(unknownType); + } + + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> DataSchema.fromBytes(ByteBuffer.wrap(bytes.toByteArray()))); + Assert.assertTrue(exception.getMessage().contains("Upgrade all brokers and servers")); + Assert.assertTrue(exception.getMessage().contains("FUTURE_LOGICAL_TYPE")); + try (PinotInputStream input = + new DataBufferPinotInputStream(PinotByteBuffer.wrap(bytes.toByteArray()))) { + exception = Assert.expectThrows(IllegalArgumentException.class, () -> DataSchema.fromBytes(input)); + Assert.assertTrue(exception.getMessage().contains("Upgrade all brokers and servers")); + Assert.assertTrue(exception.getMessage().contains("FUTURE_LOGICAL_TYPE")); + } + } + + @Test + public void testPreVariantPeerCompatibilityContract() + throws Exception { + DataSchema legacySchema = + new DataSchema(new String[]{"name", "uuid"}, new DataSchema.ColumnDataType[]{STRING, UUID}); + byte[] legacyPayload = legacySchema.toBytes(); + Assert.assertEquals(readWithPreVariantPeer(legacyPayload), + new PreVariantColumnDataType[]{PreVariantColumnDataType.STRING, PreVariantColumnDataType.UUID}); + Assert.assertEquals(DataSchema.fromBytes(ByteBuffer.wrap(legacyPayload)), legacySchema); + + byte[] variantPayload = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{VARIANT}).toBytes(); + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> readWithPreVariantPeer(variantPayload)); + Assert.assertTrue(exception.getMessage().contains("VARIANT")); + } + + private static PreVariantColumnDataType[] readWithPreVariantPeer(byte[] serialized) + throws IOException { + try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(serialized))) { + int numColumns = input.readInt(); + for (int i = 0; i < numColumns; i++) { + int length = input.readInt(); + byte[] ignoredColumnName = new byte[length]; + input.readFully(ignoredColumnName); + } + PreVariantColumnDataType[] dataTypes = new PreVariantColumnDataType[numColumns]; + for (int i = 0; i < numColumns; i++) { + int length = input.readInt(); + dataTypes[i] = PreVariantColumnDataType.valueOf(new String(input.readNBytes(length), StandardCharsets.UTF_8)); + } + return dataTypes; + } + } + + /// Frozen logical-type names understood by the build immediately before VARIANT. DataSchema is serialized by + /// name, so this small peer model exercises the actual old-reader compatibility boundary without depending on enum + /// ordinals or loading a second Pinot binary into the test JVM. + private enum PreVariantColumnDataType { + INT, + LONG, + FLOAT, + DOUBLE, + BIG_DECIMAL, + BOOLEAN, + TIMESTAMP, + STRING, + JSON, + BYTES, + UUID, + MAP, + OBJECT, + INT_ARRAY, + LONG_ARRAY, + FLOAT_ARRAY, + DOUBLE_ARRAY, + BIG_DECIMAL_ARRAY, + BOOLEAN_ARRAY, + TIMESTAMP_ARRAY, + STRING_ARRAY, + BYTES_ARRAY, + UUID_ARRAY, + UNKNOWN + } + /// The null placeholder must be resolved on the *logical* type. UUID is the only type whose placeholder differs /// from its stored type's: it needs the 16-byte nil UUID, while BYTES supplies a zero-length one. Every other /// logical type must agree with its stored type, otherwise callers that resolve the stored type (DataBlockBuilder, diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java new file mode 100644 index 000000000000..9a7cb3317112 --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java @@ -0,0 +1,541 @@ +/** + * 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; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.sql.Timestamp; +import java.util.Arrays; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils.ResultType; +import org.apache.pinot.common.utils.VariantUtils.ReusableResult; +import org.apache.pinot.common.utils.VariantUtils.VariantPath; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.VariantEnvelope; +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; + + +public class VariantUtilsTest { + @Test + public void testResultTypeContract() { + assertEquals(ResultType.BOOLEAN.getDataType(), DataType.BOOLEAN); + assertEquals(ResultType.BOOLEAN.getSqlTypeName(), SqlTypeName.BOOLEAN); + assertEquals(ResultType.INT.getDataType(), DataType.INT); + assertEquals(ResultType.INT.getSqlTypeName(), SqlTypeName.INTEGER); + assertEquals(ResultType.LONG.getDataType(), DataType.LONG); + assertEquals(ResultType.LONG.getSqlTypeName(), SqlTypeName.BIGINT); + assertEquals(ResultType.FLOAT.getDataType(), DataType.FLOAT); + assertEquals(ResultType.FLOAT.getSqlTypeName(), SqlTypeName.REAL); + assertEquals(ResultType.DOUBLE.getDataType(), DataType.DOUBLE); + assertEquals(ResultType.DOUBLE.getSqlTypeName(), SqlTypeName.DOUBLE); + assertEquals(ResultType.BIG_DECIMAL.getDataType(), DataType.BIG_DECIMAL); + assertEquals(ResultType.BIG_DECIMAL.getSqlTypeName(), SqlTypeName.DECIMAL); + assertEquals(ResultType.STRING.getDataType(), DataType.STRING); + assertEquals(ResultType.STRING.getSqlTypeName(), SqlTypeName.VARCHAR); + assertEquals(ResultType.BYTES.getDataType(), DataType.BYTES); + assertEquals(ResultType.BYTES.getSqlTypeName(), SqlTypeName.VARBINARY); + assertEquals(ResultType.UUID.getDataType(), DataType.UUID); + assertEquals(ResultType.UUID.getSqlTypeName(), SqlTypeName.UUID); + assertEquals(ResultType.TIMESTAMP.getDataType(), DataType.TIMESTAMP); + assertEquals(ResultType.TIMESTAMP.getSqlTypeName(), SqlTypeName.TIMESTAMP); + assertEquals(ResultType.VARIANT.getDataType(), DataType.VARIANT); + assertEquals(ResultType.VARIANT.getSqlTypeName(), SqlTypeName.VARIANT); + assertEquals(ResultType.JSON.getDataType(), DataType.JSON); + assertEquals(ResultType.JSON.getSqlTypeName(), SqlTypeName.VARCHAR); + } + + @Test + public void testDirectBinaryPathExtractionAndPredicates() { + byte[] variant = VariantUtils.parseJsonToVariant( + "{\"eventType\":\"click\",\"items\":[{\"price\":12.5},null],\"active\":true}"); + + assertEquals(VariantUtils.variantGet(variant, "$.eventType", "STRING"), "click"); + assertEquals((double) VariantUtils.variantGet(variant, "$.items[0].price", "DOUBLE"), 12.5); + assertEquals(VariantUtils.variantGet(variant, "$.active", "BOOLEAN"), true); + assertTrue(VariantUtils.variantExists(variant, "$.items[1]")); + assertFalse(VariantUtils.variantExists(variant, "$.missing")); + assertTrue(VariantUtils.isVariantNull(variant, "$.items[1]")); + assertFalse(VariantUtils.isVariantNull(variant, "$.missing")); + assertEquals(VariantUtils.variantTypeOf(variant, "$.items[0]"), "OBJECT"); + assertEquals(VariantUtils.variantTypeOf(variant, "$.items[0].price"), "DECIMAL"); + } + + @Test + public void testStrictAndTolerantExtraction() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\",\"score\":\"not-a-number\"}"); + + assertNull(VariantUtils.variantGet(variant, "$.missing", "STRING")); + assertThrows(IllegalArgumentException.class, () -> VariantUtils.variantGet(variant, "$.score", "DOUBLE")); + assertNull(VariantUtils.tryVariantGet(variant, "$.missing", "STRING")); + assertNull(VariantUtils.tryVariantGet(variant, "$.score", "DOUBLE")); + } + + @Test + public void testReusableResultStrictAndTolerantExtraction() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"value\":7,\"null\":null,\"text\":\"x\"}"); + VariantPath valuePath = VariantUtils.compilePath("$.value"); + ReusableResult result = new ReusableResult(); + + assertTrue(VariantUtils.extractInto(variant, valuePath, ResultType.INT, result)); + assertEquals(result.getIntValue(), 7); + assertFalse(VariantUtils.extractInto(variant, VariantUtils.compilePath("$.missing"), ResultType.INT, result)); + assertFalse(VariantUtils.extractInto(variant, VariantUtils.compilePath("$.null"), ResultType.INT, result)); + assertThrows(IllegalArgumentException.class, + () -> VariantUtils.extractInto(variant, VariantUtils.compilePath("$.text"), ResultType.DOUBLE, result)); + assertFalse( + VariantUtils.tryExtractInto(variant, VariantUtils.compilePath("$.text"), ResultType.DOUBLE, result)); + assertFalse(VariantUtils.tryExtractInto(new byte[]{1}, valuePath, ResultType.INT, result)); + assertThrows(NullPointerException.class, () -> VariantUtils.extractInto( + variant, valuePath, ResultType.INT, null)); + assertThrows(NullPointerException.class, () -> VariantUtils.tryExtractInto( + variant, valuePath, ResultType.INT, null)); + } + + @Test + public void testReusableTolerantHeterogeneousMismatchesAndNumericRange() { + byte[][] rows = { + VariantUtils.parseJsonToVariant("{\"value\":\"not-an-int\"}"), + VariantUtils.parseJsonToVariant("{\"value\":true}"), + VariantUtils.parseJsonToVariant("{\"value\":{}}"), + VariantUtils.parseJsonToVariant("{\"value\":[]}"), + VariantUtils.parseJsonToVariant("{\"value\":2147483648}"), + VariantUtils.parseJsonToVariant("{\"value\":-2147483649}"), + VariantUtils.parseJsonToVariant("{\"value\":1.5}"), + VariantUtils.parseJsonToVariant("{\"value\":9223372036854775808}"), + VariantUtils.parseJsonToVariant("{\"value\":-9223372036854775809}") + }; + ResultType[] resultTypes = { + ResultType.INT, + ResultType.INT, + ResultType.INT, + ResultType.INT, + ResultType.INT, + ResultType.INT, + ResultType.INT, + ResultType.LONG, + ResultType.LONG + }; + VariantPath path = VariantUtils.compilePath("$.value"); + ReusableResult result = new ReusableResult(); + + for (int i = 0; i < rows.length; i++) { + assertFalse(VariantUtils.tryExtractInto(rows[i], path, resultTypes[i], result), + "Expected tolerant conversion to reject row " + i); + } + + assertThrows(IllegalArgumentException.class, + () -> VariantUtils.extractInto(rows[0], path, ResultType.INT, result)); + assertThrows(ArithmeticException.class, + () -> VariantUtils.extractInto(rows[4], path, ResultType.INT, result)); + assertThrows(ArithmeticException.class, + () -> VariantUtils.extractInto(rows[7], path, ResultType.LONG, result)); + + byte[] valid = VariantUtils.parseJsonToVariant("{\"value\":17}"); + assertTrue(VariantUtils.tryExtractInto(valid, path, ResultType.INT, result)); + assertEquals(result.getIntValue(), 17); + } + + @Test + public void testReusableResultParityForEveryResultType() { + VariantBuilder builder = new VariantBuilder(); + builder.appendBoolean(true); + assertReusableParity(encode(builder), ResultType.BOOLEAN); + + builder = new VariantBuilder(); + builder.appendInt(-17); + assertReusableParity(encode(builder), ResultType.INT); + + builder = new VariantBuilder(); + builder.appendLong(9_876_543_210L); + assertReusableParity(encode(builder), ResultType.LONG); + + builder = new VariantBuilder(); + builder.appendFloat(1.25F); + assertReusableParity(encode(builder), ResultType.FLOAT); + + builder = new VariantBuilder(); + builder.appendDouble(-123.5D); + assertReusableParity(encode(builder), ResultType.DOUBLE); + + builder = new VariantBuilder(); + builder.appendDecimal(new BigDecimal("12345678901234567890.1234")); + assertReusableParity(encode(builder), ResultType.BIG_DECIMAL); + + builder = new VariantBuilder(); + builder.appendString("a UTF-8 value \uD83D\uDE00"); + assertReusableParity(encode(builder), ResultType.STRING); + + builder = new VariantBuilder(); + builder.appendBinary(ByteBuffer.wrap(new byte[]{0, 1, -1, 42})); + assertReusableParity(encode(builder), ResultType.BYTES); + + builder = new VariantBuilder(); + builder.appendUUID(UUID.fromString("00112233-4455-6677-8899-aabbccddeeff")); + assertReusableParity(encode(builder), ResultType.UUID); + + builder = new VariantBuilder(); + builder.appendTimestampNanosTz(-1_234_567_890L); + assertReusableParity(encode(builder), ResultType.TIMESTAMP); + + byte[] nested = VariantUtils.parseJsonToVariant("{\"payload\":{\"count\":7}}"); + assertReusableParity(nested, VariantUtils.compilePath("$.payload"), ResultType.VARIANT); + assertReusableParity(nested, VariantUtils.compilePath("$.payload"), ResultType.JSON); + } + + @Test + public void testReusableNumericAndTemporalEncodingParity() { + VariantBuilder builder = new VariantBuilder(); + builder.appendByte((byte) -8); + byte[] byteValue = encode(builder); + assertReusableParity(byteValue, ResultType.INT); + assertReusableParity(byteValue, ResultType.LONG); + assertReusableParity(byteValue, ResultType.FLOAT); + assertReusableParity(byteValue, ResultType.DOUBLE); + assertReusableParity(byteValue, ResultType.BIG_DECIMAL); + + builder = new VariantBuilder(); + builder.appendShort((short) 32_000); + assertReusableParity(encode(builder), ResultType.INT); + + builder = new VariantBuilder(); + builder.appendDecimal(new BigDecimal("123.45")); + assertReusableParity(encode(builder), ResultType.BIG_DECIMAL); + + builder = new VariantBuilder(); + builder.appendDecimal(new BigDecimal("123.00")); + byte[] integralDecimal = encode(builder); + assertReusableParity(integralDecimal, ResultType.INT); + assertReusableParity(integralDecimal, ResultType.LONG); + + builder = new VariantBuilder(); + builder.appendDecimal(new BigDecimal("1234567890123.45")); + assertReusableParity(encode(builder), ResultType.BIG_DECIMAL); + + builder = new VariantBuilder(); + builder.appendString("x".repeat(128)); + assertReusableParity(encode(builder), ResultType.STRING); + + builder = new VariantBuilder(); + builder.appendDate(-1); + byte[] dateValue = encode(builder); + assertReusableParity(dateValue, ResultType.TIMESTAMP); + assertEquals(((Timestamp) VariantUtils.variantGet(dateValue, "$", "TIMESTAMP")).getTime(), + -TimeUnit.DAYS.toMillis(1), "DATE conversion must use UTC epoch days"); + + builder = new VariantBuilder(); + builder.appendTimestampTz(-1_234_567L); + assertReusableParity(encode(builder), ResultType.TIMESTAMP); + + builder = new VariantBuilder(); + builder.appendTimestampNtz(1_234_567L); + assertReusableParity(encode(builder), ResultType.TIMESTAMP); + + builder = new VariantBuilder(); + builder.appendTimestampNanosNtz(-1_234_567L); + assertReusableParity(encode(builder), ResultType.TIMESTAMP); + } + + @Test + public void testReusableTolerantNonFiniteNumericParity() { + VariantBuilder builder = new VariantBuilder(); + builder.appendDouble(Double.POSITIVE_INFINITY); + byte[] positiveInfinity = encode(builder); + ReusableResult result = new ReusableResult(); + VariantPath rootPath = VariantUtils.compilePath("$"); + + assertTrue(VariantUtils.tryExtractInto(positiveInfinity, rootPath, ResultType.FLOAT, result)); + assertEquals(result.getFloatValue(), Float.POSITIVE_INFINITY); + assertTrue(VariantUtils.tryExtractInto(positiveInfinity, rootPath, ResultType.DOUBLE, result)); + assertEquals(result.getDoubleValue(), Double.POSITIVE_INFINITY); + assertThrows(NumberFormatException.class, + () -> VariantUtils.extractInto(positiveInfinity, rootPath, ResultType.BIG_DECIMAL, result)); + assertFalse(VariantUtils.tryExtractInto(positiveInfinity, rootPath, ResultType.BIG_DECIMAL, result)); + } + + @Test + public void testReusableNestedNavigationAndCompiledPredicates() { + byte[] variant = VariantUtils.parseJsonToVariant( + "{\"\uD83D\uDE00\":{\"items\":[{\"value\":11},null]},\"empty\":\"\"}"); + VariantPath valuePath = VariantUtils.compilePath("$.\uD83D\uDE00.items[0].value"); + VariantPath nullPath = VariantUtils.compilePath("$.\uD83D\uDE00.items[1]"); + VariantPath missingPath = VariantUtils.compilePath("$.\uD83D\uDE00.items[2]"); + ReusableResult result = new ReusableResult(); + + assertTrue(VariantUtils.extractInto(variant, valuePath, ResultType.INT, result)); + assertEquals(result.getIntValue(), 11); + assertTrue(VariantUtils.extractInto(variant, VariantUtils.compilePath("$.empty"), ResultType.STRING, result)); + assertEquals(result.getStringValue(), ""); + + assertEquals(VariantUtils.variantExists(variant, valuePath, result), + VariantUtils.variantExists(variant, valuePath)); + assertEquals(VariantUtils.variantExists(variant, missingPath, result), + VariantUtils.variantExists(variant, missingPath)); + assertEquals(VariantUtils.isVariantNull(variant, nullPath, result), + VariantUtils.isVariantNull(variant, nullPath)); + assertEquals(VariantUtils.variantTypeOf(variant, valuePath, result), + VariantUtils.variantTypeOf(variant, valuePath)); + assertNull(VariantUtils.variantExists(new byte[0], valuePath, result)); + assertFalse(VariantUtils.isVariantNull(new byte[0], nullPath, result)); + assertNull(VariantUtils.variantTypeOf(new byte[0], valuePath, result)); + } + + @Test + public void testReusableResultMalformedInputAndReuse() { + VariantPath path = VariantUtils.compilePath("$.items[0]"); + ReusableResult result = new ReusableResult(); + byte[] first = VariantUtils.parseJsonToVariant("{\"items\":[7]}"); + byte[] second = VariantUtils.parseJsonToVariant("{\"items\":[9]}"); + + assertTrue(VariantUtils.extractInto(first, path, ResultType.INT, result)); + assertEquals(result.getIntValue(), 7); + assertFalse(VariantUtils.extractInto(first, VariantUtils.compilePath("$.missing"), ResultType.INT, result)); + + byte[] badMagic = Arrays.copyOf(first, first.length); + badMagic[0] = 0; + assertThrows(IllegalArgumentException.class, + () -> VariantUtils.extractInto(badMagic, path, ResultType.INT, result)); + assertFalse(VariantUtils.tryExtractInto(badMagic, path, ResultType.INT, result)); + + byte[] badMetadataVersion = Arrays.copyOf(first, first.length); + badMetadataVersion[VariantEnvelope.HEADER_SIZE] = + (byte) ((badMetadataVersion[VariantEnvelope.HEADER_SIZE] & 0xF0) | 2); + assertThrows(UnsupportedOperationException.class, + () -> VariantUtils.extractInto(badMetadataVersion, path, ResultType.INT, result)); + assertFalse(VariantUtils.tryExtractInto(badMetadataVersion, path, ResultType.INT, result)); + + byte[] badArrayOffset = VariantUtils.parseJsonToVariant("[7]"); + int valueOffset = VariantEnvelope.HEADER_SIZE + readBigEndianInt(badArrayOffset, 8); + badArrayOffset[valueOffset + 3] = 0x7F; + assertThrows(IllegalArgumentException.class, + () -> VariantUtils.extractInto(badArrayOffset, VariantUtils.compilePath("$[0]"), ResultType.INT, result)); + assertFalse(VariantUtils.tryExtractInto( + badArrayOffset, VariantUtils.compilePath("$[0]"), ResultType.INT, result)); + + assertTrue(VariantUtils.extractInto(second, path, ResultType.INT, result)); + assertEquals(result.getIntValue(), 9); + } + + @Test + public void testSqlNullAndVariantNullRemainDistinct() { + byte[] variantNull = VariantUtils.parseJsonToVariant("null"); + byte[] sqlNullPlaceholder = new byte[0]; + + assertTrue(VariantUtils.isVariantNull(variantNull)); + assertEquals(VariantUtils.variantTypeOf(variantNull), "NULL"); + assertEquals(VariantUtils.variantToJson(variantNull), "null"); + assertTrue(VariantUtils.isVariantNull(VariantUtils.variantGet(variantNull, "$"))); + assertNull(VariantUtils.variantGet(variantNull, "$", "STRING")); + + assertFalse(VariantUtils.isVariantNull(sqlNullPlaceholder)); + assertNull(VariantUtils.variantTypeOf(sqlNullPlaceholder)); + assertNull(VariantUtils.variantToJson(sqlNullPlaceholder)); + assertNull(VariantUtils.variantGet(sqlNullPlaceholder, "$.anything", "STRING")); + } + + @Test + public void testParseRenderAndLogicalDataSchema() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"a\":[1,true,null],\"b\":\"text\"}"); + assertEquals(VariantUtils.variantToJson(variant), "{\"a\":[1,true,null],\"b\":\"text\"}"); + assertNull(VariantUtils.tryParseJsonToVariant("{not-json")); + assertThrows(IllegalArgumentException.class, () -> VariantUtils.parseJsonToVariant("{not-json")); + + assertEquals(ColumnDataType.VARIANT.getStoredType(), ColumnDataType.BYTES); + assertEquals(ColumnDataType.VARIANT.toDataType(), DataType.VARIANT); + assertEquals(ColumnDataType.VARIANT.toExternal(new ByteArray(variant)), variant); + assertEquals(ColumnDataType.VARIANT.convertAndFormat(new ByteArray(variant)), + "{\"a\":[1,true,null],\"b\":\"text\"}"); + } + + @Test + public void testJsonIntegerBoundsAndBigIntegerFallback() { + byte[] variant = VariantUtils.parseJsonToVariant( + "{\"min\":-9223372036854775808,\"max\":9223372036854775807," + + "\"below\":-9223372036854775809,\"above\":9223372036854775808}"); + + assertEquals(VariantUtils.variantTypeOf(variant, "$.min"), "LONG"); + assertEquals(VariantUtils.variantGet(variant, "$.min", "LONG"), Long.MIN_VALUE); + assertEquals(VariantUtils.variantTypeOf(variant, "$.max"), "LONG"); + assertEquals(VariantUtils.variantGet(variant, "$.max", "LONG"), Long.MAX_VALUE); + + assertEquals(VariantUtils.variantTypeOf(variant, "$.below"), "DECIMAL"); + assertEquals(VariantUtils.variantGet(variant, "$.below", "BIG_DECIMAL"), + new BigDecimal("-9223372036854775809")); + assertEquals(VariantUtils.variantTypeOf(variant, "$.above"), "DECIMAL"); + assertEquals(VariantUtils.variantGet(variant, "$.above", "BIG_DECIMAL"), + new BigDecimal("9223372036854775808")); + assertEquals(VariantUtils.variantToJson(VariantUtils.variantGet(variant, "$.below")), + "-9223372036854775809"); + assertEquals(VariantUtils.variantToJson(VariantUtils.variantGet(variant, "$.above")), + "9223372036854775808"); + + VariantPath minPath = VariantUtils.compilePath("$.min"); + VariantPath maxPath = VariantUtils.compilePath("$.max"); + VariantPath abovePath = VariantUtils.compilePath("$.above"); + ReusableResult result = new ReusableResult(); + assertTrue(VariantUtils.tryExtractInto(variant, minPath, ResultType.LONG, result)); + assertEquals(result.getLongValue(), Long.MIN_VALUE); + assertTrue(VariantUtils.tryExtractInto(variant, maxPath, ResultType.LONG, result)); + assertEquals(result.getLongValue(), Long.MAX_VALUE); + assertThrows(ArithmeticException.class, + () -> VariantUtils.extractInto(variant, abovePath, ResultType.LONG, result)); + assertFalse(VariantUtils.tryExtractInto(variant, abovePath, ResultType.LONG, result)); + } + + @Test + public void testJsonDecimalEncodingBoundsAndExponentNormalization() { + String maxDecimal = "9".repeat(38); + String minDecimal = "-" + maxDecimal; + byte[] boundary = VariantUtils.parseJsonToVariant( + "{\"max\":" + maxDecimal + ",\"min\":" + minDecimal + ",\"positiveExponent\":1e3," + + "\"negativeExponent\":-1.25e3,\"maxScale\":1e-38,\"trailingZeros\":1." + + "0".repeat(39) + ",\"zeroExponent\":0e100000000}"); + + assertEquals(VariantUtils.variantGet(boundary, "$.max", "BIG_DECIMAL"), new BigDecimal(maxDecimal)); + assertEquals(VariantUtils.variantGet(boundary, "$.min", "BIG_DECIMAL"), new BigDecimal(minDecimal)); + assertEquals(VariantUtils.variantGet(boundary, "$.positiveExponent", "BIG_DECIMAL"), new BigDecimal("1000")); + assertEquals(VariantUtils.variantGet(boundary, "$.negativeExponent", "BIG_DECIMAL"), new BigDecimal("-1250")); + assertEquals(VariantUtils.variantGet(boundary, "$.maxScale", "BIG_DECIMAL"), new BigDecimal("1e-38")); + assertEquals(VariantUtils.variantGet(boundary, "$.trailingZeros", "BIG_DECIMAL"), BigDecimal.ONE); + assertEquals(VariantUtils.variantGet(boundary, "$.zeroExponent", "BIG_DECIMAL"), BigDecimal.ZERO); + + String oversizedPositive = "9".repeat(39); + String oversizedNegative = "-" + oversizedPositive; + for (String unsupported + : List.of(oversizedPositive, oversizedNegative, "1e38", "1e-39", "1e100000000")) { + assertThrows(IllegalArgumentException.class, () -> VariantUtils.parseJsonToVariant(unsupported)); + assertNull(VariantUtils.tryParseJsonToVariant(unsupported)); + } + } + + @Test + public void testVariantSubtreeExtraction() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"payload\":{\"count\":7}}"); + byte[] subtree = VariantUtils.variantGet(variant, "$.payload"); + + assertEquals(VariantUtils.variantToJson(subtree), "{\"count\":7}"); + assertEquals(VariantUtils.variantGet(subtree, "$.count", "INT"), 7); + assertNull(VariantUtils.variantGet(variant, "$.missing")); + assertNull(VariantUtils.tryVariantGet(variant, "$.missing")); + } + + private static void assertReusableParity(byte[] envelope, ResultType resultType) { + assertReusableParity(envelope, VariantUtils.compilePath("$"), resultType); + } + + private static void assertReusableParity(byte[] envelope, VariantPath path, ResultType resultType) { + Object expected = VariantUtils.variantGet(envelope, path, resultType); + Object tolerantExpected = VariantUtils.tryVariantGet(envelope, path, resultType); + ReusableResult result = new ReusableResult(); + assertTrue(VariantUtils.extractInto(envelope, path, resultType, result)); + assertTrue(VariantUtils.tryExtractInto(envelope, path, resultType, result)); + Object externalValue = result.getExternalValue(resultType); + Object internalValue = result.getInternalValue(resultType); + switch (resultType) { + case BOOLEAN: + assertEquals(result.getIntValue() != 0, expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, (boolean) expected ? 1 : 0); + break; + case INT: + assertEquals(result.getIntValue(), expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, expected); + break; + case LONG: + assertEquals(result.getLongValue(), expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, expected); + break; + case FLOAT: + assertEquals(result.getFloatValue(), expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, expected); + break; + case DOUBLE: + assertEquals(result.getDoubleValue(), expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, expected); + break; + case BIG_DECIMAL: + assertEquals(result.getBigDecimalValue(), expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, expected); + break; + case STRING: + case JSON: + assertEquals(result.getStringValue(), expected); + assertEquals(externalValue, expected); + assertEquals(internalValue, expected); + break; + case BYTES: + case VARIANT: + assertTrue(Arrays.equals(result.getBytesValue(), (byte[]) expected)); + assertTrue(Arrays.equals((byte[]) externalValue, (byte[]) expected)); + assertSame(((ByteArray) internalValue).getBytes(), result.getBytesValue()); + assertTrue(Arrays.equals(((ByteArray) internalValue).getBytes(), (byte[]) expected)); + break; + case UUID: + assertEquals(result.getUuidValue(), expected); + assertEquals(externalValue, expected); + assertTrue(Arrays.equals(result.getBytesValue(), UuidUtils.toBytes((UUID) expected))); + assertSame(((ByteArray) internalValue).getBytes(), result.getBytesValue()); + assertTrue(Arrays.equals(((ByteArray) internalValue).getBytes(), UuidUtils.toBytes((UUID) expected))); + break; + case TIMESTAMP: + assertEquals(result.getLongValue(), ((Timestamp) expected).getTime()); + assertEquals(externalValue, expected); + assertEquals(internalValue, ((Timestamp) expected).getTime()); + break; + default: + throw new IllegalStateException("Unhandled result type: " + resultType); + } + if (expected instanceof byte[]) { + assertTrue(Arrays.equals((byte[]) tolerantExpected, (byte[]) expected)); + } else { + assertEquals(tolerantExpected, expected); + } + } + + private static byte[] encode(VariantBuilder builder) { + Variant variant = builder.build(); + return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); + } + + private static int readBigEndianInt(byte[] bytes, int offset) { + return Byte.toUnsignedInt(bytes[offset]) << 24 + | Byte.toUnsignedInt(bytes[offset + 1]) << 16 + | Byte.toUnsignedInt(bytes[offset + 2]) << 8 + | Byte.toUnsignedInt(bytes[offset + 3]); + } +} diff --git a/pinot-common/src/test/proto/legacy_expressions.proto b/pinot-common/src/test/proto/legacy_expressions.proto new file mode 100644 index 000000000000..1e6e0039ff95 --- /dev/null +++ b/pinot-common/src/test/proto/legacy_expressions.proto @@ -0,0 +1,58 @@ +// +// 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. +// + +syntax = "proto3"; + +package org.apache.pinot.common.proto.legacy; + +option java_package = "org.apache.pinot.common.proto.legacy"; +option java_outer_classname = "LegacyExpressions"; + +// Frozen pre-VARIANT view of the expression wire enum. This intentionally omits +// VARIANT so tests exercise how an un-upgraded peer handles wire value 24. +enum ColumnDataType { + INT = 0; + LONG = 1; + FLOAT = 2; + DOUBLE = 3; + BIG_DECIMAL = 4; + BOOLEAN = 5; + TIMESTAMP = 6; + STRING = 7; + JSON = 8; + BYTES = 9; + OBJECT = 10; + INT_ARRAY = 11; + LONG_ARRAY = 12; + FLOAT_ARRAY = 13; + DOUBLE_ARRAY = 14; + BOOLEAN_ARRAY = 15; + TIMESTAMP_ARRAY = 16; + STRING_ARRAY = 17; + BYTES_ARRAY = 18; + UNKNOWN = 19; + MAP = 20; + BIG_DECIMAL_ARRAY = 21; + UUID = 22; + UUID_ARRAY = 23; +} + +message FunctionCall { + ColumnDataType data_type = 1; +} diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotDdlRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotDdlRestletResource.java index 9a7c72caadd9..388999fec2db 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotDdlRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotDdlRestletResource.java @@ -30,6 +30,7 @@ import io.swagger.annotations.SwaggerDefinition; import java.math.BigDecimal; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -589,6 +590,10 @@ private static boolean defaultValuesEqual(FieldSpec.DataType dataType, && compiledDefault instanceof BigDecimal) { return ((BigDecimal) storedDefault).compareTo((BigDecimal) compiledDefault) == 0; } + if (dataType == FieldSpec.DataType.VARIANT) { + // Schema defaults are structural metadata; this does not enable semantic equality for Variant values. + return Arrays.equals((byte[]) storedDefault, (byte[]) compiledDefault); + } return dataType.equals(storedDefault, compiledDefault); } diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotDdlRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotDdlRestletResourceTest.java index 8d1b9d6a857a..d2ebab69d8d4 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotDdlRestletResourceTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotDdlRestletResourceTest.java @@ -46,6 +46,7 @@ public class PinotDdlRestletResourceTest extends ControllerTest { private static final String TBL_BASIC = "ddlBasicOffline"; private static final String TBL_DRY_RUN = "ddlDryRunOffline"; private static final String TBL_IF_NOT_EXISTS = "ddlIfNotExistsOffline"; + private static final String TBL_VARIANT = "ddlVariantRoundtrip"; private static final String TBL_DROP = "ddlDropOffline"; @BeforeClass @@ -118,6 +119,46 @@ public void createIfNotExistsIsIdempotent() "Expected idempotent message, got: " + second.get("message").asText()); } + @Test + public void variantCreateIfNotExistsAndShowCreateRoundTrips() + throws IOException { + String createSql = "CREATE TABLE IF NOT EXISTS " + TBL_VARIANT + " (" + + " id INT NOT NULL DIMENSION," + + " payload VARIANT DIMENSION" + + ") TABLE_TYPE = OFFLINE PROPERTIES (" + + " 'replication' = '1'," + + " 'nullHandlingEnabled' = 'true'," + + " 'fieldConfigs' = " + + "'[{\"name\":\"payload\",\"encodingType\":\"RAW\"," + + "\"indexes\":{\"forward\":{\"compressionCodec\":\"ZSTANDARD\"}}}]'" + + ")"; + + JsonNode first = postDdl(createSql, false); + assertEquals(first.get("operation").asText(), "CREATE_TABLE"); + assertEquals(first.get("tableName").asText(), TBL_VARIANT + "_OFFLINE"); + + JsonNode second = postDdl(createSql, false); + assertEquals(second.get("operation").asText(), "CREATE_TABLE"); + assertTrue(second.get("message").asText().toLowerCase().contains("exist"), + "Expected idempotent message, got: " + second.get("message").asText()); + + JsonNode show = postDdl("SHOW CREATE TABLE " + TBL_VARIANT, false); + assertEquals(show.get("operation").asText(), "SHOW_CREATE_TABLE"); + String canonicalDdl = show.get("ddl").asText(); + assertTrue(canonicalDdl.contains("payload VARIANT DIMENSION"), canonicalDdl); + assertFalse(canonicalDdl.contains("payload VARIANT DEFAULT"), + "The reserved VARIANT SQL-null sentinel must not be emitted as a user default:\n" + canonicalDdl); + + // DROP intentionally leaves the shared schema in place. Replaying SHOW CREATE therefore exercises the + // controller's stored-vs-compiled schema comparison, where VARIANT's byte[] SQL-null sentinel must compare by + // content instead of array identity. + postDdl("DROP TABLE " + TBL_VARIANT, false); + DEFAULT_INSTANCE.waitForEVToDisappear(TBL_VARIANT + "_OFFLINE"); + JsonNode recreated = postDdl(canonicalDdl, false); + assertEquals(recreated.get("operation").asText(), "CREATE_TABLE"); + assertEquals(recreated.get("tableName").asText(), TBL_VARIANT + "_OFFLINE"); + } + @Test public void createWithoutIfNotExistsConflicts() throws IOException { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProvider.java index 822a61b05c92..99992abaaab7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProvider.java @@ -67,6 +67,45 @@ public static PredicateEvaluator getPredicateEvaluator(Predicate predicate, @Nul private static PredicateEvaluator buildEvaluator(Predicate predicate, @Nullable Dictionary dictionary, DataType dataType, @Nullable QueryContext queryContext, @Nullable DataSource dataSource) { try { + boolean predicateSupported; + switch (predicate.getType()) { + case EQ: + case NOT_EQ: + predicateSupported = dataType.supportsEquality(); + break; + case IN: + case NOT_IN: + predicateSupported = dataType.supportsEquality() && dataType.supportsHashing(); + break; + case RANGE: + predicateSupported = dataType.supportsOrdering(); + break; + case REGEXP_LIKE: + predicateSupported = dataType.supportsPatternMatching(); + break; + default: + predicateSupported = true; + break; + } + if (!predicateSupported) { + String operation; + switch (predicate.getType()) { + case IN: + case NOT_IN: + operation = "IN"; + break; + case EQ: + case NOT_EQ: + case RANGE: + operation = "comparison"; + break; + default: + operation = predicate.getType().name(); + break; + } + throw new IllegalArgumentException( + "Raw VARIANT values do not support " + operation + "; extract a typed path with variantGet first"); + } if (dictionary != null) { // dictionary based predicate evaluators switch (predicate.getType()) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java index f910cd0bda6e..a4fa711b8c45 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java @@ -49,6 +49,7 @@ import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor; import org.apache.pinot.core.util.GroupByUtils; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.query.QueryScanCostContext; import org.apache.pinot.spi.trace.Tracing; import org.slf4j.Logger; @@ -97,8 +98,12 @@ public FilteredGroupByOperator(QueryContext queryContext, List for (int i = 0; i < numGroupByExpressions; i++) { ExpressionContext groupByExpression = _groupByExpressions[i]; columnNames[i] = groupByExpression.toString(); - columnDataTypes[i] = DataSchema.ColumnDataType.fromDataTypeSV( - projectOperator.getResultColumnContext(groupByExpression).getDataType()); + DataType dataType = projectOperator.getResultColumnContext(groupByExpression).getDataType(); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw new IllegalArgumentException( + "Raw VARIANT values do not support GROUP BY; extract a typed path with variantGet first"); + } + columnDataTypes[i] = DataSchema.ColumnDataType.fromDataTypeSV(dataType); } /// Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java index 50afa03d62a7..02d326ea652a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java @@ -43,6 +43,7 @@ import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor; import org.apache.pinot.core.util.GroupByUtils; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.query.QueryScanCostContext; import org.apache.pinot.spi.trace.Tracing; import org.slf4j.Logger; @@ -89,8 +90,12 @@ public GroupByOperator(QueryContext queryContext, AggregationInfo aggregationInf for (int i = 0; i < numGroupByExpressions; i++) { ExpressionContext groupByExpression = _groupByExpressions[i]; columnNames[i] = groupByExpression.toString(); - columnDataTypes[i] = DataSchema.ColumnDataType.fromDataTypeSV( - _projectOperator.getResultColumnContext(groupByExpression).getDataType()); + DataType dataType = _projectOperator.getResultColumnContext(groupByExpression).getDataType(); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw new IllegalArgumentException( + "Raw VARIANT values do not support GROUP BY; extract a typed path with variantGet first"); + } + columnDataTypes[i] = DataSchema.ColumnDataType.fromDataTypeSV(dataType); } /// Synthetic grouping-id discriminator column for GROUP BY GROUPING SETS / ROLLUP / CUBE diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java new file mode 100644 index 000000000000..33756eaa0f99 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java @@ -0,0 +1,161 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.common.utils.VariantUtils.ReusableResult; +import org.apache.pinot.common.utils.VariantUtils.VariantPath; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.roaringbitmap.RoaringBitmap; + + +/** + * Shared single-stage lifecycle for functions that inspect or extract a VARIANT value. + * + *

The first argument is consistently validated as a single-value VARIANT-compatible operand, the optional path is + * compiled once, and each input block is evaluated at most once even when the engine asks for both values and a null + * bitmap. Instances are query-local and not thread-safe. + */ +abstract class BaseVariantTransformFunction extends BaseTransformFunction { + private static final VariantPath ROOT_PATH = VariantUtils.compilePath("$"); + + private TransformFunction _variantTransformFunction; + private VariantPath _path; + private final ReusableResult _reusableResult = new ReusableResult(); + @Nullable + private ValueBlock _cachedValueBlock; + @Nullable + private RoaringBitmap _cachedNullBitmap; + + /** + * Initializes the common VARIANT operand and path contract. + * + * @param arguments function arguments + * @param minArguments minimum accepted argument count + * @param maxArguments maximum accepted argument count + * @param pathRequired whether argument 1 is mandatory + */ + protected final void initVariantArguments(List arguments, int minArguments, int maxArguments, + boolean pathRequired) { + int numArguments = arguments.size(); + if (numArguments < minArguments || numArguments > maxArguments) { + throw new IllegalArgumentException(argumentCountMessage(minArguments, maxArguments)); + } + + _variantTransformFunction = arguments.get(0); + TransformResultMetadata inputMetadata = _variantTransformFunction.getResultMetadata(); + DataType inputType = inputMetadata.getDataType(); + if (!inputMetadata.isSingleValue() + || (inputType != DataType.VARIANT && inputType != DataType.BYTES && inputType != DataType.UNKNOWN)) { + throw new IllegalArgumentException(getName() + " first argument must be a single-value VARIANT"); + } + + if (pathRequired || numArguments >= 2) { + TransformFunction pathTransformFunction = arguments.get(1); + if (!(pathTransformFunction instanceof LiteralTransformFunction) + || pathTransformFunction.getResultMetadata().getDataType() != DataType.STRING) { + throw new IllegalArgumentException(getName() + " path must be a string literal"); + } + _path = VariantUtils.compilePath(((LiteralTransformFunction) pathTransformFunction).getStringLiteral()); + } else { + _path = ROOT_PATH; + } + _cachedValueBlock = null; + _cachedNullBitmap = null; + } + + private String argumentCountMessage(int minArguments, int maxArguments) { + if (minArguments == maxArguments) { + return getName() + " expects exactly " + minArguments + " arguments"; + } + return getName() + " expects " + minArguments + " to " + maxArguments + " arguments"; + } + + protected final VariantPath getVariantPath() { + return _path; + } + + protected final ReusableResult getReusableResult() { + return _reusableResult; + } + + @Nullable + @Override + public final RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + ensureEvaluated(valueBlock); + return _cachedNullBitmap; + } + + protected final void ensureEvaluated(ValueBlock valueBlock) { + if (_cachedValueBlock == valueBlock) { + return; + } + + // Invalidate before mutating reusable output arrays. A strict extraction failure must never leave a partially + // overwritten block marked as cached. + _cachedValueBlock = null; + _cachedNullBitmap = null; + int numDocs = valueBlock.getNumDocs(); + initResultValues(numDocs); + + RoaringBitmap inputNulls = _variantTransformFunction.getNullBitmap(valueBlock); + RoaringBitmap resultNulls = + inputNullIsResultNull() && inputNulls != null ? inputNulls.clone() : new RoaringBitmap(); + byte[][] variants = _variantTransformFunction.transformToBytesValuesSV(valueBlock); + for (int i = 0; i < numDocs; i++) { + if (inputNulls != null && inputNulls.contains(i)) { + setNullValue(i); + continue; + } + if (!evaluateVariant(variants[i], i)) { + resultNulls.add(i); + setNullValue(i); + } + } + + _cachedNullBitmap = resultNulls.isEmpty() ? null : resultNulls; + _cachedValueBlock = valueBlock; + } + + /** + * Returns whether an input SQL null should remain SQL null in the result. + */ + protected boolean inputNullIsResultNull() { + return true; + } + + /** + * Initializes the output array for a block. + */ + protected abstract void initResultValues(int numDocs); + + /** + * Evaluates one non-SQL-null VARIANT. Returns {@code false} when the result should be SQL null. + */ + protected abstract boolean evaluateVariant(@Nullable byte[] variant, int index); + + /** + * Stores the type-specific placeholder for a SQL-null result. + */ + protected abstract void setNullValue(int index); +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java index f6a2101d43b3..a7fed5b82b89 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java @@ -108,6 +108,8 @@ public void init(List arguments, Map c _rightTransformFunction = arguments.get(1); DataType leftDataType = _leftTransformFunction.getResultMetadata().getDataType(); DataType rightDataType = _rightTransformFunction.getResultMetadata().getDataType(); + Preconditions.checkArgument(leftDataType.supportsOrdering() && rightDataType.supportsOrdering(), + "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); _leftStoredType = leftDataType.getStoredType(); _rightStoredType = rightDataType.getStoredType(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java index 00dc81eb3ad2..c75b52d9313a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java @@ -64,6 +64,11 @@ public void init(List arguments, Map c Preconditions.checkArgument(numArguments >= 2, "At least 2 arguments are required for [%s] " + "transform function: (expression, values)", getName()); _mainFunction = arguments.get(0); + for (TransformFunction argument : arguments) { + DataType dataType = argument.getResultMetadata().getDataType(); + Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), + "Raw VARIANT values do not support IN; extract a typed path with variantGet first"); + } boolean allLiteralValues = true; ObjectOpenHashSet stringValues = new ObjectOpenHashSet<>(numArguments - 1); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java new file mode 100644 index 000000000000..b99be20e70d5 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java @@ -0,0 +1,82 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; + + +/** + * Returns whether a Variant root value or a value selected by a literal path is an encoded Variant null. + * + *

SQL null and missing paths return a non-null {@code false}; only a present encoded Variant null returns + * {@code true}. Instances are query-local and not thread-safe. + */ +public class IsVariantNullTransformFunction extends BaseVariantTransformFunction { + + @Override + public String getName() { + return TransformFunctionType.IS_VARIANT_NULL.getName(); + } + + @Override + public void init(List arguments, Map columnContextMap, + boolean nullHandlingEnabled) { + super.init(arguments, columnContextMap, nullHandlingEnabled); + initVariantArguments(arguments, 1, 2, false); + } + + @Override + public TransformResultMetadata getResultMetadata() { + return BOOLEAN_SV_NO_DICTIONARY_METADATA; + } + + @Override + public int[] transformToIntValuesSV(ValueBlock valueBlock) { + ensureEvaluated(valueBlock); + return _intValuesSV; + } + + @Override + protected void initResultValues(int numDocs) { + initIntValuesSV(numDocs); + } + + @Override + protected boolean evaluateVariant(@Nullable byte[] variant, int index) { + _intValuesSV[index] = VariantUtils.isVariantNull(variant, getVariantPath(), getReusableResult()) ? 1 : 0; + return true; + } + + @Override + protected void setNullValue(int index) { + _intValuesSV[index] = 0; + } + + @Override + protected boolean inputNullIsResultNull() { + return false; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java new file mode 100644 index 000000000000..eb8486e5ba53 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java @@ -0,0 +1,168 @@ +/** + * 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.operator.transform.function; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder; +import org.roaringbitmap.RoaringBitmap; + + +/** + * Parses JSON text into a logical Variant value for single-stage queries and ingestion transforms. + * + *

Instances are query-local and not thread-safe. + */ +public class ParseJsonToVariantTransformFunction extends BaseTransformFunction { + public static final String FUNCTION_NAME = "parseJson"; + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.VARIANT, true, false); + + private final boolean _tolerant; + private TransformFunction _jsonTransformFunction; + private boolean _literalInput; + @Nullable + private byte[] _literalResult; + @Nullable + private ValueBlock _cachedValueBlock; + @Nullable + private RoaringBitmap _cachedNullBitmap; + + public ParseJsonToVariantTransformFunction() { + this(false); + } + + protected ParseJsonToVariantTransformFunction(boolean tolerant) { + _tolerant = tolerant; + } + + @Override + public String getName() { + return _tolerant ? Try.FUNCTION_NAME : FUNCTION_NAME; + } + + @Override + public void init(List arguments, Map columnContextMap, + boolean nullHandlingEnabled) { + super.init(arguments, columnContextMap, nullHandlingEnabled); + if (arguments.size() != 1) { + throw new IllegalArgumentException(getName() + " expects exactly one argument"); + } + _jsonTransformFunction = arguments.get(0); + TransformResultMetadata inputMetadata = _jsonTransformFunction.getResultMetadata(); + boolean sqlNullLiteral = + _jsonTransformFunction instanceof LiteralTransformFunction + && ((LiteralTransformFunction) _jsonTransformFunction).isNull(); + if (!inputMetadata.isSingleValue() + || (inputMetadata.getDataType().getStoredType() != DataType.STRING && !sqlNullLiteral)) { + throw new IllegalArgumentException(getName() + " argument must be a single-value STRING or JSON"); + } + _cachedValueBlock = null; + _cachedNullBitmap = null; + _literalInput = _jsonTransformFunction instanceof LiteralTransformFunction; + if (_literalInput) { + LiteralTransformFunction literal = (LiteralTransformFunction) _jsonTransformFunction; + _literalResult = literal.isNull() ? null : parse(literal.getStringLiteral()); + } else { + _literalResult = null; + } + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + ensureParsed(valueBlock); + return _cachedNullBitmap; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + ensureParsed(valueBlock); + return _bytesValuesSV; + } + + private void ensureParsed(ValueBlock valueBlock) { + if (_cachedValueBlock == valueBlock) { + return; + } + // Invalidate before mutating the reusable result buffer. A strict parse failure must not leave a previous block + // marked as cached after its values have been partially overwritten. + _cachedValueBlock = null; + _cachedNullBitmap = null; + int numDocs = valueBlock.getNumDocs(); + initBytesValuesSV(numDocs); + if (_literalInput) { + byte[] result = _literalResult; + Arrays.fill(_bytesValuesSV, 0, numDocs, result != null ? result : NullValuePlaceHolder.BYTES); + if (result == null && numDocs != 0) { + RoaringBitmap resultNulls = new RoaringBitmap(); + resultNulls.add(0L, numDocs); + _cachedNullBitmap = resultNulls; + } + _cachedValueBlock = valueBlock; + return; + } + RoaringBitmap inputNulls = _jsonTransformFunction.getNullBitmap(valueBlock); + RoaringBitmap resultNulls = inputNulls != null ? inputNulls.clone() : new RoaringBitmap(); + String[] jsonValues = _jsonTransformFunction.transformToStringValuesSV(valueBlock); + for (int i = 0; i < numDocs; i++) { + if (resultNulls.contains(i)) { + _bytesValuesSV[i] = NullValuePlaceHolder.BYTES; + continue; + } + byte[] value = parse(jsonValues[i]); + if (value == null) { + resultNulls.add(i); + _bytesValuesSV[i] = NullValuePlaceHolder.BYTES; + } else { + _bytesValuesSV[i] = value; + } + } + _cachedNullBitmap = resultNulls.isEmpty() ? null : resultNulls; + _cachedValueBlock = valueBlock; + } + + @Nullable + private byte[] parse(@Nullable String json) { + return _tolerant ? VariantUtils.tryParseJsonToVariant(json) : VariantUtils.parseJsonToVariant(json); + } + + /** + * Tolerant JSON-to-Variant parsing. + */ + public static final class Try extends ParseJsonToVariantTransformFunction { + public static final String FUNCTION_NAME = "tryParseJson"; + + public Try() { + super(true); + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java index c4647f9a5bc8..f9e34c948167 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java @@ -27,7 +27,6 @@ import java.util.List; import java.util.Map; import java.util.Set; -import org.apache.commons.lang3.StringUtils; import org.apache.pinot.common.function.FunctionInfo; import org.apache.pinot.common.function.FunctionRegistry; import org.apache.pinot.common.function.TransformFunctionType; @@ -135,6 +134,14 @@ private static Map> createRegistry() typeToImplementation.put(TransformFunctionType.JSON_EXTRACT_SCALAR_FORY, JsonExtractScalarTransformFunction.Fory.class); typeToImplementation.put(TransformFunctionType.JSON_EXTRACT_KEY, JsonExtractKeyTransformFunction.class); + typeToImplementation.put(TransformFunctionType.VARIANT_GET, VariantGetTransformFunction.class); + typeToImplementation.put(TransformFunctionType.TRY_VARIANT_GET, VariantGetTransformFunction.Try.class); + typeToImplementation.put(TransformFunctionType.VARIANT_EXISTS, VariantExistsTransformFunction.class); + typeToImplementation.put(TransformFunctionType.IS_VARIANT_NULL, IsVariantNullTransformFunction.class); + typeToImplementation.put(TransformFunctionType.VARIANT_TYPE_OF, VariantTypeOfTransformFunction.class); + typeToImplementation.put(TransformFunctionType.PARSE_JSON_TO_VARIANT, ParseJsonToVariantTransformFunction.class); + typeToImplementation.put(TransformFunctionType.TRY_PARSE_JSON_TO_VARIANT, + ParseJsonToVariantTransformFunction.Try.class); typeToImplementation.put(TransformFunctionType.TIME_CONVERT, TimeConversionTransformFunction.class); typeToImplementation.put(TransformFunctionType.DATE_TIME_CONVERT, DateTimeConversionTransformFunction.class); typeToImplementation.put(TransformFunctionType.DATE_TIME_CONVERT_WINDOW_HOP, @@ -312,7 +319,15 @@ public static TransformFunction get(ExpressionContext expression, Map> getAllFunctions() { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java new file mode 100644 index 000000000000..0dec8fe85882 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java @@ -0,0 +1,82 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; + + +/** + * Vectorized single-stage implementation of {@code variantExists}. + * + *

The path must be a string literal and is compiled once during initialization. A present encoded Variant null + * counts as present, a missing path returns {@code false}, and SQL null remains SQL null. Instances are query-local + * and not thread-safe. + */ +public class VariantExistsTransformFunction extends BaseVariantTransformFunction { + public static final String FUNCTION_NAME = "variantExists"; + + @Override + public String getName() { + return FUNCTION_NAME; + } + + @Override + public void init(List arguments, Map columnContextMap, + boolean nullHandlingEnabled) { + super.init(arguments, columnContextMap, nullHandlingEnabled); + initVariantArguments(arguments, 2, 2, true); + } + + @Override + public TransformResultMetadata getResultMetadata() { + return BOOLEAN_SV_NO_DICTIONARY_METADATA; + } + + @Override + public int[] transformToIntValuesSV(ValueBlock valueBlock) { + ensureEvaluated(valueBlock); + return _intValuesSV; + } + + @Override + protected void initResultValues(int numDocs) { + initIntValuesSV(numDocs); + } + + @Override + protected boolean evaluateVariant(@Nullable byte[] variant, int index) { + Boolean exists = VariantUtils.variantExists(variant, getVariantPath(), getReusableResult()); + if (exists != null) { + _intValuesSV[index] = exists ? 1 : 0; + return true; + } + return false; + } + + @Override + protected void setNullValue(int index) { + _intValuesSV[index] = 0; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java new file mode 100644 index 000000000000..43d8e688505d --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java @@ -0,0 +1,253 @@ +/** + * 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.operator.transform.function; + +import java.math.BigDecimal; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.common.utils.VariantUtils.ResultType; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder; + + +/** + * Single-stage typed extraction from a Variant envelope. + * + *

The path and optional target type must be literals, so they are parsed once during initialization. Omitting the + * target type returns a Variant. A missing path returns SQL null; the strict form throws for an incompatible non-null + * value, while {@link Try} maps that failure to SQL null. Instances are query-local and not thread-safe. + */ +public class VariantGetTransformFunction extends BaseVariantTransformFunction { + public static final String FUNCTION_NAME = "variantGet"; + private final boolean _tolerant; + private ResultType _targetType; + private DataType _storedType; + private TransformResultMetadata _resultMetadata; + + public VariantGetTransformFunction() { + this(false); + } + + protected VariantGetTransformFunction(boolean tolerant) { + _tolerant = tolerant; + } + + @Override + public String getName() { + return _tolerant ? Try.FUNCTION_NAME : FUNCTION_NAME; + } + + @Override + public void init(List arguments, Map columnContextMap, + boolean nullHandlingEnabled) { + super.init(arguments, columnContextMap, nullHandlingEnabled); + initVariantArguments(arguments, 2, 3, true); + int numArguments = arguments.size(); + if (numArguments == 3 && (!(arguments.get(2) instanceof LiteralTransformFunction) + || arguments.get(2).getResultMetadata().getDataType() != DataType.STRING)) { + throw new IllegalArgumentException(getName() + " target type must be a string literal"); + } + _targetType = numArguments == 2 ? ResultType.VARIANT + : VariantUtils.parseResultType(((LiteralTransformFunction) arguments.get(2)).getStringLiteral()); + DataType resultType = _targetType.getDataType(); + _storedType = resultType.getStoredType(); + _resultMetadata = new TransformResultMetadata(resultType, true, false); + } + + @Override + public TransformResultMetadata getResultMetadata() { + return _resultMetadata; + } + + @Override + public int[] transformToIntValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.INT) { + return super.transformToIntValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _intValuesSV; + } + + @Override + public long[] transformToLongValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.LONG) { + return super.transformToLongValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _longValuesSV; + } + + @Override + public float[] transformToFloatValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.FLOAT) { + return super.transformToFloatValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _floatValuesSV; + } + + @Override + public double[] transformToDoubleValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.DOUBLE) { + return super.transformToDoubleValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _doubleValuesSV; + } + + @Override + public BigDecimal[] transformToBigDecimalValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.BIG_DECIMAL) { + return super.transformToBigDecimalValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _bigDecimalValuesSV; + } + + @Override + public String[] transformToStringValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.STRING) { + return super.transformToStringValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _stringValuesSV; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + if (_storedType != DataType.BYTES) { + return super.transformToBytesValuesSV(valueBlock); + } + ensureEvaluated(valueBlock); + return _bytesValuesSV; + } + + @Override + protected void initResultValues(int numDocs) { + switch (_storedType) { + case INT: + initIntValuesSV(numDocs); + break; + case LONG: + initLongValuesSV(numDocs); + break; + case FLOAT: + initFloatValuesSV(numDocs); + break; + case DOUBLE: + initDoubleValuesSV(numDocs); + break; + case BIG_DECIMAL: + initBigDecimalValuesSV(numDocs); + break; + case STRING: + initStringValuesSV(numDocs); + break; + case BYTES: + initBytesValuesSV(numDocs); + break; + default: + throw new IllegalStateException("Unsupported Variant result storage type: " + _storedType); + } + } + + @Override + protected boolean evaluateVariant(@Nullable byte[] variant, int index) { + boolean extracted = _tolerant + ? VariantUtils.tryExtractInto(variant, getVariantPath(), _targetType, getReusableResult()) + : VariantUtils.extractInto(variant, getVariantPath(), _targetType, getReusableResult()); + if (extracted) { + setExtractedValue(index); + } + return extracted; + } + + @Override + protected void setNullValue(int index) { + switch (_storedType) { + case INT: + _intValuesSV[index] = NullValuePlaceHolder.INT; + break; + case LONG: + _longValuesSV[index] = NullValuePlaceHolder.LONG; + break; + case FLOAT: + _floatValuesSV[index] = NullValuePlaceHolder.FLOAT; + break; + case DOUBLE: + _doubleValuesSV[index] = NullValuePlaceHolder.DOUBLE; + break; + case BIG_DECIMAL: + _bigDecimalValuesSV[index] = NullValuePlaceHolder.BIG_DECIMAL; + break; + case STRING: + _stringValuesSV[index] = NullValuePlaceHolder.STRING; + break; + case BYTES: + _bytesValuesSV[index] = NullValuePlaceHolder.BYTES; + break; + default: + throw new IllegalStateException("Unsupported Variant result storage type: " + _storedType); + } + } + + private void setExtractedValue(int index) { + switch (_storedType) { + case INT: + _intValuesSV[index] = getReusableResult().getIntValue(); + break; + case LONG: + _longValuesSV[index] = getReusableResult().getLongValue(); + break; + case FLOAT: + _floatValuesSV[index] = getReusableResult().getFloatValue(); + break; + case DOUBLE: + _doubleValuesSV[index] = getReusableResult().getDoubleValue(); + break; + case BIG_DECIMAL: + _bigDecimalValuesSV[index] = getReusableResult().getBigDecimalValue(); + break; + case STRING: + _stringValuesSV[index] = getReusableResult().getStringValue(); + break; + case BYTES: + _bytesValuesSV[index] = getReusableResult().getBytesValue(); + break; + default: + throw new IllegalStateException("Unsupported Variant result storage type: " + _storedType); + } + } + + /** + * Tolerant Variant extraction. + */ + public static final class Try extends VariantGetTransformFunction { + public static final String FUNCTION_NAME = "tryVariantGet"; + + public Try() { + super(true); + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java new file mode 100644 index 000000000000..75bd1f9af027 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java @@ -0,0 +1,82 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder; + + +/** + * Returns the type name of a Variant root value or a value selected by a literal path. + * + *

SQL null and missing paths produce SQL null. An encoded Variant null is a present value whose type name is + * {@code NULL}. Instances are query-local and not thread-safe. + */ +public class VariantTypeOfTransformFunction extends BaseVariantTransformFunction { + public static final String FUNCTION_NAME = "variantTypeOf"; + + @Override + public String getName() { + return FUNCTION_NAME; + } + + @Override + public void init(List arguments, Map columnContextMap, + boolean nullHandlingEnabled) { + super.init(arguments, columnContextMap, nullHandlingEnabled); + initVariantArguments(arguments, 1, 2, false); + } + + @Override + public TransformResultMetadata getResultMetadata() { + return STRING_SV_NO_DICTIONARY_METADATA; + } + + @Override + public String[] transformToStringValuesSV(ValueBlock valueBlock) { + ensureEvaluated(valueBlock); + return _stringValuesSV; + } + + @Override + protected void initResultValues(int numDocs) { + initStringValuesSV(numDocs); + } + + @Override + protected boolean evaluateVariant(@Nullable byte[] variant, int index) { + String typeName = VariantUtils.variantTypeOf(variant, getVariantPath(), getReusableResult()); + if (typeName != null) { + _stringValuesSV[index] = typeName; + return true; + } + return false; + } + + @Override + protected void setNullValue(int index) { + _stringValuesSV[index] = NullValuePlaceHolder.STRING; + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/DistinctPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/DistinctPlanNode.java index 5abec423520b..3fa6ec55a0eb 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/plan/DistinctPlanNode.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/DistinctPlanNode.java @@ -37,6 +37,7 @@ import org.apache.pinot.segment.spi.index.reader.NullValueVectorReader; import org.apache.pinot.segment.spi.index.reader.SortedIndexReader; import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.data.FieldSpec.DataType; /// Execution plan for distinct queries on a single segment. @@ -54,6 +55,18 @@ public DistinctPlanNode(SegmentContext segmentContext, QueryContext queryContext @Override public Operator run() { List expressions = _queryContext.getSelectExpressions(); + for (ExpressionContext expression : expressions) { + String column = expression.getIdentifier(); + if (column != null) { + DataType dataType = _indexSegment.getDataSource(column, _queryContext.getSchema()) + .getDataSourceMetadata().getDataType(); + if (dataType.supportsEquality() && dataType.supportsHashing()) { + continue; + } + throw new IllegalArgumentException( + "Raw VARIANT values do not support DISTINCT; extract a typed path with variantGet first"); + } + } // Use dictionary to solve the query if possible if (_queryContext.getFilter() == null && expressions.size() == 1) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 527b7817621c..ee61ce12dc6b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -55,6 +55,7 @@ import org.apache.pinot.core.common.ObjectSerDeUtils; import org.apache.pinot.core.common.datatable.DataTableBuilder; import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.core.operator.filter.BaseFilterOperator; import org.apache.pinot.core.operator.filter.CombinedFilterOperator; @@ -149,6 +150,51 @@ public static F mergeFinalResult(AggregationFunction[] aggregationFunctions, + BaseProjectOperator projectOperator) { + for (AggregationFunction aggregationFunction : aggregationFunctions) { + if (aggregationFunction.getType() == AggregationFunctionType.COUNT) { + continue; + } + for (ExpressionContext inputExpression : aggregationFunction.getInputExpressions()) { + ColumnContext columnContext = projectOperator.getResultColumnContext(inputExpression); + if (columnContext != null && !columnContext.getDataType().supportsDirectAggregation()) { + throw rawVariantUnsupported(aggregationFunction); + } + } + } + } + + private static void validateRawVariantIdentifierInputs(AggregationFunction[] aggregationFunctions, + SegmentContext segmentContext, QueryContext queryContext) { + for (AggregationFunction aggregationFunction : aggregationFunctions) { + if (aggregationFunction.getType() == AggregationFunctionType.COUNT) { + continue; + } + for (ExpressionContext inputExpression : aggregationFunction.getInputExpressions()) { + if (inputExpression.getType() == ExpressionContext.Type.IDENTIFIER) { + FieldSpec.DataType dataType = + segmentContext.getIndexSegment().getDataSource(inputExpression.getIdentifier(), queryContext.getSchema()) + .getDataSourceMetadata().getDataType(); + if (!dataType.supportsDirectAggregation()) { + throw rawVariantUnsupported(aggregationFunction); + } + } + } + } + } + + private static IllegalArgumentException rawVariantUnsupported(AggregationFunction aggregationFunction) { + return new IllegalArgumentException( + "Aggregation function " + aggregationFunction.getType().getName() + + " does not support raw VARIANT values; extract a typed path with variantGet first"); + } + /// Creates a map from expression required by the [AggregationFunction] to [BlockValSet] fetched from the /// [ValueBlock]. public static Map getBlockValSetMap(AggregationFunction aggregationFunction, @@ -349,6 +395,9 @@ public AggregationInfo(AggregationFunction[] functions, BaseProjectOperator p _functions = functions; _projectOperator = projectOperator; _useStarTree = useStarTree; + if (!useStarTree) { + validateRawVariantAggregationInputs(functions, projectOperator); + } } public AggregationFunction[] getFunctions() { @@ -382,6 +431,10 @@ public static AggregationInfo buildAggregationInfo(SegmentContext segmentContext public static AggregationInfo buildAggregationInfoWithStarTree(SegmentContext segmentContext, QueryContext queryContext, AggregationFunction[] aggregationFunctions, @Nullable FilterContext filter, BaseFilterOperator filterOperator, List> predicateEvaluators) { + // Star-tree project operators expose pre-aggregated columns instead of the original logical operands, so validate + // direct identifiers before attempting the star-tree path. Function and literal operands cannot use star-tree and + // are validated against the regular project operator after fallback. + validateRawVariantIdentifierInputs(aggregationFunctions, segmentContext, queryContext); /// Star-tree stores pre-aggregated values per group key and cannot expand a row across multiple grouping /// sets, so it cannot serve GROUP BY GROUPING SETS / ROLLUP / CUBE queries. Fall back to the regular path. if (queryContext.isGroupingSets()) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java index a649fe6d7a17..87da2f8877f1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java @@ -324,10 +324,15 @@ private byte[] deserializeVariableBytes(ByteBuffer buffer) { } private void ensureResultType(BlockValSet bvs) { + FieldSpec.DataType valueType = bvs.getValueType(); + if (!valueType.supportsDirectAggregation()) { + throw new IllegalArgumentException( + "ANY_VALUE does not support raw VARIANT values; extract a typed path with variantGet first"); + } if (_resultType != null) { return; } - switch (bvs.getValueType().getStoredType()) { + switch (valueType.getStoredType()) { case INT: _resultType = ColumnDataType.INT; return; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/DistinctExecutorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/DistinctExecutorFactory.java index 6d8122a72f74..ca0e06ece9ec 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/DistinctExecutorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/DistinctExecutorFactory.java @@ -60,6 +60,10 @@ public static DistinctExecutor getDistinctExecutor(BaseProjectOperator projec ExpressionContext expression = expressions.get(0); ColumnContext columnContext = projectOperator.getResultColumnContext(expression); DataType dataType = columnContext.getDataType(); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw new IllegalArgumentException( + "Raw VARIANT values do not support DISTINCT; extract a typed path with variantGet first"); + } OrderByExpressionContext orderByExpression; if (orderByExpressions != null) { assert orderByExpressions.size() == 1; @@ -107,6 +111,10 @@ public static DistinctExecutor getDistinctExecutor(BaseProjectOperator projec for (int i = 0; i < numExpressions; i++) { ExpressionContext expression = expressions.get(i); ColumnContext columnContext = projectOperator.getResultColumnContext(expression); + if (!columnContext.getDataType().supportsEquality() || !columnContext.getDataType().supportsHashing()) { + throw new IllegalArgumentException( + "Raw VARIANT values do not support DISTINCT; extract a typed path with variantGet first"); + } if (!columnContext.isSingleValue()) { hasMVExpression = true; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java index 6c044df5606e..a11c206751ea 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java @@ -51,6 +51,10 @@ public static Comparator getComparator(List throw new BadQueryRequestException("MV expression: " + orderByExpressions.get(i) + " should not be included in the ORDER-BY clause"); } + if (!orderByColumnContexts[i].getDataType().supportsOrdering()) { + throw new BadQueryRequestException( + "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first"); + } } return getComparator(orderByExpressions, nullHandlingEnabled, from, to); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/common/datablock/DataBlockBuilderTest.java b/pinot-core/src/test/java/org/apache/pinot/core/common/datablock/DataBlockBuilderTest.java index 24e4c03b11cf..cad781c5413d 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/common/datablock/DataBlockBuilderTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/common/datablock/DataBlockBuilderTest.java @@ -30,6 +30,7 @@ import org.apache.pinot.common.datablock.DataBlock; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.core.query.aggregation.function.AggregationFunction; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.UuidUtils; @@ -109,7 +110,7 @@ private List generateRows(ColumnDataType type, int numRows) { break; case BYTES: for (int i = 0; i < numRows; i++) { - result.add(new Object[]{new ByteArray(String.valueOf(r.nextInt()).getBytes())}); + result.add(new Object[]{bytesValue(type, r)}); } break; case BIG_DECIMAL: @@ -262,7 +263,7 @@ Object[] generateColumns(ColumnDataType type, int numRows) { break; case BYTES: for (int i = 0; i < numRows; i++) { - result[i] = new ByteArray(String.valueOf(r.nextInt()).getBytes()); + result[i] = bytesValue(type, r); } break; case MAP: @@ -477,4 +478,11 @@ private void checkEquals(ColumnDataType type, DataBlock block, IntFunction PredicateEvaluatorProvider.getPredicateEvaluator(predicate, null, DataType.VARIANT, null)); + String message = exception.getCause().getMessage(); + assertTrue(message.contains("Raw VARIANT values do not support")); + assertTrue(message.contains("extract a typed path with variantGet first")); + } + } + /// RAW forward index with dictionary but no inverted/range/sorted index — the planner must drop the dictionary so /// that the scan-based filter operator receives a raw-value evaluator. Otherwise, the scan iterator would call /// `applySV(rawValue)` on a dict-based evaluator and throw [UnsupportedOperationException]. diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunctionTest.java new file mode 100644 index 000000000000..8fe0a00fe804 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunctionTest.java @@ -0,0 +1,162 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.LiteralContext; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.roaringbitmap.RoaringBitmap; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; + + +public class IsVariantNullTransformFunctionTest { + @Test + public void testSqlNullReturnsNonNullFalse() { + byte[] encodedVariantNull = VariantUtils.parseJsonToVariant("null"); + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{encodedVariantNull}, RoaringBitmap.bitmapOf(0)); + IsVariantNullTransformFunction function = new IsVariantNullTransformFunction(); + function.init(List.of(input), Map.of(), true); + ValueBlock block = valueBlock(1); + + assertEquals(function.getResultMetadata().getDataType(), DataType.BOOLEAN); + assertEquals(function.transformToIntValuesSV(block)[0], 0, + "The SQL-null bitmap must override any physical placeholder bytes"); + assertNull(function.getNullBitmap(block), "is_variant_null(SQL NULL) must be a non-null false"); + } + + @Test + public void testLiteralSqlNullReturnsNonNullFalse() { + IsVariantNullTransformFunction function = new IsVariantNullTransformFunction(); + function.init( + List.of(new LiteralTransformFunction(new LiteralContext(DataType.UNKNOWN, null))), Map.of(), true); + ValueBlock block = valueBlock(2); + + assertEquals(function.transformToIntValuesSV(block), new int[]{0, 0}); + assertNull(function.getNullBitmap(block)); + } + + @Test + public void testRootAndPathVariantNullSemantics() { + byte[] encodedVariantNull = VariantUtils.parseJsonToVariant("null"); + byte[] objectWithNull = VariantUtils.parseJsonToVariant("{\"coupon\":null}"); + byte[] objectWithMissingPath = VariantUtils.parseJsonToVariant("{\"eventType\":\"checkout\"}"); + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{encodedVariantNull, objectWithNull, objectWithMissingPath}, null); + ValueBlock block = valueBlock(3); + + IsVariantNullTransformFunction rootFunction = new IsVariantNullTransformFunction(); + rootFunction.init(List.of(input), Map.of(), true); + int[] rootValues = rootFunction.transformToIntValuesSV(block); + assertEquals(rootValues, new int[]{1, 0, 0}); + assertSame(rootFunction.transformToIntValuesSV(block), rootValues); + assertNull(rootFunction.getNullBitmap(block)); + assertEquals(input._valueCallCount, 1); + assertEquals(input._nullCallCount, 1); + + IsVariantNullTransformFunction pathFunction = new IsVariantNullTransformFunction(); + pathFunction.init(List.of(input, stringLiteral("$.coupon")), Map.of(), true); + int[] pathValues = pathFunction.transformToIntValuesSV(block); + assertEquals(pathValues, new int[]{0, 1, 0}); + assertSame(pathFunction.transformToIntValuesSV(block), pathValues); + assertNull(pathFunction.getNullBitmap(block)); + assertEquals(input._valueCallCount, 2); + assertEquals(input._nullCallCount, 2); + } + + @Test + public void testArgumentValidationAndFactoryRegistration() { + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{VariantUtils.parseJsonToVariant("true")}, null); + IsVariantNullTransformFunction function = new IsVariantNullTransformFunction(); + + assertThrows(IllegalArgumentException.class, + () -> function.init(List.of(input, mock(TransformFunction.class)), Map.of(), true)); + assertThrows(IllegalArgumentException.class, + () -> function.init(List.of(input, new LiteralTransformFunction(new LiteralContext(DataType.INT, 1))), + Map.of(), true)); + assertThrows("An invalid path must fail during initialization, before any row is evaluated", + IllegalArgumentException.class, + () -> function.init(List.of(input, stringLiteral("coupon")), Map.of(), true)); + + Map> functions = TransformFunctionFactory.getAllFunctions(); + assertSame(functions.get(TransformFunctionFactory.canonicalize("is_variant_null")), + IsVariantNullTransformFunction.class); + } + + private static LiteralTransformFunction stringLiteral(String value) { + return new LiteralTransformFunction(new LiteralContext(DataType.STRING, value)); + } + + private static ValueBlock valueBlock(int numDocs) { + ValueBlock valueBlock = mock(ValueBlock.class); + when(valueBlock.getNumDocs()).thenReturn(numDocs); + return valueBlock; + } + + private static final class BytesTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.VARIANT, true, false); + private final byte[][] _values; + @Nullable + private final RoaringBitmap _nullBitmap; + private int _valueCallCount; + private int _nullCallCount; + + private BytesTransformFunction(byte[][] values, @Nullable RoaringBitmap nullBitmap) { + _values = values; + _nullBitmap = nullBitmap; + } + + @Override + public String getName() { + return "variantInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + _valueCallCount++; + return _values; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + _nullCallCount++; + return _nullBitmap; + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunctionTest.java new file mode 100644 index 000000000000..f589fa989dfe --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunctionTest.java @@ -0,0 +1,154 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.LiteralContext; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.roaringbitmap.RoaringBitmap; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; + + +public class VariantExistsTransformFunctionTest { + @Test + public void testVectorizedPresenceAndNullSemanticsAreCachedPerBlock() { + byte[] present = VariantUtils.parseJsonToVariant("{\"payload\":{\"name\":\"alice\",\"coupon\":null}}"); + byte[] missing = VariantUtils.parseJsonToVariant("{\"payload\":{}}"); + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{present, present, missing, new byte[0]}, RoaringBitmap.bitmapOf(3)); + VariantExistsTransformFunction function = new VariantExistsTransformFunction(); + function.init(List.of(input, stringLiteral("$.payload.coupon")), Map.of(), true); + ValueBlock block = valueBlock(4); + + assertEquals(function.getResultMetadata().getDataType(), DataType.BOOLEAN); + assertEquals(function.transformToIntValuesSV(block), new int[]{1, 1, 0, 0}, + "A present Variant null counts as present, while a missing path does not"); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(3)); + assertSame(function.transformToIntValuesSV(block), function.transformToIntValuesSV(block)); + assertEquals(input._valueCallCount, 1); + assertEquals(input._nullCallCount, 1); + + function.getNullBitmap(valueBlock(4)); + assertEquals(input._valueCallCount, 2); + assertEquals(input._nullCallCount, 2); + } + + @Test + public void testPhysicalSqlNullWithoutBitmapBecomesNull() { + BytesTransformFunction input = new BytesTransformFunction(new byte[][]{new byte[0]}, null); + VariantExistsTransformFunction function = new VariantExistsTransformFunction(); + function.init(List.of(input, stringLiteral("$")), Map.of(), true); + ValueBlock block = valueBlock(1); + + assertEquals(function.transformToIntValuesSV(block), new int[]{0}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + } + + @Test + public void testLiteralSqlNullRemainsSqlNull() { + VariantExistsTransformFunction function = new VariantExistsTransformFunction(); + function.init( + List.of(new LiteralTransformFunction(new LiteralContext(DataType.UNKNOWN, null)), stringLiteral("$")), + Map.of(), true); + ValueBlock block = valueBlock(2); + + assertEquals(function.transformToIntValuesSV(block), new int[]{0, 0}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0, 1)); + } + + @Test + public void testArgumentValidationCompilesPathAndFactoryRegistration() { + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{VariantUtils.parseJsonToVariant("true")}, null); + VariantExistsTransformFunction function = new VariantExistsTransformFunction(); + + assertThrows(IllegalArgumentException.class, () -> function.init(List.of(input), Map.of(), true)); + assertThrows(IllegalArgumentException.class, + () -> function.init(List.of(input, mock(TransformFunction.class)), Map.of(), true)); + assertThrows(IllegalArgumentException.class, + () -> function.init(List.of(input, new LiteralTransformFunction(new LiteralContext(DataType.INT, 1))), + Map.of(), true)); + assertThrows("An invalid path must fail during initialization, before any row is evaluated", + IllegalArgumentException.class, + () -> function.init(List.of(input, stringLiteral("payload.name")), Map.of(), true)); + + Map> functions = TransformFunctionFactory.getAllFunctions(); + assertSame(functions.get(TransformFunctionFactory.canonicalize("variant_exists")), + VariantExistsTransformFunction.class); + } + + private static LiteralTransformFunction stringLiteral(String value) { + return new LiteralTransformFunction(new LiteralContext(DataType.STRING, value)); + } + + private static ValueBlock valueBlock(int numDocs) { + ValueBlock valueBlock = mock(ValueBlock.class); + when(valueBlock.getNumDocs()).thenReturn(numDocs); + return valueBlock; + } + + private static final class BytesTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.VARIANT, true, false); + private final byte[][] _values; + @Nullable + private final RoaringBitmap _nullBitmap; + private int _valueCallCount; + private int _nullCallCount; + + private BytesTransformFunction(byte[][] values, @Nullable RoaringBitmap nullBitmap) { + _values = values; + _nullBitmap = nullBitmap; + } + + @Override + public String getName() { + return "variantInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + _valueCallCount++; + return _values; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + _nullCallCount++; + return _nullBitmap; + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java new file mode 100644 index 000000000000..fcdc49f3cf9a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java @@ -0,0 +1,519 @@ +/** + * 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.operator.transform.function; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.function.Consumer; +import javax.annotation.Nullable; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.LiteralContext; +import org.apache.pinot.common.request.context.RequestContextUtils; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.exception.BadQueryRequestException; +import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.VariantEnvelope; +import org.roaringbitmap.RoaringBitmap; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +public class VariantGetTransformFunctionTest { + @Test + public void testLiteralSqlNullRemainsSqlNull() { + VariantGetTransformFunction function = new VariantGetTransformFunction(); + function.init( + List.of(new LiteralTransformFunction(new LiteralContext(DataType.UNKNOWN, null)), stringLiteral("$")), + Map.of(), true); + ValueBlock block = valueBlock(2); + + assertEquals(function.transformToBytesValuesSV(block), new byte[][]{new byte[0], new byte[0]}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0, 1)); + } + + @Test + public void testMetadataAndTypedExtraction() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\",\"score\":12.5}"); + ValueBlock block = valueBlock(1); + BytesTransformFunction input = new BytesTransformFunction(new byte[][]{variant}, null); + + VariantGetTransformFunction stringFunction = new VariantGetTransformFunction(); + stringFunction.init(arguments(input, "$.eventType", "STRING"), Map.of(), true); + assertEquals(stringFunction.getResultMetadata().getDataType(), DataType.STRING); + assertEquals(stringFunction.transformToStringValuesSV(block)[0], "click"); + assertNull(stringFunction.getNullBitmap(block)); + + VariantGetTransformFunction doubleFunction = new VariantGetTransformFunction(); + doubleFunction.init(arguments(input, "$.score", "DOUBLE"), Map.of(), true); + assertEquals(doubleFunction.getResultMetadata().getDataType(), DataType.DOUBLE); + assertEquals(doubleFunction.transformToDoubleValuesSV(block)[0], 12.5); + } + + @Test + public void testEverySupportedTargetType() { + ValueBlock block = valueBlock(1); + + VariantGetTransformFunction function = typedFunction(variant(builder -> builder.appendBoolean(true)), "BOOLEAN"); + assertEquals(function.getResultMetadata().getDataType(), DataType.BOOLEAN); + assertEquals(function.transformToIntValuesSV(block)[0], 1); + + function = typedFunction(variant(builder -> builder.appendInt(42)), "INT"); + assertEquals(function.getResultMetadata().getDataType(), DataType.INT); + assertEquals(function.transformToIntValuesSV(block)[0], 42); + + function = typedFunction(variant(builder -> builder.appendLong(4_294_967_296L)), "LONG"); + assertEquals(function.getResultMetadata().getDataType(), DataType.LONG); + assertEquals(function.transformToLongValuesSV(block)[0], 4_294_967_296L); + + function = typedFunction(variant(builder -> builder.appendFloat(1.25F)), "FLOAT"); + assertEquals(function.getResultMetadata().getDataType(), DataType.FLOAT); + assertEquals(function.transformToFloatValuesSV(block)[0], 1.25F); + + function = typedFunction(variant(builder -> builder.appendDouble(12.5)), "DOUBLE"); + assertEquals(function.getResultMetadata().getDataType(), DataType.DOUBLE); + assertEquals(function.transformToDoubleValuesSV(block)[0], 12.5); + + BigDecimal decimal = new BigDecimal("1234567890.12345"); + function = typedFunction(variant(builder -> builder.appendDecimal(decimal)), "BIG_DECIMAL"); + assertEquals(function.getResultMetadata().getDataType(), DataType.BIG_DECIMAL); + assertEquals(function.transformToBigDecimalValuesSV(block)[0], decimal); + + function = typedFunction(variant(builder -> builder.appendString("click")), "STRING"); + assertEquals(function.getResultMetadata().getDataType(), DataType.STRING); + assertEquals(function.transformToStringValuesSV(block)[0], "click"); + + byte[] binary = new byte[]{0, 1, (byte) 0xFF}; + function = typedFunction(variant(builder -> builder.appendBinary(ByteBuffer.wrap(binary))), "BYTES"); + assertEquals(function.getResultMetadata().getDataType(), DataType.BYTES); + assertEquals(function.transformToBytesValuesSV(block)[0], binary); + + UUID uuid = UUID.fromString("12345678-1234-5678-9abc-def012345678"); + function = typedFunction(variant(builder -> builder.appendUUID(uuid)), "UUID"); + assertEquals(function.getResultMetadata().getDataType(), DataType.UUID); + byte[] uuidBytes = function.transformToBytesValuesSV(block)[0]; + assertTrue(UuidUtils.equals(uuidBytes, UuidUtils.toBytes(uuid)), + "UUID extraction must directly expose the copied 16-byte value"); + assertEquals(UuidUtils.fromBytes(uuidBytes), uuid); + + long timestampMicros = 1_700_000_000_123_000L; + function = typedFunction(variant(builder -> builder.appendTimestampTz(timestampMicros)), "TIMESTAMP"); + assertEquals(function.getResultMetadata().getDataType(), DataType.TIMESTAMP); + assertEquals(function.transformToLongValuesSV(block)[0], 1_700_000_000_123L); + + function = typedFunction(variant(builder -> builder.appendString("nested")), "VARIANT"); + assertEquals(function.getResultMetadata().getDataType(), DataType.VARIANT); + assertEquals(VariantUtils.variantToJson(function.transformToBytesValuesSV(block)[0]), "\"nested\""); + + function = typedFunction(variant(builder -> builder.appendString("json")), "JSON"); + assertEquals(function.getResultMetadata().getDataType(), DataType.JSON); + assertEquals(function.transformToStringValuesSV(block)[0], "\"json\""); + + assertThrows(IllegalArgumentException.class, + () -> typedFunction(variant(builder -> builder.appendInt(1)), "UNSUPPORTED")); + } + + @Test + public void testStrictMissingPathReturnsNullAndCastFailureThrows() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\",\"score\":\"not-a-number\"}"); + VariantGetTransformFunction function = new VariantGetTransformFunction(); + function.init(arguments(new BytesTransformFunction(new byte[][]{variant}, null), "$.missing", "STRING"), Map.of(), + true); + + ValueBlock block = valueBlock(1); + assertEquals(function.transformToStringValuesSV(block), new String[]{""}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + VariantGetTransformFunction castFunction = new VariantGetTransformFunction(); + castFunction.init( + arguments(new BytesTransformFunction(new byte[][]{variant}, null), "$.score", "DOUBLE"), Map.of(), true); + assertThrows(IllegalArgumentException.class, () -> castFunction.transformToDoubleValuesSV(block)); + } + + @Test + public void testDefaultVariantTarget() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"payload\":{\"count\":7},\"variantNull\":null}"); + ValueBlock block = valueBlock(1); + VariantGetTransformFunction function = new VariantGetTransformFunction(); + function.init(arguments(new BytesTransformFunction(new byte[][]{variant}, null), "$.payload"), Map.of(), true); + + assertEquals(function.getResultMetadata().getDataType(), DataType.VARIANT); + assertEquals(VariantUtils.variantToJson(function.transformToBytesValuesSV(block)[0]), "{\"count\":7}"); + + VariantGetTransformFunction variantNullFunction = new VariantGetTransformFunction(); + variantNullFunction.init( + arguments(new BytesTransformFunction(new byte[][]{variant}, null), "$.variantNull"), Map.of(), true); + assertTrue(VariantUtils.isVariantNull(variantNullFunction.transformToBytesValuesSV(block)[0])); + assertNull(variantNullFunction.getNullBitmap(block)); + } + + @Test + public void testTryVariantGetNullBitmap() { + byte[] populated = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\"}"); + byte[] variantNull = VariantUtils.parseJsonToVariant("null"); + byte[] missing = VariantUtils.parseJsonToVariant("{\"other\":\"value\"}"); + RoaringBitmap inputNulls = RoaringBitmap.bitmapOf(3); + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{populated, variantNull, missing, new byte[0]}, inputNulls); + + VariantGetTransformFunction.Try function = new VariantGetTransformFunction.Try(); + function.init(arguments(input, "$.eventType", "STRING"), Map.of(), true); + ValueBlock block = valueBlock(4); + + assertEquals(function.getResultMetadata().getDataType(), DataType.STRING); + assertEquals(function.transformToStringValuesSV(block), new String[]{"click", "", "", ""}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(1, 2, 3)); + } + + @Test + public void testVariantExtractionIsCachedPerBlock() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\"}"); + BytesTransformFunction input = new BytesTransformFunction(new byte[][]{variant}, null); + VariantGetTransformFunction function = new VariantGetTransformFunction(); + function.init(arguments(input, "$.eventType", "STRING"), Map.of(), true); + + ValueBlock block = valueBlock(1); + assertNull(function.getNullBitmap(block)); + String[] values = function.transformToStringValuesSV(block); + assertEquals(values[0], "click"); + assertNull(function.getNullBitmap(block)); + assertSame(function.transformToStringValuesSV(block), values); + assertEquals(input._transformCalls, 1); + assertEquals(input._nullBitmapCalls, 1); + + assertSame(function.transformToStringValuesSV(valueBlock(1)), values); + assertEquals(input._transformCalls, 2); + assertEquals(input._nullBitmapCalls, 2); + } + + @Test + public void testJsonParsingIsCachedPerBlock() { + StringTransformFunction input = + new StringTransformFunction(new String[]{"{\"answer\":42}", "null", "{not-json", ""}, + RoaringBitmap.bitmapOf(3)); + ParseJsonToVariantTransformFunction.Try function = new ParseJsonToVariantTransformFunction.Try(); + function.init(List.of(input), Map.of(), true); + + ValueBlock block = valueBlock(4); + byte[][] values = function.transformToBytesValuesSV(block); + assertEquals(VariantUtils.variantToJson(values[0]), "{\"answer\":42}"); + assertTrue(VariantUtils.isVariantNull(values[1])); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(2, 3)); + assertSame(function.transformToBytesValuesSV(block), values); + assertEquals(input._transformCalls, 1); + assertEquals(input._nullBitmapCalls, 1); + + ValueBlock nextBlock = valueBlock(4); + function.getNullBitmap(nextBlock); + assertSame(function.transformToBytesValuesSV(nextBlock), values); + assertEquals(input._transformCalls, 2); + assertEquals(input._nullBitmapCalls, 2); + } + + @Test + public void testLiteralJsonIsParsedAtInitializationAndReused() { + ParseJsonToVariantTransformFunction strict = new ParseJsonToVariantTransformFunction(); + strict.init(List.of(stringLiteral("{\"answer\":42}")), Map.of(), true); + + byte[][] values = strict.transformToBytesValuesSV(valueBlock(3)); + assertEquals(VariantUtils.variantToJson(values[0]), "{\"answer\":42}"); + assertSame(values[1], values[0]); + assertSame(values[2], values[0]); + assertNull(strict.getNullBitmap(valueBlock(3))); + assertSame(strict.transformToBytesValuesSV(valueBlock(1))[0], values[0], + "The parsed literal must be reused across blocks"); + + ParseJsonToVariantTransformFunction invalidStrict = new ParseJsonToVariantTransformFunction(); + assertThrows("Strict literal parsing must fail during initialization", IllegalArgumentException.class, + () -> invalidStrict.init(List.of(stringLiteral("{not-json")), Map.of(), true)); + + ParseJsonToVariantTransformFunction.Try invalidTolerant = new ParseJsonToVariantTransformFunction.Try(); + invalidTolerant.init(List.of(stringLiteral("{not-json")), Map.of(), true); + ValueBlock invalidBlock = valueBlock(2); + assertEquals(invalidTolerant.transformToBytesValuesSV(invalidBlock), new byte[][]{new byte[0], new byte[0]}); + assertEquals(invalidTolerant.getNullBitmap(invalidBlock), RoaringBitmap.bitmapOf(0, 1)); + + ParseJsonToVariantTransformFunction sqlNull = new ParseJsonToVariantTransformFunction(); + sqlNull.init(List.of(new LiteralTransformFunction(new LiteralContext(DataType.UNKNOWN, null))), Map.of(), true); + ValueBlock sqlNullBlock = valueBlock(2); + assertEquals(sqlNull.transformToBytesValuesSV(sqlNullBlock), new byte[][]{new byte[0], new byte[0]}); + assertEquals(sqlNull.getNullBitmap(sqlNullBlock), RoaringBitmap.bitmapOf(0, 1)); + + ParseJsonToVariantTransformFunction variantNull = new ParseJsonToVariantTransformFunction(); + variantNull.init(List.of(stringLiteral("null")), Map.of(), true); + ValueBlock variantNullBlock = valueBlock(1); + assertTrue(VariantUtils.isVariantNull(variantNull.transformToBytesValuesSV(variantNullBlock)[0])); + assertNull(variantNull.getNullBitmap(variantNullBlock), "JSON null must remain distinct from SQL null"); + } + + @Test + public void testStrictFailureInvalidatesPreviousBlockCache() { + ValueBlock validBlock = valueBlock(1); + ValueBlock invalidBlock = valueBlock(1); + byte[] validVariant = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\"}"); + BlockAwareBytesTransformFunction variantInput = new BlockAwareBytesTransformFunction( + Map.of(validBlock, new byte[][]{validVariant}, invalidBlock, new byte[][]{new byte[]{1}})); + VariantGetTransformFunction variantGet = new VariantGetTransformFunction(); + variantGet.init(arguments(variantInput, "$.eventType", "STRING"), Map.of(), true); + + assertEquals(variantGet.transformToStringValuesSV(validBlock)[0], "click"); + assertThrows(IllegalArgumentException.class, () -> variantGet.transformToStringValuesSV(invalidBlock)); + assertEquals(variantGet.transformToStringValuesSV(validBlock)[0], "click"); + assertEquals(variantInput._transformCalls, 3); + + BlockAwareStringTransformFunction jsonInput = new BlockAwareStringTransformFunction( + Map.of(validBlock, new String[]{"{\"answer\":42}"}, invalidBlock, new String[]{"{not-json"})); + ParseJsonToVariantTransformFunction parseJson = new ParseJsonToVariantTransformFunction(); + parseJson.init(List.of(jsonInput), Map.of(), true); + + assertEquals(VariantUtils.variantToJson(parseJson.transformToBytesValuesSV(validBlock)[0]), "{\"answer\":42}"); + assertThrows(IllegalArgumentException.class, () -> parseJson.transformToBytesValuesSV(invalidBlock)); + assertEquals(VariantUtils.variantToJson(parseJson.transformToBytesValuesSV(validBlock)[0]), "{\"answer\":42}"); + assertEquals(jsonInput._transformCalls, 3); + } + + @Test + public void testFactoryRegistrations() { + Map> functions = TransformFunctionFactory.getAllFunctions(); + assertSame(functions.get(TransformFunctionFactory.canonicalize("variant_get")), + VariantGetTransformFunction.class); + assertSame(functions.get(TransformFunctionFactory.canonicalize("try_variant_get")), + VariantGetTransformFunction.Try.class); + assertSame(functions.get(TransformFunctionFactory.canonicalize("parse_json")), + ParseJsonToVariantTransformFunction.class); + assertSame(functions.get(TransformFunctionFactory.canonicalize("try_parse_json")), + ParseJsonToVariantTransformFunction.Try.class); + assertSame(functions.get(TransformFunctionFactory.canonicalize("parseJsonToVariant")), + ParseJsonToVariantTransformFunction.class); + assertSame(functions.get(TransformFunctionFactory.canonicalize("tryParseJsonToVariant")), + ParseJsonToVariantTransformFunction.Try.class); + } + + @Test + public void testVariantFunctionsRequireQueryNullHandling() { + List expressions = List.of( + "variant_get(parse_json('{}'), '$.value')", + "try_variant_get(parse_json('{}'), '$.value')", + "variant_exists(parse_json('{}'), '$.value')", + "is_variant_null(parse_json('null'))", + "variant_type_of(parse_json('{}'))", + "variant_to_json(parse_json('{}'))", + "parse_json('{}')", + "parse_json_to_variant('{}')", + "try_parse_json('{}')", + "try_parse_json_to_variant('{}')"); + for (String expressionString : expressions) { + ExpressionContext expression = RequestContextUtils.getExpression(expressionString); + BadQueryRequestException exception = expectThrows(BadQueryRequestException.class, + () -> TransformFunctionFactory.get(expression, Map.of())); + assertTrue(exception.getMessage().contains("requires query null handling"), + "Unexpected rejection for " + expressionString + ": " + exception.getMessage()); + TransformFunctionFactory.getNullHandlingEnabled(expression, Map.of()); + } + } + + private static List arguments(TransformFunction input, String path) { + return List.of(input, stringLiteral(path)); + } + + private static List arguments(TransformFunction input, String path, String targetType) { + return List.of(input, stringLiteral(path), stringLiteral(targetType)); + } + + private static VariantGetTransformFunction typedFunction(byte[] variant, String targetType) { + VariantGetTransformFunction function = new VariantGetTransformFunction(); + function.init(arguments(new BytesTransformFunction(new byte[][]{variant}, null), "$", targetType), Map.of(), true); + return function; + } + + private static byte[] variant(Consumer writer) { + VariantBuilder builder = new VariantBuilder(); + writer.accept(builder); + Variant variant = builder.build(); + return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); + } + + private static LiteralTransformFunction stringLiteral(String value) { + return new LiteralTransformFunction(new LiteralContext(DataType.STRING, value)); + } + + private static ValueBlock valueBlock(int numDocs) { + ValueBlock valueBlock = mock(ValueBlock.class); + when(valueBlock.getNumDocs()).thenReturn(numDocs); + return valueBlock; + } + + private static final class BytesTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.VARIANT, true, false); + private final byte[][] _values; + @Nullable + private final RoaringBitmap _nullBitmap; + private int _transformCalls; + private int _nullBitmapCalls; + + private BytesTransformFunction(byte[][] values, @Nullable RoaringBitmap nullBitmap) { + _values = values; + _nullBitmap = nullBitmap; + } + + @Override + public String getName() { + return "variantInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + _transformCalls++; + return _values; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + _nullBitmapCalls++; + return _nullBitmap; + } + } + + private static final class StringTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.STRING, true, false); + private final String[] _values; + @Nullable + private final RoaringBitmap _nullBitmap; + private int _transformCalls; + private int _nullBitmapCalls; + + private StringTransformFunction(String[] values, @Nullable RoaringBitmap nullBitmap) { + _values = values; + _nullBitmap = nullBitmap; + } + + @Override + public String getName() { + return "stringInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public String[] transformToStringValuesSV(ValueBlock valueBlock) { + _transformCalls++; + return _values; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + _nullBitmapCalls++; + return _nullBitmap; + } + } + + private static final class BlockAwareBytesTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.VARIANT, true, false); + private final IdentityHashMap _valuesByBlock; + private int _transformCalls; + + private BlockAwareBytesTransformFunction(Map valuesByBlock) { + _valuesByBlock = new IdentityHashMap<>(valuesByBlock); + } + + @Override + public String getName() { + return "blockAwareVariantInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + _transformCalls++; + return _valuesByBlock.get(valueBlock); + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + return null; + } + } + + private static final class BlockAwareStringTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.STRING, true, false); + private final IdentityHashMap _valuesByBlock; + private int _transformCalls; + + private BlockAwareStringTransformFunction(Map valuesByBlock) { + _valuesByBlock = new IdentityHashMap<>(valuesByBlock); + } + + @Override + public String getName() { + return "blockAwareStringInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public String[] transformToStringValuesSV(ValueBlock valueBlock) { + _transformCalls++; + return _valuesByBlock.get(valueBlock); + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + return null; + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunctionTest.java new file mode 100644 index 000000000000..a84d821c6883 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunctionTest.java @@ -0,0 +1,164 @@ +/** + * 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.operator.transform.function; + +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.LiteralContext; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.transform.TransformResultMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.roaringbitmap.RoaringBitmap; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; + + +public class VariantTypeOfTransformFunctionTest { + @Test + public void testLiteralSqlNullRemainsSqlNull() { + VariantTypeOfTransformFunction function = new VariantTypeOfTransformFunction(); + function.init( + List.of(new LiteralTransformFunction(new LiteralContext(DataType.UNKNOWN, null))), Map.of(), true); + ValueBlock block = valueBlock(2); + + assertEquals(function.transformToStringValuesSV(block), new String[]{"", ""}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0, 1)); + } + + @Test + public void testRootTypePreservesSqlNullAndVariantNull() { + byte[] variantNull = VariantUtils.parseJsonToVariant("null"); + byte[] object = VariantUtils.parseJsonToVariant("{\"name\":\"alice\"}"); + RoaringBitmap inputNulls = RoaringBitmap.bitmapOf(0); + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{new byte[0], variantNull, object, new byte[0]}, inputNulls); + + VariantTypeOfTransformFunction function = new VariantTypeOfTransformFunction(); + function.init(List.of(input), Map.of(), true); + ValueBlock block = valueBlock(4); + + assertEquals(function.getResultMetadata().getDataType(), DataType.STRING); + String[] values = function.transformToStringValuesSV(block); + assertEquals(values, new String[]{"", "NULL", "OBJECT", ""}); + assertSame(function.transformToStringValuesSV(block), values); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0, 3)); + assertEquals(input._valueCallCount, 1); + assertEquals(input._nullCallCount, 1); + + function.getNullBitmap(valueBlock(4)); + assertEquals(input._valueCallCount, 2); + assertEquals(input._nullCallCount, 2); + } + + @Test + public void testLiteralPathTypeAndMissingPathNullBitmap() { + byte[] present = VariantUtils.parseJsonToVariant("{\"payload\":{\"name\":\"alice\"}}"); + byte[] missing = VariantUtils.parseJsonToVariant("{\"payload\":{}}"); + byte[] variantNull = VariantUtils.parseJsonToVariant("{\"payload\":{\"name\":null}}"); + RoaringBitmap inputNulls = RoaringBitmap.bitmapOf(3); + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{present, missing, variantNull, new byte[0]}, inputNulls); + + VariantTypeOfTransformFunction function = new VariantTypeOfTransformFunction(); + function.init(List.of(input, stringLiteral("$.payload.name")), Map.of(), true); + ValueBlock block = valueBlock(4); + + String[] values = function.transformToStringValuesSV(block); + assertEquals(values, new String[]{"STRING", "", "NULL", ""}); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(1, 3)); + assertSame(function.transformToStringValuesSV(block), values); + assertEquals(input._valueCallCount, 1); + assertEquals(input._nullCallCount, 1); + } + + @Test + public void testArgumentValidationAndFactoryRegistration() { + BytesTransformFunction input = + new BytesTransformFunction(new byte[][]{VariantUtils.parseJsonToVariant("true")}, null); + VariantTypeOfTransformFunction function = new VariantTypeOfTransformFunction(); + + assertThrows(IllegalArgumentException.class, + () -> function.init(List.of(input, mock(TransformFunction.class)), Map.of(), true)); + assertThrows(IllegalArgumentException.class, + () -> function.init(List.of(input, new LiteralTransformFunction(new LiteralContext(DataType.INT, 1))), + Map.of(), true)); + assertThrows("An invalid path must fail during initialization, before any row is evaluated", + IllegalArgumentException.class, + () -> function.init(List.of(input, stringLiteral("payload.name")), Map.of(), true)); + + Map> functions = TransformFunctionFactory.getAllFunctions(); + assertSame(functions.get(TransformFunctionFactory.canonicalize("variant_type_of")), + VariantTypeOfTransformFunction.class); + } + + private static LiteralTransformFunction stringLiteral(String value) { + return new LiteralTransformFunction(new LiteralContext(DataType.STRING, value)); + } + + private static ValueBlock valueBlock(int numDocs) { + ValueBlock valueBlock = mock(ValueBlock.class); + when(valueBlock.getNumDocs()).thenReturn(numDocs); + return valueBlock; + } + + private static final class BytesTransformFunction extends BaseTransformFunction { + private static final TransformResultMetadata RESULT_METADATA = + new TransformResultMetadata(DataType.VARIANT, true, false); + private final byte[][] _values; + @Nullable + private final RoaringBitmap _nullBitmap; + private int _valueCallCount; + private int _nullCallCount; + + private BytesTransformFunction(byte[][] values, @Nullable RoaringBitmap nullBitmap) { + _values = values; + _nullBitmap = nullBitmap; + } + + @Override + public String getName() { + return "variantInput"; + } + + @Override + public TransformResultMetadata getResultMetadata() { + return RESULT_METADATA; + } + + @Override + public byte[][] transformToBytesValuesSV(ValueBlock valueBlock) { + _valueCallCount++; + return _values; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap(ValueBlock valueBlock) { + _nullCallCount++; + return _nullBitmap; + } + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/plan/DistinctPlanNodeTest.java b/pinot-core/src/test/java/org/apache/pinot/core/plan/DistinctPlanNodeTest.java new file mode 100644 index 000000000000..0671adb31c15 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/plan/DistinctPlanNodeTest.java @@ -0,0 +1,58 @@ +/** + * 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.plan; + +import java.util.List; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.mockito.Mockito; +import org.testng.Assert; +import org.testng.annotations.Test; + + +public class DistinctPlanNodeTest { + @Test + public void rawVariantIsRejectedBeforeDictionaryOptimizations() { + SegmentContext segmentContext = Mockito.mock(SegmentContext.class); + IndexSegment indexSegment = Mockito.mock(IndexSegment.class); + QueryContext queryContext = Mockito.mock(QueryContext.class); + DataSource dataSource = Mockito.mock(DataSource.class); + DataSourceMetadata metadata = Mockito.mock(DataSourceMetadata.class); + Schema schema = new Schema(); + + Mockito.when(segmentContext.getIndexSegment()).thenReturn(indexSegment); + Mockito.when(queryContext.getSelectExpressions()) + .thenReturn(List.of(ExpressionContext.forIdentifier("payload"))); + Mockito.when(queryContext.getSchema()).thenReturn(schema); + Mockito.when(indexSegment.getDataSource("payload", schema)).thenReturn(dataSource); + Mockito.when(dataSource.getDataSourceMetadata()).thenReturn(metadata); + Mockito.when(metadata.getDataType()).thenReturn(DataType.VARIANT); + + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> new DistinctPlanNode(segmentContext, queryContext).run()); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support DISTINCT")); + Assert.assertTrue(exception.getMessage().contains("extract a typed path with variantGet first")); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java index 29041933af2f..fe62e332aeb3 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java @@ -18,9 +18,15 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.List; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.ColumnContext; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.datasource.DataSource; import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec; +import org.testng.Assert; import org.testng.annotations.Test; import static org.mockito.Mockito.mock; @@ -33,6 +39,7 @@ /// result resolver used by the non-scan based and partial metadata based aggregation paths. @SuppressWarnings("rawtypes") public class AggregationFunctionUtilsTest { + private static final ExpressionContext PAYLOAD = ExpressionContext.forIdentifier("payload"); private static AggregationFunction mockFunction(AggregationFunctionType type) { AggregationFunction aggregationFunction = mock(AggregationFunction.class); @@ -81,4 +88,50 @@ public void testNonCountWithNullDataSourceThrows() { () -> AggregationFunctionUtils.getAggregationResult(mockFunction(AggregationFunctionType.MIN), null, 100, "TEST")); } + + @Test + public void testRejectsUnsafeRawVariantAggregationsUsingLogicalType() { + Assert.assertEquals(FieldSpec.DataType.VARIANT.getStoredType(), FieldSpec.DataType.BYTES); + BaseProjectOperator projectOperator = projectOperator(FieldSpec.DataType.VARIANT); + + for (AggregationFunctionType functionType + : List.of(AggregationFunctionType.SUM, AggregationFunctionType.ANYVALUE, + AggregationFunctionType.DISTINCTCOUNTHLL)) { + AggregationFunction aggregationFunction = aggregationFunction(functionType); + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> AggregationFunctionUtils.validateRawVariantAggregationInputs( + new AggregationFunction[]{aggregationFunction}, projectOperator)); + Assert.assertTrue(exception.getMessage().contains(functionType.getName())); + Assert.assertTrue(exception.getMessage().contains("variantGet")); + } + } + + @Test + public void testAllowsCountOfRawVariant() { + AggregationFunctionUtils.validateRawVariantAggregationInputs( + new AggregationFunction[]{aggregationFunction(AggregationFunctionType.COUNT)}, + projectOperator(FieldSpec.DataType.VARIANT)); + } + + @Test + public void testAllowsAggregationOfTypedVariantExtraction() { + AggregationFunctionUtils.validateRawVariantAggregationInputs( + new AggregationFunction[]{aggregationFunction(AggregationFunctionType.MINSTRING)}, + projectOperator(FieldSpec.DataType.STRING)); + } + + private static AggregationFunction aggregationFunction(AggregationFunctionType functionType) { + AggregationFunction aggregationFunction = mock(AggregationFunction.class); + when(aggregationFunction.getType()).thenReturn(functionType); + when(aggregationFunction.getInputExpressions()).thenReturn(List.of(PAYLOAD)); + return aggregationFunction; + } + + private static BaseProjectOperator projectOperator(FieldSpec.DataType dataType) { + BaseProjectOperator projectOperator = mock(BaseProjectOperator.class); + ColumnContext columnContext = mock(ColumnContext.class); + when(columnContext.getDataType()).thenReturn(dataType); + when(projectOperator.getResultColumnContext(PAYLOAD)).thenReturn(columnContext); + return projectOperator; + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java index f0bf5dd01336..eb0b18163611 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunctionTest.java @@ -18,11 +18,20 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.BlockValSet; import org.apache.pinot.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; +import org.testng.Assert; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + + public class AnyValueAggregationFunctionTest extends AbstractAggregationFunctionTest { // Constants for standardized test queries and expected results @@ -31,6 +40,20 @@ public class AnyValueAggregationFunctionTest extends AbstractAggregationFunction private static final String EXPECTED_COLUMN_TYPES = "STRING | STRING"; private static final String EXPECTED_NULL_RESULT = "testResult | null"; + @Test + void rejectsRawVariant() { + ExpressionContext expression = ExpressionContext.forIdentifier("myField"); + AnyValueAggregationFunction function = new AnyValueAggregationFunction(List.of(expression), true); + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(FieldSpec.DataType.VARIANT); + + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> function.aggregate(1, function.createAggregationResultHolder(), Map.of(expression, blockValSet))); + + Assert.assertEquals(exception.getMessage(), + "ANY_VALUE does not support raw VARIANT values; extract a typed path with variantGet first"); + } + @DataProvider(name = "scenarios") Object[] scenarios() { return new Object[] { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java index d0deaf5bbc5c..1ee8e7629962 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java @@ -24,6 +24,11 @@ import java.util.stream.Collectors; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.exception.BadQueryRequestException; +import org.mockito.Mockito; +import org.testng.Assert; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; @@ -104,4 +109,18 @@ public void testTwoNullsCompareNextColumn() { assertEquals(extractColumn(_rows, COLUMN2_INDEX), Arrays.asList(1, 2, 3)); } + + @Test + public void testRejectsRawVariant() { + List orderBys = + List.of(new OrderByExpressionContext(COLUMN1, ASC, NULLS_LAST)); + ColumnContext columnContext = Mockito.mock(ColumnContext.class); + Mockito.when(columnContext.isSingleValue()).thenReturn(true); + Mockito.when(columnContext.getDataType()).thenReturn(DataType.VARIANT); + + BadQueryRequestException exception = Assert.expectThrows(BadQueryRequestException.class, + () -> OrderByComparatorFactory.getComparator(orderBys, new ColumnContext[]{columnContext}, + ENABLE_NULL_HANDLING)); + Assert.assertTrue(exception.getMessage().contains("ORDER BY does not support raw VARIANT")); + } } diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java index 4ce8b26ef999..c60d1caa6a2f 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java @@ -788,6 +788,7 @@ private static JsonNode extractValue(DataSchema.ColumnDataType columnDataType, J case STRING: case BYTES: case JSON: + case VARIANT: object = jsonValue.textValue(); break; case UNKNOWN: diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java new file mode 100644 index 000000000000..10eaceeecb56 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java @@ -0,0 +1,404 @@ +/** + * 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.integration.tests.custom; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.StandardCopyOption; +import java.util.List; +import java.util.Locale; +import org.apache.pinot.integration.tests.ClusterIntegrationTestUtils; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.FileFormat; +import org.apache.pinot.spi.utils.JsonUtils; +import org.testng.Assert; +import org.testng.annotations.Test; + + +/** + * End-to-end coverage for creating a VARIANT table, ingesting Parquet VARIANT(1), materializing a hot path, and + * querying nested values with both Pinot query engines. + */ +@Test(suiteName = "CustomClusterIntegrationTest") +public class VariantTypeTest extends CustomDataQueryClusterIntegrationTest { + private static final String RESOURCE_DIRECTORY = "examples/batch/variantEvents/"; + private static final String TABLE_NAME = "variantEvents"; + private static final String EVENT_ID = "eventId"; + private static final String EVENT_TYPE = "eventType"; + private static final String PAYLOAD = "payload"; + private static final int NUM_DOCS = 5; + + @Override + public String getTableName() { + return TABLE_NAME; + } + + @Override + protected long getCountStarResult() { + return NUM_DOCS; + } + + @Override + public Schema createSchema() { + try (InputStream inputStream = openResource("variantEvents_schema.json")) { + return Schema.fromInputStream(inputStream); + } catch (IOException e) { + throw new IllegalStateException("Failed to load the VARIANT quickstart schema", e); + } + } + + @Override + public TableConfig createOfflineTableConfig() { + try (InputStream inputStream = openResource("variantEvents_offline_table_config.json")) { + return JsonUtils.inputStreamToObject(inputStream, TableConfig.class); + } catch (IOException e) { + throw new IllegalStateException("Failed to load the VARIANT quickstart table config", e); + } + } + + @Override + protected void setUpTable() + throws Exception { + Schema schema = createSchema(); + addSchema(schema); + TableConfig tableConfig = createOfflineTableConfig(); + addTableConfig(tableConfig); + + File parquetFile = new File(_tempDir, "variantEvents_data.parquet"); + try (InputStream inputStream = openResource("rawdata/variantEvents_data.parquet")) { + Files.copy(inputStream, parquetFile.toPath(), StandardCopyOption.REPLACE_EXISTING); + } + ClusterIntegrationTestUtils.buildSegmentFromFile(parquetFile, tableConfig, schema, "0", _segmentDir, _tarDir, + FileFormat.PARQUET); + uploadSegments(getTableName(), _tarDir); + } + + @Override + public List createAvroFiles() { + throw new UnsupportedOperationException("VariantTypeTest ingests Parquet, not Avro"); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testMaterializedPathAndDirectExtraction(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postVariantQuery( + "SELECT " + EVENT_ID + ", " + EVENT_TYPE + " FROM " + TABLE_NAME + + " WHERE " + EVENT_TYPE + " = 'checkout' ORDER BY " + EVENT_ID); + assertNoExceptions(response); + JsonNode rows = response.get("resultTable").get("rows"); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"STRING\",\"STRING\"]"); + Assert.assertEquals(rows.size(), 2); + Assert.assertEquals(rows.get(0).get(0).asText(), "evt-001"); + Assert.assertEquals(rows.get(0).get(1).asText(), "checkout"); + Assert.assertEquals(rows.get(1).get(0).asText(), "evt-003"); + Assert.assertEquals(rows.get(1).get(1).asText(), "checkout"); + + response = postVariantQuery( + "SELECT " + EVENT_ID + ", variantGet(" + PAYLOAD + ", '$.user.id', 'STRING'), " + + "variantGet(" + PAYLOAD + ", '$.amount', 'DOUBLE') FROM " + TABLE_NAME + + " WHERE " + EVENT_TYPE + " = 'checkout' ORDER BY " + EVENT_ID); + assertNoExceptions(response); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"STRING\",\"STRING\",\"DOUBLE\"]"); + rows = response.get("resultTable").get("rows"); + Assert.assertEquals(rows.size(), 2); + Assert.assertEquals(rows.get(0).get(0).asText(), "evt-001"); + Assert.assertEquals(rows.get(0).get(1).asText(), "u-1"); + Assert.assertEquals(rows.get(0).get(2).asDouble(), 42.5); + Assert.assertEquals(rows.get(1).get(0).asText(), "evt-003"); + Assert.assertEquals(rows.get(1).get(1).asText(), "u-3"); + Assert.assertEquals(rows.get(1).get(2).asDouble(), 19.0); + + if (!useMultiStageQueryEngine) { + Assert.assertEquals(response.get("numEntriesScannedInFilter").asLong(), 0L, + "The materialized eventType filter should use its inverted index"); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testJsonProjectionAndNullStates(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postVariantQuery( + "SELECT " + PAYLOAD + ", variantToJson(" + PAYLOAD + ") FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " = 'evt-001'"); + assertNoExceptions(response); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"VARIANT\",\"STRING\"]"); + JsonNode row = response.get("resultTable").get("rows").get(0); + String json = row.get(0).asText(); + Assert.assertEquals(row.get(1).asText(), json); + Assert.assertEquals(json, + "{\"amount\":42.5,\"eventType\":\"checkout\",\"items\":[\"sku-1\",\"sku-2\"]," + + "\"user\":{\"id\":\"u-1\"}}"); + JsonNode payload = JsonUtils.stringToJsonNode(json); + Assert.assertEquals(payload.get("eventType").asText(), "checkout"); + Assert.assertEquals(payload.get("user").get("id").asText(), "u-1"); + Assert.assertEquals(payload.get("items").size(), 2); + + response = postVariantQuery( + "SELECT " + EVENT_ID + ", variantExists(" + PAYLOAD + ", '$.coupon'), " + + "isVariantNull(" + PAYLOAD + ", '$.coupon') FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " IN ('evt-001', 'evt-002', 'evt-003') ORDER BY " + EVENT_ID); + assertNoExceptions(response); + JsonNode rows = response.get("resultTable").get("rows"); + Assert.assertEquals(rows.size(), 3); + Assert.assertFalse(rows.get(0).get(1).asBoolean()); + Assert.assertFalse(rows.get(0).get(2).asBoolean()); + Assert.assertTrue(rows.get(1).get(1).asBoolean()); + Assert.assertTrue(rows.get(1).get(2).asBoolean()); + Assert.assertTrue(rows.get(2).get(1).asBoolean()); + Assert.assertFalse(rows.get(2).get(2).asBoolean()); + + response = postVariantQuery( + "SELECT " + EVENT_ID + ", variantTypeOf(" + PAYLOAD + ", '$') FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " IN ('evt-004', 'evt-005') ORDER BY " + EVENT_ID); + assertNoExceptions(response); + rows = response.get("resultTable").get("rows"); + Assert.assertEquals(rows.size(), 2); + Assert.assertEquals(rows.get(0).get(1).asText(), "NULL"); + Assert.assertTrue(rows.get(1).get(1).isNull()); + + response = postVariantQuery( + "SELECT " + EVENT_ID + ", " + PAYLOAD + " FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " IN ('evt-004', 'evt-005') ORDER BY " + EVENT_ID); + assertNoExceptions(response); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"STRING\",\"VARIANT\"]"); + rows = response.get("resultTable").get("rows"); + Assert.assertEquals(rows.get(0).get(1).asText(), "null", + "An encoded Variant null must render as JSON text"); + Assert.assertTrue(rows.get(1).get(1).isNull(), "A missing Parquet payload must remain SQL null"); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testSparkCompatibleFunctionSemantics(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postVariantQuery( + "SELECT variant_get(" + PAYLOAD + ", '$.user'), " + + "variant_get(" + PAYLOAD + ", '$.missing'), " + + "variant_get(" + PAYLOAD + ", '$.items[1]', 'STRING'), " + + "try_variant_get(" + PAYLOAD + ", '$.eventType', 'DOUBLE') " + + "FROM " + TABLE_NAME + " WHERE " + EVENT_ID + " = 'evt-001'"); + assertNoExceptions(response); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"VARIANT\",\"VARIANT\",\"STRING\",\"DOUBLE\"]"); + JsonNode row = response.get("resultTable").get("rows").get(0); + Assert.assertEquals(row.get(0).asText(), "{\"id\":\"u-1\"}"); + Assert.assertTrue(row.get(1).isNull(), "A missing path must return SQL null"); + Assert.assertEquals(row.get(2).asText(), "sku-2"); + Assert.assertTrue(row.get(3).isNull(), "try_variant_get must return SQL null for an incompatible cast"); + + response = postVariantQuery( + "SELECT variant_get(parse_json('{\"answer\":42}'), '$.answer', 'INT') " + + "FROM " + TABLE_NAME + " LIMIT 1"); + assertNoExceptions(response); + Assert.assertEquals(response.get("resultTable").get("rows").get(0).get(0).asInt(), 42); + + response = postVariantQuery( + "SELECT is_variant_null(" + PAYLOAD + ") FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " = 'evt-005'"); + assertNoExceptions(response); + JsonNode isVariantNull = response.get("resultTable").get("rows").get(0).get(0); + Assert.assertFalse(isVariantNull.isNull(), "is_variant_null(SQL NULL) must return a non-null boolean"); + Assert.assertFalse(isVariantNull.asBoolean(), + "SQL null is not an encoded Variant null"); + + response = postVariantQuery( + "SELECT variant_get(" + PAYLOAD + ", '$.eventType', 'DOUBLE') FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " = 'evt-001'"); + assertExceptionContains(response, "cannot convert variant", "double"); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testRawVariantInAndNotInAreRejected(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + for (String operator : List.of("IN", "NOT IN")) { + JsonNode response = postVariantQuery( + "SELECT " + EVENT_ID + " FROM " + TABLE_NAME + " WHERE " + PAYLOAD + " " + operator + + " (parse_json('{\"candidate\":1}'), parse_json('{\"candidate\":2}'))"); + assertExceptionContains(response, "raw variant", "in"); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testRawVariantOrderByIsRejected(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = + postVariantQuery("SELECT " + EVENT_ID + " FROM " + TABLE_NAME + " ORDER BY " + PAYLOAD + " LIMIT 5"); + assertExceptionContains(response, "raw variant", "order by"); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testRawVariantComparisonGroupingAndDistinctAreRejected(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postVariantQuery( + "SELECT " + EVENT_ID + " FROM " + TABLE_NAME + " WHERE " + PAYLOAD + + " = parse_json('{\"eventType\":\"checkout\"}')"); + assertExceptionContains(response, "raw variant", "comparison"); + response = postVariantQuery( + "SELECT " + PAYLOAD + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + PAYLOAD); + assertExceptionContains(response, "raw variant", "group by"); + response = postVariantQuery("SELECT DISTINCT " + PAYLOAD + " FROM " + TABLE_NAME); + assertExceptionContains(response, "raw variant", useMultiStageQueryEngine ? "group by" : "distinct"); + + for (String query : List.of( + "SELECT " + EVENT_ID + " FROM " + TABLE_NAME + " WHERE variant_get(" + PAYLOAD + + ", '$.eventType', 'STRING') = 'checkout'", + "SELECT variant_get(" + PAYLOAD + ", '$.eventType', 'STRING'), COUNT(*) FROM " + TABLE_NAME + + " GROUP BY variant_get(" + PAYLOAD + ", '$.eventType', 'STRING')", + "SELECT DISTINCT variant_get(" + PAYLOAD + ", '$.eventType', 'STRING') FROM " + TABLE_NAME)) { + assertNoExceptions(postVariantQuery(query)); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testRawVariantAggregatesAreRejected(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + for (String aggregate : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTTHETASKETCH")) { + JsonNode response = postVariantQuery("SELECT " + aggregate + "(" + PAYLOAD + ") FROM " + TABLE_NAME); + assertExceptionContains(response, "raw variant", aggregate); + } + + JsonNode response = postVariantQuery("SELECT COUNT(" + PAYLOAD + ") FROM " + TABLE_NAME); + assertNoExceptions(response); + Assert.assertEquals(response.get("resultTable").get("rows").get(0).get(0).asLong(), 4L, + "COUNT is raw-value-independent and must retain SQL-null semantics"); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testVariantFunctionsRequireQueryNullHandling(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postQuery( + "SELECT variant_get(" + PAYLOAD + ", '$.eventType', 'STRING') FROM " + TABLE_NAME + " LIMIT 1"); + assertExceptionContains(response, "requires query null handling"); + } + + @Test + public void testRawVariantJoinIsRejectedButTypedPathJoinWorks() + throws Exception { + setUseMultiStageQueryEngine(true); + String rawJoin = "SELECT leftTable." + EVENT_ID + " FROM " + TABLE_NAME + " leftTable JOIN " + TABLE_NAME + + " rightTable ON leftTable." + PAYLOAD + " = rightTable." + PAYLOAD; + JsonNode response = postVariantQuery(rawJoin); + assertExceptionContains(response, "raw variant", "join"); + + String typedJoin = "SELECT leftTable." + EVENT_ID + ", rightTable." + EVENT_ID + " FROM " + TABLE_NAME + + " leftTable JOIN " + TABLE_NAME + " rightTable ON variant_get(leftTable." + PAYLOAD + + ", '$.eventType', 'STRING') = variant_get(rightTable." + PAYLOAD + ", '$.eventType', 'STRING')" + + " WHERE leftTable." + EVENT_ID + " = 'evt-001' AND rightTable." + EVENT_ID + " = 'evt-003'"; + response = postVariantQuery(typedJoin); + assertNoExceptions(response); + Assert.assertEquals(response.get("resultTable").get("rows").size(), 1); + Assert.assertEquals(response.get("resultTable").get("rows").get(0).get(0).asText(), "evt-001"); + Assert.assertEquals(response.get("resultTable").get("rows").get(0).get(1).asText(), "evt-003"); + } + + @Test + public void testRawVariantWindowKeysAreRejectedButTypedPathWorks() + throws Exception { + setUseMultiStageQueryEngine(true); + + for (String window : List.of( + "COUNT(*) OVER (PARTITION BY " + PAYLOAD + ")", + "COUNT(*) OVER (ORDER BY " + PAYLOAD + ")")) { + JsonNode response = postVariantQuery("SELECT " + window + " FROM " + TABLE_NAME); + assertExceptionContains(response, "raw variant", "window"); + } + + JsonNode response = postVariantQuery( + "SELECT " + EVENT_ID + ", COUNT(*) OVER (PARTITION BY variant_get(" + PAYLOAD + + ", '$.eventType', 'STRING')) FROM " + TABLE_NAME + " ORDER BY " + EVENT_ID); + assertNoExceptions(response); + Assert.assertEquals(response.get("resultTable").get("rows").size(), NUM_DOCS); + } + + @Test + public void testRawVariantSetOperationsAreRejected() + throws Exception { + setUseMultiStageQueryEngine(true); + String left = "SELECT " + PAYLOAD + " FROM " + TABLE_NAME + " WHERE " + EVENT_ID + " = 'evt-001'"; + String right = "SELECT " + PAYLOAD + " FROM " + TABLE_NAME + " WHERE " + EVENT_ID + " = 'evt-002'"; + + for (String operator : List.of("UNION", "INTERSECT", "INTERSECT ALL", "EXCEPT", "EXCEPT ALL")) { + JsonNode response = postVariantQuery(left + " " + operator + " " + right); + assertExceptionContains(response, "raw variant", "extract a typed path"); + } + + JsonNode response = postVariantQuery(left + " UNION ALL " + right); + assertNoExceptions(response); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), "[\"VARIANT\"]"); + Assert.assertEquals(response.get("resultTable").get("rows").size(), 2); + } + + private static InputStream openResource(String relativePath) { + String resourcePath = RESOURCE_DIRECTORY + relativePath; + InputStream inputStream = VariantTypeTest.class.getClassLoader().getResourceAsStream(resourcePath); + if (inputStream == null) { + throw new IllegalStateException("Missing VARIANT quickstart resource: " + resourcePath); + } + return inputStream; + } + + private static void assertNoExceptions(JsonNode response) { + Assert.assertEquals(response.get("exceptions").size(), 0, response.toPrettyString()); + } + + private static void assertExceptionContains(JsonNode response, String... expectedFragments) { + JsonNode exceptions = response.get("exceptions"); + Assert.assertNotNull(exceptions, response.toPrettyString()); + Assert.assertFalse(exceptions.isEmpty(), response.toPrettyString()); + String exceptionText = exceptions.toString().toLowerCase(Locale.ROOT); + for (String expectedFragment : expectedFragments) { + Assert.assertTrue(exceptionText.contains(expectedFragment.toLowerCase(Locale.ROOT)), + "Expected exception text to contain '" + expectedFragment + "': " + response.toPrettyString()); + } + } + + private JsonNode postVariantQuery(String query) + throws Exception { + return postQuery("SET enableNullHandling=true; " + query); + } +} diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/pom.xml b/pinot-plugins/pinot-input-format/pinot-parquet/pom.xml index c70f6a369e90..73afda7171ca 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/pom.xml +++ b/pinot-plugins/pinot-input-format/pinot-parquet/pom.xml @@ -43,6 +43,10 @@ org.apache.parquet parquet-avro + + org.apache.parquet + parquet-variant + org.apache.hadoop hadoop-common diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordExtractor.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordExtractor.java index 48e6eddd7392..89adbc0751ce 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordExtractor.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordExtractor.java @@ -24,9 +24,9 @@ import org.apache.pinot.spi.utils.TimestampUtils; -/// The type matrix is inherited from [AvroRecordExtractor]; the only override is the INT96 timestamp -/// (which parquet-avro surfaces as `fixed(12)` with `doc = "INT96 represented as byte[12]"`) → `Timestamp` -/// (or `Long` epoch nanos when `extractRawTimeValues` is `true`) via [ParquetUtils#convertInt96ToEpochNanos]. +/// Extends [AvroRecordExtractor] with Parquet-specific conversions: +/// - INT96 (surfaced by parquet-avro as `fixed(12)` with `doc = "INT96 represented as byte[12]"`) becomes `Timestamp`, +/// or `Long` epoch nanos when `extractRawTimeValues` is `true`. public class ParquetAvroRecordExtractor extends AvroRecordExtractor { private static final int INT96_BYTE_SIZE = 12; private static final String INT96_DOC = "INT96 represented as byte[12]"; diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordReader.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordReader.java index d8e3e3263b62..7ebbf49fc5d3 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordReader.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetAvroRecordReader.java @@ -52,16 +52,36 @@ public class ParquetAvroRecordReader implements RecordReader { public void init(File dataFile, @Nullable Set fieldsToRead, @Nullable RecordReaderConfig recordReaderConfig) throws IOException { File parquetFile = RecordReaderUtils.unpackIfRequired(dataFile, EXTENSION); - _dataFilePath = new Path(parquetFile.getAbsolutePath()); - _parquetReader = ParquetUtils.getParquetAvroReader(_dataFilePath); + Path dataFilePath = new Path(parquetFile.getAbsolutePath()); AvroRecordExtractorConfig extractorConfig = new AvroRecordExtractorConfig(); if (recordReaderConfig instanceof ParquetRecordReaderConfig) { extractorConfig.setExtractRawTimeValues( ((ParquetRecordReaderConfig) recordReaderConfig).isExtractRawTimeValues()); } - _recordExtractor = new ParquetAvroRecordExtractor(); - _recordExtractor.init(fieldsToRead, extractorConfig); - _nextRecord = _parquetReader.read(); + ParquetAvroRecordExtractor recordExtractor = new ParquetAvroRecordExtractor(); + recordExtractor.init(fieldsToRead, extractorConfig); + + ParquetReader parquetReader = ParquetUtils.getParquetAvroReader(dataFilePath); + GenericRecord nextRecord; + try { + nextRecord = parquetReader.read(); + } catch (IOException | RuntimeException e) { + try { + parquetReader.close(); + } catch (IOException | RuntimeException closeException) { + e.addSuppressed(closeException); + } + throw e; + } + ParquetReader previousReader = _parquetReader; + // Publish only the fully initialized replacement. A previous-reader close failure must not restore stale state. + _dataFilePath = dataFilePath; + _parquetReader = parquetReader; + _recordExtractor = recordExtractor; + _nextRecord = nextRecord; + if (previousReader != null) { + previousReader.close(); + } } @Override @@ -86,14 +106,35 @@ public GenericRow next(GenericRow reuse) @Override public void rewind() throws IOException { - _parquetReader.close(); - _parquetReader = ParquetUtils.getParquetAvroReader(_dataFilePath); - _nextRecord = _parquetReader.read(); + ParquetReader parquetReader = _parquetReader; + _parquetReader = null; + _nextRecord = null; + parquetReader.close(); + + ParquetReader rewoundReader = ParquetUtils.getParquetAvroReader(_dataFilePath); + try { + _nextRecord = rewoundReader.read(); + _parquetReader = rewoundReader; + } catch (IOException | RuntimeException e) { + try { + rewoundReader.close(); + } catch (IOException | RuntimeException closeException) { + e.addSuppressed(closeException); + } + throw e; + } } @Override public void close() throws IOException { - _parquetReader.close(); + ParquetReader parquetReader = _parquetReader; + _dataFilePath = null; + _recordExtractor = null; + _parquetReader = null; + _nextRecord = null; + if (parquetReader != null) { + parquetReader.close(); + } } } diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java index cbc92e6846d7..481dd2b0df17 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java @@ -23,7 +23,6 @@ import java.math.BigInteger; import java.time.LocalDate; import java.time.LocalTime; -import java.util.List; import java.util.Map; import javax.annotation.Nullable; import org.apache.parquet.example.data.Group; @@ -60,6 +59,7 @@ /// - `INT64` + `TIME_MICROS` / `TIME_NANOS` → `LocalTime`, or `Long` value-since-midnight in the column's /// declared unit when `extractRawTimeValues` is `true` /// - `FIXED_LEN_BYTE_ARRAY(16)` + `UUID` → `java.util.UUID` +/// - top-level non-repeated group + `VARIANT(1)` → Pinot VARIANT `byte[]` envelope /// /// **Complex types:** /// - `LIST`-annotated group (standard 3-level wrapper or legacy non-wrapper forms) → `Object[]` @@ -68,38 +68,93 @@ /// - field with zero repetition count → `null` public class ParquetNativeRecordExtractor extends BaseRecordExtractor { + private static final TopLevelFieldPlan[] NO_FIELD_PLANS = new TopLevelFieldPlan[0]; + private boolean _extractRawTimeValues; + private GroupType _plannedSchema; + private TopLevelFieldPlan[] _fieldPlans = NO_FIELD_PLANS; @Override protected void initConfig(@Nullable RecordExtractorConfig config) { + _extractRawTimeValues = false; + _plannedSchema = null; + _fieldPlans = NO_FIELD_PLANS; if (config instanceof ParquetNativeRecordExtractorConfig) { - _extractRawTimeValues = ((ParquetNativeRecordExtractorConfig) config).isExtractRawTimeValues(); + ParquetNativeRecordExtractorConfig parquetConfig = (ParquetNativeRecordExtractorConfig) config; + _extractRawTimeValues = parquetConfig.isExtractRawTimeValues(); + GroupType parquetSchema = parquetConfig.getParquetSchema(); + if (parquetSchema != null) { + initializeFieldPlans(parquetSchema); + } } } @Override public GenericRow extract(Group from, GenericRow to) { GroupType fromType = from.getType(); - if (_extractAll) { - List fields = fromType.getFields(); - for (Type field : fields) { - String fieldName = field.getName(); - to.putValue(fieldName, extractValue(from, fromType.getFieldIndex(fieldName))); + if (_plannedSchema == null || (_plannedSchema != fromType && !_plannedSchema.equals(fromType))) { + // Direct extractor users historically initialize without a file schema. Preserve that usage while keeping + // converter ownership inside extractor initialization for the normal record-reader path. + initializeFieldPlans(fromType); + } else if (_plannedSchema != fromType) { + // Avoid a deep schema equality check on every subsequent row when the configured and row schemas are equal + // but represented by distinct objects. + _plannedSchema = fromType; + } + for (TopLevelFieldPlan fieldPlan : _fieldPlans) { + to.putValue(fieldPlan._name, extractTopLevelValue(from, fieldPlan)); + } + return to; + } + + private void initializeFieldPlans(GroupType schema) { + ParquetVariantConverter[] variantConverters = + ParquetVariantConverter.createTopLevelVariantConverters(schema); + int fieldCount = schema.getFieldCount(); + int selectedFieldCount = 0; + for (int fieldIndex = 0; fieldIndex < fieldCount; fieldIndex++) { + if (_extractAll || _fields.contains(schema.getType(fieldIndex).getName())) { + selectedFieldCount++; } - } else { - for (String fieldName : _fields) { - if (fromType.containsField(fieldName)) { - to.putValue(fieldName, extractValue(from, fromType.getFieldIndex(fieldName))); - } + } + + TopLevelFieldPlan[] fieldPlans = new TopLevelFieldPlan[selectedFieldCount]; + int selectedFieldIndex = 0; + for (int fieldIndex = 0; fieldIndex < fieldCount; fieldIndex++) { + Type fieldType = schema.getType(fieldIndex); + if (_extractAll || _fields.contains(fieldType.getName())) { + fieldPlans[selectedFieldIndex++] = + new TopLevelFieldPlan(fieldType.getName(), fieldIndex, fieldType, variantConverters[fieldIndex]); } } - return to; + _plannedSchema = schema; + _fieldPlans = fieldPlans; + } + + @Nullable + private Object extractTopLevelValue(Group from, TopLevelFieldPlan fieldPlan) { + if (fieldPlan._variantConverter == null) { + return extractValue(from, fieldPlan._fieldIndex, fieldPlan._fieldType); + } + int numValues = from.getFieldRepetitionCount(fieldPlan._fieldIndex); + if (numValues == 0) { + return null; + } + return fieldPlan._variantConverter.convert(from.getGroup(fieldPlan._fieldIndex, 0)); } @Nullable private Object extractValue(Group from, int fieldIndex) { - int numValues = from.getFieldRepetitionCount(fieldIndex); Type fieldType = from.getType().getType(fieldIndex); + return extractValue(from, fieldIndex, fieldType); + } + + @Nullable + private Object extractValue(Group from, int fieldIndex, Type fieldType) { + int numValues = from.getFieldRepetitionCount(fieldIndex); + if (ParquetVariantConverter.isVariant(fieldType)) { + throw new UnsupportedOperationException("Nested Parquet VARIANT is not supported: " + fieldType.getName()); + } // REPEATED fields are always multi-valued — even when 0 or 1 occurrences are present, the contract is // `Object[]` (matching how LIST-annotated groups surface). For OPTIONAL / REQUIRED fields, 0 → null and // 1 → the scalar. @@ -313,4 +368,19 @@ private boolean isStandardListWrapper(Type repeatedField, String parentListName) String repeatedFieldName = repeatedField.getName(); return !"array".equals(repeatedFieldName) && !(parentListName + "_tuple").equals(repeatedFieldName); } + + private static final class TopLevelFieldPlan { + private final String _name; + private final int _fieldIndex; + private final Type _fieldType; + private final ParquetVariantConverter _variantConverter; + + private TopLevelFieldPlan(String name, int fieldIndex, Type fieldType, + @Nullable ParquetVariantConverter variantConverter) { + _name = name; + _fieldIndex = fieldIndex; + _fieldType = fieldType; + _variantConverter = variantConverter; + } + } } diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java index 86d3397d568c..b9e348305b33 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java @@ -19,6 +19,8 @@ package org.apache.pinot.plugin.inputformat.parquet; import java.util.Map; +import javax.annotation.Nullable; +import org.apache.parquet.schema.GroupType; import org.apache.pinot.spi.data.readers.RecordExtractorConfig; @@ -28,6 +30,8 @@ public class ParquetNativeRecordExtractorConfig implements RecordExtractorConfig public static final String EXTRACT_RAW_TIME_VALUES = "extractRawTimeValues"; private boolean _extractRawTimeValues; + @Nullable + private GroupType _parquetSchema; @Override public void init(Map props) { @@ -41,4 +45,20 @@ public boolean isExtractRawTimeValues() { public void setExtractRawTimeValues(boolean extractRawTimeValues) { _extractRawTimeValues = extractRawTimeValues; } + + /** + * Supplies the immutable Parquet record schema used to initialize schema-bound logical-type converters. + * + *

The native record reader sets this before initializing the extractor. Direct extractor users should do the + * same when the schema contains VARIANT columns; otherwise the extractor initializes those converters from the + * first record as a compatibility fallback. + */ + public void setParquetSchema(GroupType parquetSchema) { + _parquetSchema = parquetSchema; + } + + @Nullable + GroupType getParquetSchema() { + return _parquetSchema; + } } diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordReader.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordReader.java index 744f6b0484f9..074eab516f7b 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordReader.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordReader.java @@ -61,23 +61,60 @@ public class ParquetNativeRecordReader implements RecordReader { public void init(File dataFile, @Nullable Set fieldsToRead, @Nullable RecordReaderConfig recordReaderConfig) throws IOException { File parquetFile = RecordReaderUtils.unpackIfRequired(dataFile, EXTENSION); - _dataFilePath = new Path(parquetFile.getAbsolutePath()); - _hadoopConf = ParquetUtils.getParquetHadoopConfiguration(); - ParquetNativeRecordExtractorConfig extractorConfig = new ParquetNativeRecordExtractorConfig(); - if (recordReaderConfig instanceof ParquetRecordReaderConfig) { - extractorConfig.setExtractRawTimeValues( - ((ParquetRecordReaderConfig) recordReaderConfig).isExtractRawTimeValues()); - } - _recordExtractor = new ParquetNativeRecordExtractor(); - _recordExtractor.init(fieldsToRead, extractorConfig); + Path dataFilePath = new Path(parquetFile.getAbsolutePath()); + Configuration hadoopConf = ParquetUtils.getParquetHadoopConfiguration(); + ParquetReadOptions parquetReadOptions = + ParquetReadOptions.builder().withMetadataFilter(ParquetMetadataConverter.NO_FILTER).build(); - _parquetReadOptions = ParquetReadOptions.builder().withMetadataFilter(ParquetMetadataConverter.NO_FILTER).build(); + ParquetFileReader previousReader = _parquetFileReader; + ParquetFileReader parquetFileReader = + ParquetFileReader.open(HadoopInputFile.fromPath(dataFilePath, hadoopConf), parquetReadOptions); + MessageType schema; + ParquetNativeRecordExtractor recordExtractor; + MessageColumnIO columnIO; + PageReadStore pageReadStore; + org.apache.parquet.io.RecordReader parquetRecordReader; + try { + schema = parquetFileReader.getFooter().getFileMetaData().getSchema(); + ParquetNativeRecordExtractorConfig extractorConfig = new ParquetNativeRecordExtractorConfig(); + extractorConfig.setParquetSchema(schema); + if (recordReaderConfig instanceof ParquetRecordReaderConfig) { + extractorConfig.setExtractRawTimeValues( + ((ParquetRecordReaderConfig) recordReaderConfig).isExtractRawTimeValues()); + } + recordExtractor = new ParquetNativeRecordExtractor(); + recordExtractor.init(fieldsToRead, extractorConfig); - _parquetFileReader = - ParquetFileReader.open(HadoopInputFile.fromPath(_dataFilePath, _hadoopConf), _parquetReadOptions); - _schema = _parquetFileReader.getFooter().getFileMetaData().getSchema(); - _columnIO = new ColumnIOFactory().getColumnIO(_schema); - init(); + columnIO = new ColumnIOFactory().getColumnIO(schema); + pageReadStore = parquetFileReader.readNextRowGroup(); + parquetRecordReader = pageReadStore != null + ? columnIO.getRecordReader(pageReadStore, new GroupRecordConverter(schema)) : null; + } catch (IOException | RuntimeException e) { + try { + parquetFileReader.close(); + } catch (IOException closeException) { + e.addSuppressed(closeException); + } + throw e; + } + + // Publish the fully initialized replacement before cleaning up the old reader. If old-reader cleanup fails, the + // caller sees the failure but this instance remains attached to the usable replacement instead of stale, + // partially closed state. + _dataFilePath = dataFilePath; + _hadoopConf = hadoopConf; + _parquetReadOptions = parquetReadOptions; + _parquetFileReader = parquetFileReader; + _schema = schema; + _recordExtractor = recordExtractor; + _columnIO = columnIO; + _pageReadStore = pageReadStore; + _parquetRecordReader = parquetRecordReader; + _nextRecord = null; + _currentPageIdx = 0; + if (previousReader != null) { + previousReader.close(); + } } private void init() @@ -122,9 +159,11 @@ public GenericRow next(GenericRow reuse) } catch (Exception e) { throw new RecordFetchException("Failed to read next Parquet native record", e); } + // The physical reader has consumed the row. Advance before logical extraction so continueOnError callers can + // recover from malformed values without retrying the same row or leaving hasNext() stuck at end-of-file. + _currentPageIdx++; // Data parsing: extract into GenericRow. _recordExtractor.extract(_nextRecord, reuse); - _currentPageIdx++; return reuse; } @@ -140,6 +179,8 @@ public void rewind() @Override public void close() throws IOException { - _parquetFileReader.close(); + if (_parquetFileReader != null) { + _parquetFileReader.close(); + } } } diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReader.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReader.java index 4296b5f02592..250dff890998 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReader.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReader.java @@ -36,28 +36,42 @@ public class ParquetRecordReader implements RecordReader { private static final String EXTENSION = "parquet"; private RecordReader _internalParquetRecordReader; - private boolean _useAvroParquetRecordReader = true; @Override public void init(File dataFile, @Nullable Set fieldsToRead, @Nullable RecordReaderConfig recordReaderConfig) throws IOException { File parquetFile = RecordReaderUtils.unpackIfRequired(dataFile, EXTENSION); + RecordReader nextReader; if (recordReaderConfig != null && ((ParquetRecordReaderConfig) recordReaderConfig).useParquetAvroRecordReader()) { - _internalParquetRecordReader = new ParquetAvroRecordReader(); + nextReader = new ParquetAvroRecordReader(); } else if (recordReaderConfig != null && ((ParquetRecordReaderConfig) recordReaderConfig).useParquetNativeRecordReader()) { - _useAvroParquetRecordReader = false; - _internalParquetRecordReader = new ParquetNativeRecordReader(); + nextReader = new ParquetNativeRecordReader(); } else { // No reader type specified. Determine using file metadata if (ParquetUtils.hasAvroSchemaInFileMetadata(new Path(parquetFile.getAbsolutePath()))) { - _internalParquetRecordReader = new ParquetAvroRecordReader(); + nextReader = new ParquetAvroRecordReader(); } else { - _useAvroParquetRecordReader = false; - _internalParquetRecordReader = new ParquetNativeRecordReader(); + nextReader = new ParquetNativeRecordReader(); } } - _internalParquetRecordReader.init(parquetFile, fieldsToRead, recordReaderConfig); + try { + nextReader.init(parquetFile, fieldsToRead, recordReaderConfig); + } catch (IOException | RuntimeException e) { + try { + nextReader.close(); + } catch (IOException | RuntimeException closeException) { + e.addSuppressed(closeException); + } + throw e; + } + RecordReader previousReader = _internalParquetRecordReader; + // Publish only the fully initialized replacement. A previous-reader close failure must not restore a stale + // delegate. + _internalParquetRecordReader = nextReader; + if (previousReader != null) { + previousReader.close(); + } } @Override @@ -80,10 +94,14 @@ public void rewind() @Override public void close() throws IOException { - _internalParquetRecordReader.close(); + RecordReader reader = _internalParquetRecordReader; + _internalParquetRecordReader = null; + if (reader != null) { + reader.close(); + } } public boolean useAvroParquetRecordReader() { - return _useAvroParquetRecordReader; + return _internalParquetRecordReader instanceof ParquetAvroRecordReader; } } diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReaderConfig.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReaderConfig.java index 70f8d152046c..db6c19fcb6b2 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReaderConfig.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetRecordReaderConfig.java @@ -25,7 +25,8 @@ /// Three settings, all default `false`: /// - `useParquetAvroRecordReader` — force the parquet-avro reader. /// - `useParquetNativeRecordReader` — force the native parquet reader. When neither flag is set, the dispatcher -/// auto-detects via the file's `avro.schema` metadata (Avro reader if present, native otherwise). +/// auto-detects via the file's `avro.schema` metadata (Avro reader if present, native otherwise). Set this flag for +/// VARIANT files that also carry Avro schema metadata so that the native reader retains the physical VARIANT type. /// - `extractRawTimeValues` — opt out of TIMESTAMP / DATE / TIME conversion at the extractor boundary, /// surfacing the raw underlying integer in the column's declared unit instead of the contract Java type. /// DECIMAL and UUID always convert. See [ParquetAvroRecordExtractor] / [ParquetNativeRecordExtractor] for diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java index f741095ff613..d928485ffec3 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java @@ -34,6 +34,7 @@ import org.apache.parquet.avro.AvroSchemaConverter; import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.ParquetReader; +import org.apache.parquet.hadoop.metadata.FileMetaData; import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.io.InputFile; import org.apache.parquet.schema.LogicalTypeAnnotation; @@ -82,15 +83,34 @@ public static Schema getParquetAvroSchema(Path path) } } - public static boolean hasAvroSchemaInFileMetadata(Path path) + /** + * Returns the physical Parquet schema for the given file path. + */ + public static MessageType getParquetSchema(Path path) + throws IOException { + return getParquetFileMetadata(path).getSchema(); + } + + /** + * Returns the immutable footer metadata for the given Parquet file path. + */ + public static FileMetaData getParquetFileMetadata(Path path) throws IOException { InputFile inputFile = HadoopInputFile.fromPath(path, getParquetHadoopConfiguration()); try (ParquetFileReader reader = ParquetFileReader.open(inputFile)) { - Map metaData = reader.getFileMetaData().getKeyValueMetaData(); - return metaData.containsKey(AVRO_SCHEMA_METADATA_KEY) || metaData.containsKey(OLD_AVRO_SCHEMA_METADATA_KEY); + return reader.getFileMetaData(); } } + public static boolean hasAvroSchemaInFileMetadata(Path path) + throws IOException { + return hasAvroSchemaInFileMetadata(getParquetFileMetadata(path).getKeyValueMetaData()); + } + + static boolean hasAvroSchemaInFileMetadata(Map metadata) { + return metadata.containsKey(AVRO_SCHEMA_METADATA_KEY) || metadata.containsKey(OLD_AVRO_SCHEMA_METADATA_KEY); + } + public static Configuration getParquetHadoopConfiguration() { // The file path used in ParquetRecordReader is a local file path without prefix 'file:///', // so we have to make sure that the configuration item 'fs.defaultFS' is set to 'file:///' diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java new file mode 100644 index 000000000000..6862e49331af --- /dev/null +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java @@ -0,0 +1,305 @@ +/** + * 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.plugin.inputformat.parquet; + +import java.nio.ByteBuffer; +import java.util.HashSet; +import java.util.Set; +import java.util.function.Consumer; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.io.api.Converter; +import org.apache.parquet.io.api.GroupConverter; +import org.apache.parquet.io.api.PrimitiveConverter; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.apache.parquet.variant.ImmutableMetadata; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.parquet.variant.VariantConverters; +import org.apache.pinot.spi.utils.VariantEnvelope; + + +/// Converts a top-level Parquet `VARIANT(1)` group into Pinot's self-describing VARIANT byte envelope. +/// +/// Detection is based only on [LogicalTypeAnnotation.VariantLogicalTypeAnnotation], never on field names. +/// The parquet-java [VariantConverters] implementation performs the logical reconstruction so unshredded, +/// shredded, and partially shredded encodings all produce complete Variant metadata and value buffers. +/// +/// Instances retain parquet-java's converter tree and are intentionally scoped to one record reader because the +/// converter and builder holder are mutable and not thread-safe. +final class ParquetVariantConverter { + private static final byte SUPPORTED_SPEC_VERSION = 1; + + private final GroupType _variantType; + private final int _metadataIndex; + private final int _valueIndex; + private final int _typedValueIndex; + private final BuilderHolder _holder = new BuilderHolder(); + private final GroupConverter _converter; + + private ParquetVariantConverter(GroupType variantType) { + _variantType = variantType; + _metadataIndex = variantType.getFieldIndex("metadata"); + _valueIndex = variantType.containsField("value") ? variantType.getFieldIndex("value") : -1; + _typedValueIndex = variantType.containsField("typed_value") ? variantType.getFieldIndex("typed_value") : -1; + _converter = VariantConverters.newVariantConverter(variantType, _holder::setMetadata, _holder::build); + } + + /// Validates every VARIANT annotation in the file and returns the supported top-level field names. + /// + /// Pinot currently supports only non-repeated, top-level `VARIANT(1)` values. Failing during reader + /// initialization avoids silently surfacing an unsupported nested/repeated Variant as an ordinary struct. + static Set validateAndGetTopLevelVariantFields(GroupType schema) { + Set variantFields = new HashSet<>(); + for (Type field : schema.getFields()) { + if (isVariant(field)) { + validateTopLevelVariant(field); + variantFields.add(field.getName()); + } else if (!field.isPrimitive()) { + rejectNestedVariants(field.asGroupType(), field.getName()); + } + } + return Set.copyOf(variantFields); + } + + /// Validates the file schema and builds an index-aligned reusable converter tree for each top-level VARIANT column. + static ParquetVariantConverter[] createTopLevelVariantConverters(GroupType schema) { + Set variantFields = validateAndGetTopLevelVariantFields(schema); + ParquetVariantConverter[] variantConverters = new ParquetVariantConverter[schema.getFieldCount()]; + for (int fieldIndex = 0; fieldIndex < schema.getFieldCount(); fieldIndex++) { + Type field = schema.getType(fieldIndex); + if (variantFields.contains(field.getName())) { + variantConverters[fieldIndex] = new ParquetVariantConverter(field.asGroupType()); + } + } + return variantConverters; + } + + static boolean isVariant(Type type) { + return type.getLogicalTypeAnnotation() instanceof LogicalTypeAnnotation.VariantLogicalTypeAnnotation; + } + + byte[] convert(Group group) { + if (_valueIndex >= 0 && group.getFieldRepetitionCount(_valueIndex) == 1 + && (_typedValueIndex < 0 || group.getFieldRepetitionCount(_typedValueIndex) == 0)) { + return pack(group.getBinary(_metadataIndex, 0), group.getBinary(_valueIndex, 0)); + } + _holder.reset(); + feedVariantGroup(group, _converter, _variantType, _metadataIndex); + Variant variant = _holder.finish(); + return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); + } + + static byte[] pack(ByteBuffer metadata, ByteBuffer value) { + return VariantEnvelope.encode(metadata, value); + } + + private byte[] pack(Binary metadata, Binary value) { + // Binary.writeTo(OutputStream) materializes the full payload for a direct or read-only ByteBuffer-backed Binary. + // Consume the independent views returned by toByteBuffer() so both payloads go straight into the final envelope, + // without intermediate byte arrays or another duplicate view. + int metadataLength = metadata.length(); + int valueLength = value.length(); + byte[] envelope = VariantEnvelope.allocate(metadataLength, valueLength); + copyInto(metadata, envelope, VariantEnvelope.HEADER_SIZE, metadataLength, "metadata"); + copyInto(value, envelope, VariantEnvelope.HEADER_SIZE + metadataLength, valueLength, "value"); + return envelope; + } + + private static void copyInto(Binary source, byte[] target, int targetOffset, int expectedLength, String payload) { + ByteBuffer sourceView = source.toByteBuffer(); + if (sourceView.remaining() != expectedLength) { + throw new IllegalStateException( + "Parquet VARIANT " + payload + " length changed while copying: expected " + expectedLength + " but found " + + sourceView.remaining()); + } + if (sourceView.hasArray()) { + System.arraycopy(sourceView.array(), sourceView.arrayOffset() + sourceView.position(), target, targetOffset, + expectedLength); + } else { + sourceView.get(target, targetOffset, expectedLength); + } + } + + private static void validateTopLevelVariant(Type field) { + if (!isVariant(field) || field.isPrimitive()) { + throw new IllegalArgumentException("Parquet VARIANT must be a group: " + field); + } + LogicalTypeAnnotation.VariantLogicalTypeAnnotation annotation = + (LogicalTypeAnnotation.VariantLogicalTypeAnnotation) field.getLogicalTypeAnnotation(); + if (annotation.getSpecVersion() != SUPPORTED_SPEC_VERSION) { + throw new UnsupportedOperationException( + "Unsupported Parquet VARIANT spec version: " + annotation.getSpecVersion()); + } + if (field.isRepetition(Type.Repetition.REPEATED)) { + throw new UnsupportedOperationException("Repeated Parquet VARIANT is not supported: " + field.getName()); + } + + GroupType group = field.asGroupType(); + for (Type child : group.getFields()) { + String name = child.getName(); + if (!"metadata".equals(name) && !"value".equals(name) && !"typed_value".equals(name)) { + throw new IllegalArgumentException("Invalid Parquet VARIANT field: " + child); + } + } + if (!group.containsField("metadata")) { + throw new IllegalArgumentException("Invalid Parquet VARIANT: missing metadata field"); + } + Type metadata = group.getType("metadata"); + if (!isBinary(metadata) || !metadata.isRepetition(Type.Repetition.REQUIRED)) { + throw new IllegalArgumentException("Invalid Parquet VARIANT metadata field: " + metadata); + } + + boolean hasValue = group.containsField("value"); + boolean hasTypedValue = group.containsField("typed_value"); + if (!hasValue && !hasTypedValue) { + throw new IllegalArgumentException("Invalid Parquet VARIANT: missing value and typed_value fields"); + } + if (hasValue) { + Type value = group.getType("value"); + if (!isBinary(value) || value.isRepetition(Type.Repetition.REPEATED)) { + throw new IllegalArgumentException("Invalid Parquet VARIANT value field: " + value); + } + } + if (hasTypedValue) { + Type typedValue = group.getType("typed_value"); + if (!typedValue.isRepetition(Type.Repetition.OPTIONAL)) { + throw new IllegalArgumentException("Invalid Parquet VARIANT typed_value field: " + typedValue); + } + } + } + + private static void rejectNestedVariants(GroupType group, String path) { + for (Type child : group.getFields()) { + String childPath = path + "." + child.getName(); + if (isVariant(child)) { + throw new UnsupportedOperationException("Nested Parquet VARIANT is not supported: " + childPath); + } + if (!child.isPrimitive()) { + rejectNestedVariants(child.asGroupType(), childPath); + } + } + } + + private static boolean isBinary(Type field) { + return field.isPrimitive() + && field.asPrimitiveType().getPrimitiveTypeName() == PrimitiveType.PrimitiveTypeName.BINARY; + } + + private static void feedVariantGroup(Group group, GroupConverter converter, GroupType type, int metadataIndex) { + converter.start(); + + // Metadata initializes VariantBuilder and must be delivered before value/typed_value. File field order is not + // semantically significant (early Spark writers emitted value first). The immutable field index is resolved once + // when the converter is built instead of by name for every row. + feedField(group, metadataIndex, converter.getConverter(metadataIndex)); + for (int fieldIndex = 0; fieldIndex < type.getFieldCount(); fieldIndex++) { + if (fieldIndex != metadataIndex) { + feedField(group, fieldIndex, converter.getConverter(fieldIndex)); + } + } + converter.end(); + } + + private static void feedGroup(Group group, GroupConverter converter) { + converter.start(); + GroupType type = group.getType(); + for (int fieldIndex = 0; fieldIndex < type.getFieldCount(); fieldIndex++) { + feedField(group, fieldIndex, converter.getConverter(fieldIndex)); + } + converter.end(); + } + + private static void feedField(Group group, int fieldIndex, Converter converter) { + int repetitionCount = group.getFieldRepetitionCount(fieldIndex); + if (repetitionCount == 0) { + return; + } + if (converter == null) { + throw new IllegalArgumentException( + "Unsupported Parquet VARIANT shape at field: " + group.getType().getType(fieldIndex)); + } + + Type field = group.getType().getType(fieldIndex); + for (int valueIndex = 0; valueIndex < repetitionCount; valueIndex++) { + if (field.isPrimitive()) { + feedPrimitive(group, fieldIndex, valueIndex, field.asPrimitiveType(), converter.asPrimitiveConverter()); + } else { + feedGroup(group.getGroup(fieldIndex, valueIndex), converter.asGroupConverter()); + } + } + } + + private static void feedPrimitive(Group group, int fieldIndex, int valueIndex, PrimitiveType type, + PrimitiveConverter converter) { + switch (type.getPrimitiveTypeName()) { + case BOOLEAN: + converter.addBoolean(group.getBoolean(fieldIndex, valueIndex)); + break; + case INT32: + converter.addInt(group.getInteger(fieldIndex, valueIndex)); + break; + case INT64: + converter.addLong(group.getLong(fieldIndex, valueIndex)); + break; + case FLOAT: + converter.addFloat(group.getFloat(fieldIndex, valueIndex)); + break; + case DOUBLE: + converter.addDouble(group.getDouble(fieldIndex, valueIndex)); + break; + case BINARY: + case FIXED_LEN_BYTE_ARRAY: + converter.addBinary(group.getBinary(fieldIndex, valueIndex)); + break; + default: + throw new IllegalArgumentException("Unsupported Parquet VARIANT physical type: " + type); + } + } + + private static final class BuilderHolder { + private VariantBuilder _builder; + + private void reset() { + _builder = null; + } + + private void setMetadata(ByteBuffer metadata) { + _builder = new VariantBuilder(new ImmutableMetadata(metadata)); + } + + private void build(Consumer consumer) { + if (_builder == null) { + throw new IllegalStateException("Cannot build Parquet VARIANT: metadata has not been read"); + } + consumer.accept(_builder); + } + + private Variant finish() { + if (_builder == null) { + throw new IllegalStateException("Cannot build Parquet VARIANT: missing metadata"); + } + _builder.appendNullIfEmpty(); + return _builder.build(); + } + } +} diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java new file mode 100644 index 000000000000..e6a2058f238d --- /dev/null +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java @@ -0,0 +1,876 @@ +/** + * 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.plugin.inputformat.parquet; + +import java.io.File; +import java.io.IOException; +import java.lang.reflect.Field; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.function.Consumer; +import org.apache.commons.io.FileUtils; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.format.converter.ParquetMetadataConverter; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; +import org.apache.parquet.variant.ImmutableMetadata; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantArrayBuilder; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.parquet.variant.VariantObjectBuilder; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.utils.VariantEnvelope; +import org.testng.annotations.AfterClass; +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.assertTrue; +import static org.testng.Assert.expectThrows; + + +/** + * End-to-end coverage for Parquet `VARIANT(1)` reconstruction and reader selection. + */ +public class ParquetVariantRecordReaderTest { + private static final String VARIANT_FIELD = "variant_col"; + + private final File _tempDir = new File(FileUtils.getTempDirectory(), getClass().getSimpleName()); + + @AfterClass + public void cleanUp() { + FileUtils.deleteQuietly(_tempDir); + } + + @Test + public void testUnshreddedShreddedNullProjectionAndRewind() + throws Exception { + Variant objectVariant = objectVariant("name", "pinot"); + File dataFile = writeScalarVariantFile(objectVariant); + assertFalse(ParquetUtils.hasAvroSchemaInFileMetadata(new Path(dataFile.getAbsolutePath()))); + + assertScalarRows(new ParquetNativeRecordReader(), dataFile); + ParquetRecordReader autoSelectingReader = new ParquetRecordReader(); + assertScalarRows(autoSelectingReader, dataFile); + assertFalse(autoSelectingReader.useAvroParquetRecordReader()); + } + + @Test + public void testPartiallyShreddedObject() + throws Exception { + Variant expected = objectVariant("cold", "kept", "hot", 7); + VariantBuilder baseBuilder = new VariantBuilder(new ImmutableMetadata(expected.getMetadataBuffer())); + VariantObjectBuilder baseObject = baseBuilder.startObject(); + baseObject.appendKey("cold"); + baseObject.appendString("kept"); + baseObject.appendKey("hot"); + baseObject.appendInt(1); + baseBuilder.endObject(); + Variant base = baseBuilder.build(); + + File dataFile = writePartiallyShreddedFile(expected.getMetadataBuffer(), base.getValueBuffer()); + assertShreddedObjectRows(new ParquetNativeRecordReader(), dataFile); + } + + @Test + public void testUnshreddedValuePreservesEncodedBuffers() { + Variant expected = objectVariant("name", "pinot"); + byte[] expectedMetadata = remainingBytes(expected.getMetadataBuffer()); + byte[] metadataBacking = addSentinels(expectedMetadata); + byte[] expectedValue = remainingBytes(expected.getValueBuffer()); + byte[] valueBacking = addSentinels(expectedValue); + MessageType schema = MessageTypeParser.parseMessageType( + "message direct_variant {" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " optional int32 typed_value;" + + " }" + + "}"); + Group variantGroup = new SimpleGroupFactory(schema).newGroup().addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteArray(metadataBacking, 1, expectedMetadata.length)) + .append("value", Binary.fromConstantByteArray(valueBacking, 1, expectedValue.length)); + + ParquetVariantConverter converter = + ParquetVariantConverter.createTopLevelVariantConverters(schema)[schema.getFieldIndex(VARIANT_FIELD)]; + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(converter.convert(variantGroup)); + assertEquals(remainingBytes(decoded.getMetadata()), expectedMetadata); + assertEquals(remainingBytes(decoded.getValue()), expectedValue); + } + + @Test + public void testUnshreddedValueSupportsDirectAndReadOnlyBuffers() { + Variant expected = objectVariant("name", "pinot"); + byte[] expectedMetadata = remainingBytes(expected.getMetadataBuffer()); + byte[] metadataBacking = addSentinels(expectedMetadata); + byte[] expectedValue = remainingBytes(expected.getValueBuffer()); + byte[] valueBacking = addSentinels(expectedValue); + MessageType schema = MessageTypeParser.parseMessageType( + "message direct_variant {" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " }" + + "}"); + + ByteBuffer directMetadata = ByteBuffer.allocateDirect(metadataBacking.length); + directMetadata.put(metadataBacking).flip(); + directMetadata.position(1).limit(1 + expectedMetadata.length); + ByteBuffer directValue = ByteBuffer.allocateDirect(valueBacking.length); + directValue.put(valueBacking).flip(); + directValue.position(1).limit(1 + expectedValue.length); + assertUnshreddedBufferRoundTrip(schema, directMetadata, directValue, expectedMetadata, expectedValue); + + ByteBuffer readOnlyMetadata = ByteBuffer.wrap(metadataBacking).asReadOnlyBuffer(); + readOnlyMetadata.position(1).limit(1 + expectedMetadata.length); + ByteBuffer readOnlyValue = ByteBuffer.wrap(valueBacking).asReadOnlyBuffer(); + readOnlyValue.position(1).limit(1 + expectedValue.length); + assertUnshreddedBufferRoundTrip(schema, readOnlyMetadata, readOnlyValue, expectedMetadata, expectedValue); + } + + @Test + public void testStandaloneExtractorOwnsSchemaBoundVariantConverters() { + Variant expected = objectVariant("name", "pinot"); + MessageType schema = MessageTypeParser.parseMessageType( + "message direct_variant {" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " }" + + "}"); + Group root = new SimpleGroupFactory(schema).newGroup(); + root.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(expected.getMetadataBuffer())) + .append("value", Binary.fromConstantByteBuffer(expected.getValueBuffer())); + + ParquetNativeRecordExtractorConfig config = new ParquetNativeRecordExtractorConfig(); + config.setParquetSchema(schema); + ParquetNativeRecordExtractor extractor = new ParquetNativeRecordExtractor(); + extractor.init(Set.of(VARIANT_FIELD), config); + + GenericRow row = extractor.extract(root, new GenericRow()); + Variant actual = decode((byte[]) row.getValue(VARIANT_FIELD)); + assertEquals(actual.getFieldByKey("name").getString(), "pinot"); + + ParquetNativeRecordExtractor compatibilityExtractor = new ParquetNativeRecordExtractor(); + compatibilityExtractor.init(Set.of(VARIANT_FIELD), null); + GenericRow compatibilityRow = compatibilityExtractor.extract(root, new GenericRow()); + assertEquals(decode((byte[]) compatibilityRow.getValue(VARIANT_FIELD)).getFieldByKey("name").getString(), + "pinot"); + } + + @Test + public void testShreddedPrimitiveTypes() + throws Exception { + File dataFile = writeShreddedPrimitiveFile(); + assertShreddedPrimitiveRow(new ParquetNativeRecordReader(), dataFile); + } + + @Test + public void testShreddedArrayWithRepeatedAndNullElements() + throws Exception { + File dataFile = writeShreddedArrayFile(); + try (ParquetNativeRecordReader reader = new ParquetNativeRecordReader()) { + reader.init(dataFile, null, null); + List rows = readAll(reader); + assertEquals(rows.size(), 1); + + Variant array = decode((byte[]) rows.get(0).getValue(VARIANT_FIELD)); + assertEquals(array.getType(), Variant.Type.ARRAY); + assertEquals(array.numArrayElements(), 3); + assertEquals(array.getElementAtIndex(0).getInt(), 7); + assertEquals(array.getElementAtIndex(1).getType(), Variant.Type.NULL); + assertEquals(array.getElementAtIndex(2).getInt(), 9); + } + } + + @Test + public void testAutomaticReaderReinitializesFromVariantToOrdinary() + throws Exception { + File variantFile = writeScalarVariantFile(objectVariant("name", "pinot")); + File ordinaryFile = getResourceFile("data-avro.parquet"); + try (ParquetRecordReader reader = new ParquetRecordReader()) { + reader.init(variantFile, Set.of("id", VARIANT_FIELD), null); + assertFalse(reader.useAvroParquetRecordReader()); + assertTrue(reader.hasNext()); + assertEquals(reader.next().getValue("id"), 1); + + reader.init(ordinaryFile, null, null); + assertTrue(reader.useAvroParquetRecordReader()); + assertTrue(reader.hasNext()); + } + } + + @Test + public void testAutomaticReaderReinitializesFromOrdinaryToVariant() + throws Exception { + File variantFile = writeScalarVariantFile(objectVariant("name", "pinot")); + File ordinaryFile = getResourceFile("data-avro.parquet"); + try (ParquetRecordReader reader = new ParquetRecordReader()) { + reader.init(ordinaryFile, null, null); + assertTrue(reader.useAvroParquetRecordReader()); + assertTrue(reader.hasNext()); + reader.next(); + + reader.init(variantFile, Set.of("id", VARIANT_FIELD), null); + assertFalse(reader.useAvroParquetRecordReader()); + assertTrue(reader.hasNext()); + assertEquals(reader.next().getValue("id"), 1); + } + } + + @Test + public void testNativeReaderReinitializesOnlyAfterCandidateSucceeds() + throws Exception { + File firstFile = writeScalarVariantFile(objectVariant("name", "first"), "native-first.parquet", Map.of()); + File secondFile = writeScalarVariantFile(objectVariant("name", "second"), "native-second.parquet", Map.of()); + try (ParquetNativeRecordReader reader = new ParquetNativeRecordReader()) { + reader.init(firstFile, Set.of("id"), null); + assertEquals(reader.next().getValue("id"), 1); + + reader.init(secondFile, Set.of(VARIANT_FIELD), null); + List rows = readAll(reader); + assertEquals(rows.size(), 4); + assertFalse(rows.get(2).getFieldToValueMap().containsKey("id")); + Variant second = decode((byte[]) rows.get(2).getValue(VARIANT_FIELD)); + assertEquals(second.getFieldByKey("name").getString(), "second"); + } + } + + @Test + public void testNativeReaderFailedReinitializationRetainsPreviousState() + throws Exception { + File variantFile = writeScalarVariantFile(objectVariant("name", "pinot")); + File invalidFile = prepareFile("invalid-native.parquet"); + FileUtils.writeByteArrayToFile(invalidFile, new byte[]{'N', 'O', 'P', 'E'}); + try (ParquetNativeRecordReader reader = new ParquetNativeRecordReader()) { + reader.init(variantFile, Set.of("id", VARIANT_FIELD), null); + assertEquals(reader.next().getValue("id"), 1); + + expectThrows(Exception.class, () -> reader.init(invalidFile, null, null)); + assertTrue(reader.hasNext()); + assertEquals(reader.next().getValue("id"), 2); + } + } + + @Test + public void testNativeReaderPublishesReplacementBeforePreviousCloseFailure() + throws Exception { + File variantFile = writeScalarVariantFile(objectVariant("name", "replacement")); + FailOnceOnCloseParquetFileReader previousReader = new FailOnceOnCloseParquetFileReader(variantFile); + + try { + try (ParquetNativeRecordReader reader = new ParquetNativeRecordReader()) { + Field readerField = ParquetNativeRecordReader.class.getDeclaredField("_parquetFileReader"); + readerField.setAccessible(true); + readerField.set(reader, previousReader); + + IOException exception = + expectThrows(IOException.class, () -> reader.init(variantFile, Set.of("id", VARIANT_FIELD), null)); + assertEquals(exception.getMessage(), "previous close failed"); + + assertTrue(reader.hasNext()); + assertEquals(reader.next().getValue("id"), 1); + } + } finally { + previousReader.close(); + } + } + + @Test + public void testAvroMetadataSelectionRemainsBackwardCompatibleWithExplicitNativeOptIn() + throws Exception { + File variantFile = writeScalarVariantFileWithAvroMetadata(objectVariant("name", "pinot")); + ParquetRecordReaderConfig forceAvro = new ParquetRecordReaderConfig(); + forceAvro.setUseParquetAvroRecordReader(true); + + try (ParquetRecordReader reader = new ParquetRecordReader()) { + reader.init(variantFile, Set.of("id"), null); + assertTrue(reader.useAvroParquetRecordReader()); + assertEquals(reader.next().getValue("id"), 1); + + reader.init(variantFile, Set.of("id"), forceAvro); + assertTrue(reader.useAvroParquetRecordReader()); + assertEquals(reader.next().getValue("id"), 1); + } + + ParquetRecordReaderConfig forceNative = new ParquetRecordReaderConfig(); + forceNative.setUseParquetNativeRecordReader(true); + try (ParquetRecordReader reader = new ParquetRecordReader()) { + reader.init(variantFile, Set.of("id", VARIANT_FIELD), forceNative); + assertFalse(reader.useAvroParquetRecordReader()); + assertScalarRows(readAll(reader)); + } + } + + @Test + public void testFailedReinitializationRetainsPreviousDelegate() + throws Exception { + File ordinaryFile = getResourceFile("data-avro.parquet"); + File invalidFile = prepareFile("invalid.parquet"); + FileUtils.writeByteArrayToFile(invalidFile, new byte[]{'N', 'O', 'P', 'E'}); + try (ParquetRecordReader reader = new ParquetRecordReader()) { + reader.init(ordinaryFile, null, null); + expectThrows(Exception.class, () -> reader.init(invalidFile, null, null)); + assertTrue(reader.useAvroParquetRecordReader()); + assertTrue(reader.hasNext()); + reader.next(); + } + try (ParquetAvroRecordReader reader = new ParquetAvroRecordReader()) { + reader.init(ordinaryFile, null, null); + expectThrows(Exception.class, () -> reader.init(invalidFile, null, null)); + assertTrue(reader.hasNext()); + reader.next(); + } + } + + @Test + public void testMalformedShreddedRowsAdvanceToNextRowAndEof() + throws Exception { + File dataFile = writeMalformedShreddedRows(); + try (ParquetNativeRecordReader reader = new ParquetNativeRecordReader()) { + reader.init(dataFile, null, null); + + assertEquals(reader.next().getValue("id"), 1); + assertTrue(reader.hasNext()); + expectThrows(RuntimeException.class, reader::next); + + assertTrue(reader.hasNext()); + assertEquals(reader.next().getValue("id"), 3); + assertTrue(reader.hasNext()); + expectThrows(RuntimeException.class, reader::next); + assertFalse(reader.hasNext(), "A terminal malformed row must still advance the physical reader to EOF"); + } + } + + @Test + public void testVariantDetectionAndUnsupportedShapes() { + MessageType lookalike = MessageTypeParser.parseMessageType( + "message lookalike { optional group variant_col { required binary metadata; optional binary value; } }"); + assertTrue(ParquetVariantConverter.validateAndGetTopLevelVariantFields(lookalike).isEmpty()); + + MessageType unsupportedVersion = MessageTypeParser.parseMessageType( + "message unsupported { optional group variant_col (VARIANT(2)) {" + + " required binary metadata; optional binary value; } }"); + UnsupportedOperationException versionException = expectThrows(UnsupportedOperationException.class, + () -> ParquetVariantConverter.validateAndGetTopLevelVariantFields(unsupportedVersion)); + assertTrue(versionException.getMessage().contains("spec version")); + + MessageType repeated = Types.buildMessage() + .addField(Types.buildGroup(Type.Repetition.REPEATED) + .as(LogicalTypeAnnotation.variantType((byte) 1)) + .required(PrimitiveType.PrimitiveTypeName.BINARY) + .named("metadata") + .optional(PrimitiveType.PrimitiveTypeName.BINARY) + .named("value") + .named(VARIANT_FIELD)) + .named("repeated_variant"); + UnsupportedOperationException repeatedException = expectThrows(UnsupportedOperationException.class, + () -> ParquetVariantConverter.validateAndGetTopLevelVariantFields(repeated)); + assertTrue(repeatedException.getMessage().contains("Repeated")); + + MessageType missingMetadata = MessageTypeParser.parseMessageType( + "message malformed { optional group variant_col (VARIANT(1)) { optional binary value; } }"); + IllegalArgumentException metadataException = expectThrows(IllegalArgumentException.class, + () -> ParquetVariantConverter.validateAndGetTopLevelVariantFields(missingMetadata)); + assertTrue(metadataException.getMessage().contains("metadata")); + + MessageType nested = MessageTypeParser.parseMessageType( + "message nested { optional group wrapper { optional group variant_col (VARIANT(1)) {" + + " required binary metadata; optional binary value; } } }"); + UnsupportedOperationException nestedException = expectThrows(UnsupportedOperationException.class, + () -> ParquetVariantConverter.validateAndGetTopLevelVariantFields(nested)); + assertTrue(nestedException.getMessage().contains("Nested")); + + MessageType unsupportedInt96 = MessageTypeParser.parseMessageType( + "message unsupported_int96 { optional group variant_col (VARIANT(1)) {" + + " required binary metadata; optional int96 typed_value; } }"); + UnsupportedOperationException int96Exception = expectThrows(UnsupportedOperationException.class, + () -> ParquetVariantConverter.createTopLevelVariantConverters(unsupportedInt96)); + assertTrue(int96Exception.getMessage().contains("Unsupported shredded value type")); + assertTrue(int96Exception.getMessage().matches("(?i).*int96.*")); + } + + private void assertScalarRows(RecordReader reader, File dataFile) + throws IOException { + try (reader) { + reader.init(dataFile, Set.of("id", VARIANT_FIELD), null); + List firstPass = readAll(reader); + assertScalarRows(firstPass); + + reader.rewind(); + List secondPass = readAll(reader); + assertScalarRows(secondPass); + } + } + + private void assertScalarRows(List rows) { + assertEquals(rows.size(), 4); + for (GenericRow row : rows) { + assertFalse(row.getFieldToValueMap().containsKey("note")); + } + + assertEquals(rows.get(0).getValue("id"), 1); + assertNull(rows.get(0).getValue(VARIANT_FIELD)); + + assertEquals(rows.get(1).getValue("id"), 2); + byte[] encodedNull = (byte[]) rows.get(1).getValue(VARIANT_FIELD); + assertTrue(VariantEnvelope.isEnvelope(encodedNull)); + assertEquals(decode(encodedNull).getType(), Variant.Type.NULL); + + assertEquals(rows.get(2).getValue("id"), 3); + Variant object = decode((byte[]) rows.get(2).getValue(VARIANT_FIELD)); + assertEquals(object.getType(), Variant.Type.OBJECT); + assertEquals(object.getFieldByKey("name").getString(), "pinot"); + + assertEquals(rows.get(3).getValue("id"), 4); + Variant shredded = decode((byte[]) rows.get(3).getValue(VARIANT_FIELD)); + assertEquals(shredded.getType(), Variant.Type.INT); + assertEquals(shredded.getInt(), 42); + } + + private void assertShreddedObjectRows(RecordReader reader, File dataFile) + throws IOException { + try (reader) { + reader.init(dataFile, null, null); + assertShreddedObjectRows(readAll(reader)); + reader.rewind(); + assertShreddedObjectRows(readAll(reader)); + } + } + + private void assertShreddedObjectRows(List rows) { + assertEquals(rows.size(), 2); + + Variant fullyShredded = decode((byte[]) rows.get(0).getValue(VARIANT_FIELD)); + assertEquals(fullyShredded.getType(), Variant.Type.OBJECT); + assertNull(fullyShredded.getFieldByKey("cold")); + assertEquals(fullyShredded.getFieldByKey("hot").getInt(), 9); + + Variant partiallyShredded = decode((byte[]) rows.get(1).getValue(VARIANT_FIELD)); + assertEquals(partiallyShredded.getType(), Variant.Type.OBJECT); + assertEquals(partiallyShredded.getFieldByKey("cold").getString(), "kept"); + assertEquals(partiallyShredded.getFieldByKey("hot").getInt(), 7); + } + + private void assertShreddedPrimitiveRow(RecordReader reader, File dataFile) + throws IOException { + try (reader) { + reader.init(dataFile, null, null); + List rows = readAll(reader); + assertEquals(rows.size(), 1); + GenericRow row = rows.get(0); + + Variant booleanValue = decode((byte[]) row.getValue("variant_boolean")); + assertEquals(booleanValue.getType(), Variant.Type.BOOLEAN); + assertTrue(booleanValue.getBoolean()); + + Variant longValue = decode((byte[]) row.getValue("variant_long")); + assertEquals(longValue.getType(), Variant.Type.LONG); + assertEquals(longValue.getLong(), 9_876_543_210L); + + Variant floatValue = decode((byte[]) row.getValue("variant_float")); + assertEquals(floatValue.getType(), Variant.Type.FLOAT); + assertEquals(floatValue.getFloat(), 1.25f); + + Variant doubleValue = decode((byte[]) row.getValue("variant_double")); + assertEquals(doubleValue.getType(), Variant.Type.DOUBLE); + assertEquals(doubleValue.getDouble(), 2.5d); + + Variant binaryValue = decode((byte[]) row.getValue("variant_binary")); + assertEquals(binaryValue.getType(), Variant.Type.BINARY); + assertEquals(remainingBytes(binaryValue.getBinary()), new byte[]{1, 2, 3}); + + Variant decimalValue = decode((byte[]) row.getValue("variant_decimal")); + assertEquals(decimalValue.getType(), Variant.Type.DECIMAL4); + assertEquals(decimalValue.getDecimal(), new BigDecimal("123.45")); + } + } + + private List readAll(RecordReader reader) + throws IOException { + List rows = new ArrayList<>(); + while (reader.hasNext()) { + rows.add(reader.next()); + } + return rows; + } + + private static File getResourceFile(String resourceName) { + return new File(Objects.requireNonNull( + ParquetVariantRecordReaderTest.class.getClassLoader().getResource(resourceName), + "Missing test resource: " + resourceName).getFile()); + } + + private File writeScalarVariantFile(Variant objectVariant) + throws IOException { + return writeScalarVariantFile(objectVariant, "scalar-variants.parquet", Map.of()); + } + + private File writeScalarVariantFileWithAvroMetadata(Variant objectVariant) + throws IOException { + String avroSchema = "{\"type\":\"record\",\"name\":\"scalar_variants\",\"fields\":[" + + "{\"name\":\"id\",\"type\":\"int\"}," + + "{\"name\":\"note\",\"type\":[\"null\",\"string\"],\"default\":null}," + + "{\"name\":\"variant_col\",\"type\":[\"null\"," + + "{\"type\":\"record\",\"name\":\"variant_value\",\"fields\":[" + + "{\"name\":\"value\",\"type\":[\"null\",\"bytes\"],\"default\":null}," + + "{\"name\":\"metadata\",\"type\":\"bytes\"}," + + "{\"name\":\"typed_value\",\"type\":[\"null\",\"int\"],\"default\":null}" + + "]}],\"default\":null}]}"; + return writeScalarVariantFile(objectVariant, "scalar-variants-with-avro-metadata.parquet", + Map.of("parquet.avro.schema", avroSchema)); + } + + private File writeScalarVariantFile(Variant objectVariant, String fileName, Map extraMetadata) + throws IOException { + MessageType schema = MessageTypeParser.parseMessageType( + "message scalar_variants {" + + " required int32 id;" + + " optional binary note (STRING);" + + " optional group variant_col (VARIANT(1)) {" + + " optional binary value;" + + " required binary metadata;" + + " optional int32 typed_value;" + + " }" + + "}"); + File dataFile = prepareFile(fileName); + try (ParquetWriter writer = newWriter(dataFile, schema, extraMetadata)) { + SimpleGroupFactory groups = new SimpleGroupFactory(schema); + + writer.write(groups.newGroup().append("id", 1).append("note", "absent")); + + VariantBuilder nullBuilder = new VariantBuilder(); + nullBuilder.appendNull(); + Variant nullVariant = nullBuilder.build(); + Group encodedNull = groups.newGroup().append("id", 2).append("note", "variant-null"); + encodedNull.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(nullVariant.getMetadataBuffer())); + writer.write(encodedNull); + + Group unshredded = groups.newGroup().append("id", 3).append("note", "unshredded"); + unshredded.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(objectVariant.getMetadataBuffer())) + .append("value", Binary.fromConstantByteBuffer(objectVariant.getValueBuffer())); + writer.write(unshredded); + + VariantBuilder scalarBuilder = new VariantBuilder(); + scalarBuilder.appendInt(42); + Variant scalarVariant = scalarBuilder.build(); + Group shredded = groups.newGroup().append("id", 4).append("note", "shredded"); + shredded.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(scalarVariant.getMetadataBuffer())) + .append("typed_value", 42); + writer.write(shredded); + } + return dataFile; + } + + private File writePartiallyShreddedFile(ByteBuffer metadata, ByteBuffer baseValue) + throws IOException { + MessageType schema = MessageTypeParser.parseMessageType( + "message partial_variant {" + + " required int32 id;" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " optional group typed_value {" + + " required group hot {" + + " optional binary value;" + + " optional int32 typed_value;" + + " }" + + " }" + + " }" + + "}"); + File dataFile = prepareFile("partial-variant.parquet"); + try (ParquetWriter writer = newWriter(dataFile, schema)) { + SimpleGroupFactory groups = new SimpleGroupFactory(schema); + + Group fullyShreddedRow = groups.newGroup().append("id", 1); + Group fullyShredded = fullyShreddedRow.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(metadata)); + fullyShredded.addGroup("typed_value").addGroup("hot").append("typed_value", 9); + writer.write(fullyShreddedRow); + + Group partiallyShreddedRow = groups.newGroup().append("id", 2); + Group partiallyShredded = partiallyShreddedRow.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(metadata)) + .append("value", Binary.fromConstantByteBuffer(baseValue)); + partiallyShredded.addGroup("typed_value").addGroup("hot").append("typed_value", 7); + writer.write(partiallyShreddedRow); + } + return dataFile; + } + + private File writeShreddedPrimitiveFile() + throws IOException { + MessageType schema = MessageTypeParser.parseMessageType( + "message shredded_primitives {" + + " required int32 id;" + + variantGroup("variant_boolean", "boolean typed_value") + + variantGroup("variant_long", "int64 typed_value") + + variantGroup("variant_float", "float typed_value") + + variantGroup("variant_double", "double typed_value") + + variantGroup("variant_binary", "binary typed_value") + + variantGroup("variant_decimal", "fixed_len_byte_array(4) typed_value (DECIMAL(9,2))") + + "}"); + File dataFile = prepareFile("shredded-primitives.parquet"); + try (ParquetWriter writer = newWriter(dataFile, schema)) { + Group row = new SimpleGroupFactory(schema).newGroup().append("id", 1); + addShreddedValue(row, "variant_boolean", variantMetadata(builder -> builder.appendBoolean(true))) + .append("typed_value", true); + addShreddedValue(row, "variant_long", variantMetadata(builder -> builder.appendLong(9_876_543_210L))) + .append("typed_value", 9_876_543_210L); + addShreddedValue(row, "variant_float", variantMetadata(builder -> builder.appendFloat(1.25f))) + .append("typed_value", 1.25f); + addShreddedValue(row, "variant_double", variantMetadata(builder -> builder.appendDouble(2.5d))) + .append("typed_value", 2.5d); + addShreddedValue(row, "variant_binary", + variantMetadata(builder -> builder.appendBinary(ByteBuffer.wrap(new byte[]{1, 2, 3})))) + .append("typed_value", Binary.fromConstantByteArray(new byte[]{1, 2, 3})); + addShreddedValue(row, "variant_decimal", + variantMetadata(builder -> builder.appendDecimal(new BigDecimal("123.45")))) + .append("typed_value", Binary.fromConstantByteArray(fixedLengthBytes(BigInteger.valueOf(12_345), 4))); + writer.write(row); + } + return dataFile; + } + + private File writeShreddedArrayFile() + throws IOException { + MessageType schema = MessageTypeParser.parseMessageType( + "message shredded_array {" + + " required int32 id;" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional group typed_value (LIST) {" + + " repeated group list {" + + " optional group element {" + + " optional binary value;" + + " optional int32 typed_value;" + + " }" + + " }" + + " }" + + " }" + + "}"); + VariantBuilder metadataBuilder = new VariantBuilder(); + VariantArrayBuilder arrayBuilder = metadataBuilder.startArray(); + arrayBuilder.appendInt(7); + arrayBuilder.appendNull(); + arrayBuilder.appendInt(9); + metadataBuilder.endArray(); + Variant metadataSource = metadataBuilder.build(); + + File dataFile = prepareFile("shredded-array.parquet"); + try (ParquetWriter writer = newWriter(dataFile, schema)) { + Group row = new SimpleGroupFactory(schema).newGroup().append("id", 1); + Group variant = row.addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(metadataSource.getMetadataBuffer())); + Group list = variant.addGroup("typed_value"); + list.addGroup("list").addGroup("element").append("typed_value", 7); + list.addGroup("list"); + list.addGroup("list").addGroup("element").append("typed_value", 9); + writer.write(row); + } + return dataFile; + } + + private File writeMalformedShreddedRows() + throws IOException { + MessageType schema = MessageTypeParser.parseMessageType( + "message malformed_shredded_rows {" + + " required int32 id;" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " optional int32 typed_value;" + + " }" + + "}"); + VariantBuilder builder = new VariantBuilder(); + builder.appendInt(7); + Variant valid = builder.build(); + File dataFile = prepareFile("malformed-shredded-rows.parquet"); + try (ParquetWriter writer = newWriter(dataFile, schema)) { + SimpleGroupFactory groups = new SimpleGroupFactory(schema); + for (int id = 1; id <= 4; id++) { + Group row = groups.newGroup().append("id", id); + Group variant = row.addGroup(VARIANT_FIELD); + if ((id & 1) == 0) { + variant.append("metadata", Binary.fromConstantByteArray(new byte[0])) + .append("typed_value", id); + } else { + variant.append("metadata", Binary.fromConstantByteBuffer(valid.getMetadataBuffer())) + .append("value", Binary.fromConstantByteBuffer(valid.getValueBuffer())); + } + writer.write(row); + } + } + return dataFile; + } + + private File prepareFile(String name) + throws IOException { + FileUtils.forceMkdir(_tempDir); + File dataFile = new File(_tempDir, name); + FileUtils.deleteQuietly(dataFile); + return dataFile; + } + + private ParquetWriter newWriter(File dataFile, MessageType schema) + throws IOException { + return newWriter(dataFile, schema, Map.of()); + } + + private ParquetWriter newWriter(File dataFile, MessageType schema, Map extraMetadata) + throws IOException { + return ExampleParquetWriter.builder(new Path(dataFile.getAbsolutePath())) + .withType(schema) + .withExtraMetaData(extraMetadata) + .build(); + } + + private static String variantGroup(String fieldName, String typedValueDeclaration) { + return " optional group " + fieldName + " (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " optional " + typedValueDeclaration + ";" + + " }"; + } + + private static Group addShreddedValue(Group row, String fieldName, ByteBuffer metadata) { + return row.addGroup(fieldName).append("metadata", Binary.fromConstantByteBuffer(metadata)); + } + + private static ByteBuffer variantMetadata(Consumer appender) { + VariantBuilder builder = new VariantBuilder(); + appender.accept(builder); + return builder.build().getMetadataBuffer(); + } + + private static byte[] fixedLengthBytes(BigInteger value, int length) { + byte[] source = value.toByteArray(); + byte[] result = new byte[length]; + byte signExtension = value.signum() < 0 ? (byte) 0xff : 0; + Arrays.fill(result, signExtension); + System.arraycopy(source, Math.max(0, source.length - length), result, Math.max(0, length - source.length), + Math.min(source.length, length)); + return result; + } + + private static Variant objectVariant(Object... keysAndValues) { + VariantBuilder builder = new VariantBuilder(); + VariantObjectBuilder object = builder.startObject(); + for (int i = 0; i < keysAndValues.length; i += 2) { + object.appendKey((String) keysAndValues[i]); + Object value = keysAndValues[i + 1]; + if (value instanceof String) { + object.appendString((String) value); + } else if (value instanceof Integer) { + object.appendInt((Integer) value); + } else { + throw new IllegalArgumentException("Unsupported test Variant value: " + value); + } + } + builder.endObject(); + return builder.build(); + } + + private static Variant decode(byte[] envelope) { + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + return new Variant(decoded.getValue(), decoded.getMetadata()); + } + + private static byte[] addSentinels(byte[] value) { + byte[] backing = new byte[value.length + 2]; + backing[0] = 99; + backing[backing.length - 1] = 98; + System.arraycopy(value, 0, backing, 1, value.length); + return backing; + } + + private static byte[] remainingBytes(ByteBuffer buffer) { + ByteBuffer view = buffer.duplicate(); + byte[] bytes = new byte[view.remaining()]; + view.get(bytes); + return bytes; + } + + private static void assertUnshreddedBufferRoundTrip(MessageType schema, ByteBuffer metadata, ByteBuffer value, + byte[] expectedMetadata, byte[] expectedValue) { + int metadataPosition = metadata.position(); + int metadataLimit = metadata.limit(); + int valuePosition = value.position(); + int valueLimit = value.limit(); + Group variantGroup = new SimpleGroupFactory(schema).newGroup().addGroup(VARIANT_FIELD) + .append("metadata", Binary.fromConstantByteBuffer(metadata)) + .append("value", Binary.fromConstantByteBuffer(value)); + + ParquetVariantConverter converter = + ParquetVariantConverter.createTopLevelVariantConverters(schema)[schema.getFieldIndex(VARIANT_FIELD)]; + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(converter.convert(variantGroup)); + assertEquals(remainingBytes(decoded.getMetadata()), expectedMetadata); + assertEquals(remainingBytes(decoded.getValue()), expectedValue); + assertEquals(metadata.position(), metadataPosition); + assertEquals(metadata.limit(), metadataLimit); + assertEquals(value.position(), valuePosition); + assertEquals(value.limit(), valueLimit); + } + + private static final class FailOnceOnCloseParquetFileReader extends ParquetFileReader { + private boolean _failOnClose = true; + + private FailOnceOnCloseParquetFileReader(File file) + throws IOException { + super(HadoopInputFile.fromPath(new Path(file.getAbsolutePath()), + ParquetUtils.getParquetHadoopConfiguration()), + ParquetReadOptions.builder().withMetadataFilter(ParquetMetadataConverter.NO_FILTER).build()); + } + + @Override + public void close() + throws IOException { + if (_failOnClose) { + _failOnClose = false; + throw new IOException("previous close failed"); + } + super.close(); + } + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java index 2363cb54794c..0a28bb724368 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java @@ -154,6 +154,13 @@ private static RexNode evaluateLiteralOnlyFunction(RexCall rexCall, RexBuilder r operand -> operand instanceof RexLiteral || (operand instanceof RexCall && ((RexCall) operand).getOperands() .stream().allMatch(op -> op instanceof RexLiteral))); + // Calcite does not support RexLiteral values with SqlTypeName.VARIANT. Keep these calls in the plan so that + // Pinot can evaluate them at runtime; otherwise constant folding parseJson(...) fails while trying to construct + // a VARIANT literal and prevents valid compositions such as variantGet(parseJson(...), ...). + if (rexCall.getType().getSqlTypeName() == SqlTypeName.VARIANT) { + return rexCall; + } + int numArguments = operands.size(); ColumnDataType[] argumentTypes = new ColumnDataType[numArguments]; Object[] arguments = new Object[numArguments]; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java index 773386badcad..ec7268e7c7c8 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql/fun/PinotOperatorTable.java @@ -334,8 +334,10 @@ public static PinotOperatorTable instance(boolean nullHandlingEnabled) { // Key is canonical name. Multiple operators can share the same name (e.g. binary "-" and unary "-"). private final Map> _operatorMap; private final List _operatorList; + private final boolean _nullHandlingEnabled; private PinotOperatorTable(boolean nullHandlingEnabled) { + _nullHandlingEnabled = nullHandlingEnabled; Map> operatorMap = new HashMap<>(); // Register standard operators @@ -467,6 +469,11 @@ public void lookupOperatorOverloads(SqlIdentifier opName, @Nullable SqlFunctionC if (!opName.isSimple()) { return; } + if (!_nullHandlingEnabled && TransformFunctionType.requiresNullHandling(opName.getSimple())) { + throw new IllegalStateException( + "VARIANT function " + opName.getSimple() + + " requires query null handling to be enabled; set enableNullHandling=true"); + } String canonicalName = FunctionRegistry.canonicalize(opName.getSimple()); List operators = _operatorMap.get(canonicalName); if (operators != null) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java index 432ea445a517..978dde5513ba 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java @@ -144,7 +144,7 @@ public static Literal toLiteral(RexExpression.Literal literal) { ColumnDataType dataType = literal.getDataType(); if (dataType == ColumnDataType.BOOLEAN) { value = BooleanUtils.isTrueInternalValue(value); - } else if (dataType == ColumnDataType.BYTES) { + } else if (dataType == ColumnDataType.BYTES || dataType == ColumnDataType.VARIANT) { value = ((ByteArray) value).getBytes(); } return RequestUtils.getLiteral(value); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java index 6198f1cbcb0f..92e3fe22a539 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java @@ -1085,6 +1085,11 @@ public static ColumnDataType convertToColumnDataType(RelDataType relDataType) { return isArray ? ColumnDataType.BYTES_ARRAY : ColumnDataType.BYTES; case UUID: return isArray ? ColumnDataType.UUID_ARRAY : ColumnDataType.UUID; + case VARIANT: + if (isArray) { + throw new IllegalArgumentException("ARRAY is not supported"); + } + return ColumnDataType.VARIANT; case MAP: return ColumnDataType.MAP; case OTHER: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java index 01ba0c1d8dc7..25bd660ac868 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java @@ -142,7 +142,8 @@ public static RexLiteral toRexLiteral(RelBuilder builder, RexExpression.Literal assert value != null; return rexBuilder.makeLiteral((String) value); } - case BYTES: { + case BYTES: + case VARIANT: { assert value != null; ByteArray byteArray = (ByteArray) value; byte[] bytes = byteArray.getBytes(); @@ -262,6 +263,7 @@ private static RexExpression.Literal fromRexLiteralValue(ColumnDataType dataType value = ((NlsString) value).getValue(); break; case BYTES: + case VARIANT: value = new ByteArray(((ByteString) value).getBytes()); break; default: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java index a5baa46428b2..09e0a735467a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java @@ -37,6 +37,7 @@ import org.apache.pinot.query.planner.plannode.TableScanNode; import org.apache.pinot.query.planner.plannode.ValueNode; import org.apache.pinot.query.planner.validation.ArrayToMvValidationVisitor; +import org.apache.pinot.query.planner.validation.VariantTypeValidationVisitor; import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.query.routing.WorkerMetadata; @@ -152,12 +153,13 @@ private static void trackEmptyLeafStages(DispatchablePlanContext context) { } } - /// Run validations on the plan. Since there is only one validator right now, don't try to over-engineer it. + /// Runs validations on the plan. private void runValidations(PlanFragment planFragment, DispatchablePlanContext context) { PlanNode rootPlanNode = planFragment.getFragmentRoot(); boolean isIntermediateStage = context.getDispatchablePlanMetadataMap().get(rootPlanNode.getStageId()).getScannedTables().isEmpty(); rootPlanNode.visit(ArrayToMvValidationVisitor.INSTANCE, isIntermediateStage); + rootPlanNode.visit(VariantTypeValidationVisitor.INSTANCE, null); for (PlanFragment child : planFragment.getChildren()) { runValidations(child, context); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java index cc9d44165932..f600abecf7e2 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java @@ -345,6 +345,11 @@ public static ColumnDataType convertToColumnDataType(RelDataType relDataType) { return isArray ? ColumnDataType.BYTES_ARRAY : ColumnDataType.BYTES; case UUID: return isArray ? ColumnDataType.UUID_ARRAY : ColumnDataType.UUID; + case VARIANT: + if (isArray) { + throw new IllegalArgumentException("ARRAY is not supported"); + } + return ColumnDataType.VARIANT; case MAP: return ColumnDataType.MAP; case OTHER: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java index 060ca8447214..fc36c840dffd 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java @@ -182,6 +182,8 @@ public static ColumnDataType convertColumnDataType(Expressions.ColumnDataType da return ColumnDataType.BYTES; case UUID: return ColumnDataType.UUID; + case VARIANT: + return ColumnDataType.VARIANT; case INT_ARRAY: return ColumnDataType.INT_ARRAY; case LONG_ARRAY: @@ -208,6 +210,10 @@ public static ColumnDataType convertColumnDataType(Expressions.ColumnDataType da return ColumnDataType.OBJECT; case UNKNOWN: return ColumnDataType.UNKNOWN; + case UNRECOGNIZED: + throw new IllegalArgumentException( + "Unrecognized query-plan ColumnDataType received from a peer node. Upgrade all brokers and servers " + + "before querying columns whose logical type was introduced by the newer node."); default: throw new IllegalStateException("Unsupported proto ColumnDataType: " + dataType); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java index 7b0c79ec5d81..f2ef154a274e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java @@ -170,6 +170,8 @@ public static Expressions.ColumnDataType convertColumnDataType(ColumnDataType da return Expressions.ColumnDataType.BYTES; case UUID: return Expressions.ColumnDataType.UUID; + case VARIANT: + return Expressions.ColumnDataType.VARIANT; case MAP: return Expressions.ColumnDataType.MAP; case INT_ARRAY: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java new file mode 100644 index 000000000000..ff6405ce5cd4 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java @@ -0,0 +1,201 @@ +/** + * 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.query.planner.validation; + +import java.util.List; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.query.planner.plannode.AggregateNode; +import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.PlanNodeVisitor; +import org.apache.pinot.query.planner.plannode.SetOpNode; +import org.apache.pinot.query.planner.plannode.SortNode; +import org.apache.pinot.query.planner.plannode.WindowNode; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; + + +/** + * Rejects operations that would otherwise assign physical byte ordering, equality, or hashing semantics to a raw + * VARIANT value. + */ +public class VariantTypeValidationVisitor extends PlanNodeVisitor.DepthFirstVisitor { + public static final VariantTypeValidationVisitor INSTANCE = new VariantTypeValidationVisitor(); + + private VariantTypeValidationVisitor() { + } + + @Override + protected boolean traverseStageBoundary() { + return false; + } + + @Override + public Void visitAggregate(AggregateNode node, Void context) { + List inputs = node.getInputs(); + if (inputs.size() == 1) { + validateAggregateInputs(node, inputs.get(0).getDataSchema()); + } + return super.visitAggregate(node, context); + } + + /** + * Validates aggregate operands against their logical input schema. + * + *

This method is also invoked by the runtime as a defensive check for plans that did not pass through the + * current broker planner. + */ + public static void validateAggregateInputs(AggregateNode node, DataSchema inputSchema) { + validateAggregateInputs(node.getAggCalls(), inputSchema); + } + + /** + * Validates aggregate or window-function operands against their logical input schema. + */ + public static void validateAggregateInputs(List aggCalls, DataSchema inputSchema) { + for (RexExpression.FunctionCall aggCall : aggCalls) { + if (isRawVariantIndependent(aggCall)) { + continue; + } + for (RexExpression operand : aggCall.getFunctionOperands()) { + if (!getLogicalType(operand, inputSchema).supportsDirectAggregation()) { + throw unsupported("Aggregate function " + aggCall.getFunctionName()); + } + } + } + } + + @Override + public Void visitSort(SortNode node, Void context) { + DataSchema dataSchema = node.getDataSchema(); + for (RelFieldCollation collation : node.getCollations()) { + int fieldIndex = collation.getFieldIndex(); + if (!dataSchema.getColumnDataType(fieldIndex).supportsOrdering()) { + throw unsupported("ORDER BY"); + } + } + return super.visitSort(node, context); + } + + @Override + public Void visitSetOp(SetOpNode node, Void context) { + if (!(node.getSetOpType() == SetOpNode.SetOpType.UNION && node.isAll())) { + for (DataSchema.ColumnDataType dataType : node.getDataSchema().getColumnDataTypes()) { + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw unsupported(node.explain().replace('_', ' ')); + } + } + } + return super.visitSetOp(node, context); + } + + @Override + public Void visitJoin(JoinNode node, Void context) { + List inputs = node.getInputs(); + if (inputs.size() == 2) { + validateJoinInputs(node, inputs.get(0).getDataSchema(), inputs.get(1).getDataSchema()); + } + return super.visitJoin(node, context); + } + + /** + * Validates equality/hash join keys against both logical input schemas. + * + *

The runtime also invokes this for mixed-version plans, including LOOKUP joins that do not construct a + * {@code HashJoinOperator}. + */ + public static void validateJoinInputs(JoinNode node, DataSchema leftSchema, DataSchema rightSchema) { + for (int leftKey : node.getLeftKeys()) { + DataSchema.ColumnDataType dataType = leftSchema.getColumnDataType(leftKey); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw unsupported("JOIN keys"); + } + } + for (int rightKey : node.getRightKeys()) { + DataSchema.ColumnDataType dataType = rightSchema.getColumnDataType(rightKey); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw unsupported("JOIN keys"); + } + } + if (node.getJoinStrategy() == JoinNode.JoinStrategy.ASOF) { + RexExpression.FunctionCall matchCondition = (RexExpression.FunctionCall) node.getMatchCondition(); + List matchKeys = matchCondition.getFunctionOperands(); + int leftMatchKey = ((RexExpression.InputRef) matchKeys.get(0)).getIndex(); + int rightMatchKey = ((RexExpression.InputRef) matchKeys.get(1)).getIndex() - leftSchema.size(); + if (!leftSchema.getColumnDataType(leftMatchKey).supportsOrdering() + || !rightSchema.getColumnDataType(rightMatchKey).supportsOrdering()) { + throw unsupported("ASOF JOIN MATCH_CONDITION"); + } + } + } + + @Override + public Void visitWindow(WindowNode node, Void context) { + List inputs = node.getInputs(); + if (inputs.size() == 1) { + validateWindowInputs(node, inputs.get(0).getDataSchema()); + } + return super.visitWindow(node, context); + } + + /** + * Validates window partition keys, ordering keys, and function operands against their logical input schema. + * + *

This method is also invoked by the runtime as a defensive check for plans that did not pass through the + * current broker planner. + */ + public static void validateWindowInputs(WindowNode node, DataSchema inputSchema) { + for (int key : node.getKeys()) { + DataSchema.ColumnDataType dataType = inputSchema.getColumnDataType(key); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw unsupported("Window PARTITION BY"); + } + } + for (RelFieldCollation collation : node.getCollations()) { + if (!inputSchema.getColumnDataType(collation.getFieldIndex()).supportsOrdering()) { + throw unsupported("Window ORDER BY"); + } + } + validateAggregateInputs(node.getAggCalls(), inputSchema); + } + + private static boolean isRawVariantIndependent(RexExpression.FunctionCall aggCall) { + return !aggCall.isDistinct() && aggCall.getFunctionName().equalsIgnoreCase("COUNT"); + } + + private static DataSchema.ColumnDataType getLogicalType(RexExpression expression, DataSchema inputSchema) { + if (expression instanceof RexExpression.InputRef) { + return inputSchema.getColumnDataType(((RexExpression.InputRef) expression).getIndex()); + } + if (expression instanceof RexExpression.Literal) { + return ((RexExpression.Literal) expression).getDataType(); + } + if (expression instanceof RexExpression.FunctionCall) { + return ((RexExpression.FunctionCall) expression).getDataType(); + } + throw new IllegalStateException("Unsupported aggregate operand: " + expression.getClass().getName()); + } + + private static QueryException unsupported(String operation) { + return new QueryException(QueryErrorCode.QUERY_PLANNING, + operation + " does not support raw VARIANT values; extract a typed path with variantGet first"); + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java index 3f04dce72832..cd1ac91a0aa3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/type/TypeFactory.java @@ -102,6 +102,8 @@ private static SqlTypeName getSqlTypeName(FieldSpec fieldSpec) { return SqlTypeName.VARCHAR; case UUID: return SqlTypeName.UUID; + case VARIANT: + return SqlTypeName.VARIANT; case BYTES: return SqlTypeName.VARBINARY; case BIG_DECIMAL: diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java index 4416e47bb52e..0714a38393ec 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java @@ -166,6 +166,40 @@ public void testJsonExtractScalarAcceptsFoldableJsonPath() { } } + @Test + public void testVariantReturningFunctionIsNotConstantFolded() { + // Calcite cannot represent a VARIANT RexLiteral. The literal-only evaluation rule must leave parseJson as a + // runtime call so it can be composed with functions that consume VARIANT. + DispatchableSubPlan dispatchableSubPlan = _queryEnvironment.planQuery( + "SELECT variant_get(parse_json('{\"answer\":42}'), '$.answer', 'INT') FROM a"); + assertNotNull(dispatchableSubPlan); + } + + @Test + public void testVariantFunctionsRequireQueryNullHandling() { + QueryEnvironment nullHandlingDisabled = getQueryEnvironment( + 13, 11, 12, TABLE_SCHEMAS, SERVER1_SEGMENTS, SERVER2_SEGMENTS, PARTITIONED_SEGMENTS_MAP, false); + List expressions = List.of( + "variant_get(col1, '$.value')", + "try_variant_get(col1, '$.value')", + "variant_exists(col1, '$.value')", + "is_variant_null(col1)", + "variant_type_of(col1)", + "variant_to_json(col1)", + "parse_json(col1)", + "parse_json_to_variant(col1)", + "try_parse_json(col1)", + "try_parse_json_to_variant(col1)"); + for (String expression : expressions) { + String query = "SELECT " + expression + " FROM a"; + Throwable thrown = expectThrows(RuntimeException.class, () -> nullHandlingDisabled.compile(query)); + assertTrue(Throwables.getStackTraceAsString(thrown).contains("requires query null handling"), + "Unexpected rejection for " + query + ": " + Throwables.getStackTraceAsString(thrown)); + } + + assertNotNull(nullHandlingDisabled.compile("SELECT col1 FROM a")); + } + @Test public void testPolymorphicArithmeticScalarFunctionsPlanQuery() { DispatchableSubPlan dispatchableSubPlan = _queryEnvironment.planQuery( diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java index fa42e993f273..07a4395cf816 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java @@ -304,6 +304,13 @@ protected Object[][] provideQueries() { public static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, int port2, Map schemaMap, Map> segmentMap1, Map> segmentMap2, @Nullable Map>>> partitionedSegmentsMap) { + return getQueryEnvironment(reducerPort, port1, port2, schemaMap, segmentMap1, segmentMap2, + partitionedSegmentsMap, true); + } + + public static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, int port2, + Map schemaMap, Map> segmentMap1, Map> segmentMap2, + @Nullable Map>>> partitionedSegmentsMap, boolean nullHandlingEnabled) { MockRoutingManagerFactory factory = new MockRoutingManagerFactory(port1, port2); for (Map.Entry entry : schemaMap.entrySet()) { factory.registerTable(entry.getValue(), entry.getKey()); @@ -342,7 +349,7 @@ public static QueryEnvironment getQueryEnvironment(int reducerPort, int port1, i RoutingManager routingManager = factory.buildRoutingManager(partitionInfoMap); TableCache tableCache = factory.buildTableCache(); return new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache, - new WorkerManager("Broker_localhost", "localhost", reducerPort, routingManager)); + new WorkerManager("Broker_localhost", "localhost", reducerPort, routingManager), nullHandlingEnabled); } /// JSON test case definition for query planner test cases. Tables and schemas will come from those already defined diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java index ffd8a62e76c8..4bc5ed32f133 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java @@ -22,7 +22,9 @@ import java.util.List; import java.util.Random; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.pinot.common.proto.Expressions; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.spi.utils.BooleanUtils; import org.apache.pinot.spi.utils.ByteArray; @@ -30,13 +32,16 @@ import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; public class RexExpressionSerDeTest { private static final List SUPPORTED_DATE_TYPES = List.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE, ColumnDataType.BIG_DECIMAL, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP, ColumnDataType.STRING, - ColumnDataType.BYTES, ColumnDataType.UUID, ColumnDataType.INT_ARRAY, ColumnDataType.LONG_ARRAY, + ColumnDataType.BYTES, ColumnDataType.UUID, ColumnDataType.VARIANT, ColumnDataType.INT_ARRAY, + ColumnDataType.LONG_ARRAY, ColumnDataType.FLOAT_ARRAY, ColumnDataType.DOUBLE_ARRAY, ColumnDataType.BOOLEAN_ARRAY, ColumnDataType.TIMESTAMP_ARRAY, ColumnDataType.STRING_ARRAY, ColumnDataType.UUID_ARRAY, ColumnDataType.UNKNOWN); @@ -104,6 +109,21 @@ public void testUuidLiteral() { new ByteArray(UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000")))); } + @Test + public void testVariantLiteral() { + byte[] variant = VariantUtils.parseJsonToVariant("{\"key\":\"value\"}"); + verifyLiteralSerDe(new RexExpression.Literal(ColumnDataType.VARIANT, new ByteArray(variant))); + } + + @Test + public void testUnknownPeerDataTypeFailsWithUpgradeGuidance() { + Expressions.Literal literal = + Expressions.Literal.newBuilder().setDataTypeValue(999).setNull(true).build(); + IllegalArgumentException exception = + expectThrows(IllegalArgumentException.class, () -> ProtoExpressionToRexExpression.convertLiteral(literal)); + assertTrue(exception.getMessage().contains("Upgrade all brokers and servers")); + } + @Test public void testIntArrayLiteral() { int[] values = new int[RANDOM.nextInt(10)]; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java new file mode 100644 index 000000000000..321694d718e5 --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.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.query.planner.validation; + +import java.util.List; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.query.planner.plannode.AggregateNode; +import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.SetOpNode; +import org.apache.pinot.query.planner.plannode.SortNode; +import org.apache.pinot.query.planner.plannode.ValueNode; +import org.apache.pinot.query.planner.plannode.WindowNode; +import org.apache.pinot.spi.exception.QueryException; +import org.testng.Assert; +import org.testng.annotations.Test; + + +public class VariantTypeValidationVisitorTest { + private static final DataSchema VARIANT_SCHEMA = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.VARIANT}); + private static final DataSchema TYPED_EXTRACTION_SCHEMA = + new DataSchema(new String[]{"typedPayload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + + @Test + public void testRejectsVariantOrderBy() { + SortNode sortNode = new SortNode(0, VARIANT_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(), + List.of(new RelFieldCollation(0)), 10, 0); + + QueryException exception = + Assert.expectThrows(QueryException.class, () -> sortNode.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("ORDER BY")); + } + + @Test + public void testRejectsEqualityDependentSetOperations() { + List unsupportedNodes = List.of( + setOp(SetOpNode.SetOpType.UNION, false), + setOp(SetOpNode.SetOpType.INTERSECT, false), + setOp(SetOpNode.SetOpType.INTERSECT, true), + setOp(SetOpNode.SetOpType.MINUS, false), + setOp(SetOpNode.SetOpType.MINUS, true)); + + for (SetOpNode node : unsupportedNodes) { + QueryException exception = + Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("raw VARIANT")); + } + } + + @Test + public void testAllowsVariantUnionAll() { + setOp(SetOpNode.SetOpType.UNION, true).visit(VariantTypeValidationVisitor.INSTANCE, null); + } + + @Test + public void testRejectsRawVariantJoinKeysForEveryStrategy() { + for (JoinNode.JoinStrategy strategy : JoinNode.JoinStrategy.values()) { + JoinNode leftVariant = join(strategy, VARIANT_SCHEMA, TYPED_EXTRACTION_SCHEMA); + QueryException exception = Assert.expectThrows(QueryException.class, + () -> leftVariant.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("JOIN keys")); + + JoinNode rightVariant = join(strategy, TYPED_EXTRACTION_SCHEMA, VARIANT_SCHEMA); + exception = Assert.expectThrows(QueryException.class, + () -> rightVariant.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("JOIN keys")); + } + } + + @Test + public void testRejectsRawVariantAsofMatchKeys() { + DataSchema leftVariant = new DataSchema(new String[]{"key", "match"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.VARIANT}); + DataSchema rightTyped = new DataSchema(new String[]{"key", "match"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.LONG}); + QueryException exception = Assert.expectThrows(QueryException.class, + () -> asofJoin(leftVariant, rightTyped).visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("ASOF JOIN MATCH_CONDITION")); + + DataSchema leftTyped = new DataSchema(new String[]{"key", "match"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.LONG}); + DataSchema rightVariant = new DataSchema(new String[]{"key", "match"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.VARIANT}); + exception = Assert.expectThrows(QueryException.class, + () -> asofJoin(leftTyped, rightVariant).visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("ASOF JOIN MATCH_CONDITION")); + } + + @Test + public void testRejectsAggregatesThatConsumeRawVariant() { + for (String functionName : List.of("SUM", "ANYVALUE", "DISTINCTCOUNTHLL")) { + AggregateNode node = aggregate(functionName, false, VARIANT_SCHEMA); + QueryException exception = + Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Aggregate function " + functionName)); + Assert.assertTrue(exception.getMessage().contains("variantGet")); + } + } + + @Test + public void testRejectsDistinctCountOfRawVariant() { + AggregateNode node = aggregate("COUNT", true, VARIANT_SCHEMA); + + QueryException exception = + Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Aggregate function COUNT")); + } + + @Test + public void testAllowsRawVariantCount() { + aggregate("COUNT", false, VARIANT_SCHEMA).visit(VariantTypeValidationVisitor.INSTANCE, null); + } + + @Test + public void testUsesLogicalTypeInsteadOfVariantStorageType() { + DataSchema bytesSchema = + new DataSchema(new String[]{"bytes"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BYTES}); + aggregate("ANYVALUE", false, bytesSchema).visit(VariantTypeValidationVisitor.INSTANCE, null); + } + + @Test + public void testAllowsAggregateOverTypedVariantExtraction() { + RexExpression.FunctionCall typedExtraction = + new RexExpression.FunctionCall(DataSchema.ColumnDataType.STRING, "variantGet", + List.of(new RexExpression.InputRef(0))); + AggregateNode node = aggregate("MINSTRING", false, VARIANT_SCHEMA, typedExtraction); + + node.visit(VariantTypeValidationVisitor.INSTANCE, null); + } + + @Test + public void testValidatesWindowAggregateInputs() { + QueryException exception = Assert.expectThrows(QueryException.class, + () -> window("SUM").visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Aggregate function SUM")); + + window("COUNT").visit(VariantTypeValidationVisitor.INSTANCE, null); + } + + @Test + public void testRejectsRawVariantWindowKeys() { + QueryException exception = Assert.expectThrows(QueryException.class, + () -> window("COUNT", VARIANT_SCHEMA, List.of(0), List.of()) + .visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Window PARTITION BY")); + + exception = Assert.expectThrows(QueryException.class, + () -> window("COUNT", VARIANT_SCHEMA, List.of(), List.of(new RelFieldCollation(0))) + .visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Window ORDER BY")); + } + + @Test + public void testAllowsWindowKeysOverTypedVariantExtraction() { + window("COUNT", TYPED_EXTRACTION_SCHEMA, List.of(0), List.of(new RelFieldCollation(0))) + .visit(VariantTypeValidationVisitor.INSTANCE, null); + } + + private static SetOpNode setOp(SetOpNode.SetOpType setOpType, boolean all) { + return new SetOpNode(0, VARIANT_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(), setOpType, all); + } + + private static JoinNode join(JoinNode.JoinStrategy strategy, DataSchema leftSchema, DataSchema rightSchema) { + ValueNode left = new ValueNode(0, leftSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + ValueNode right = new ValueNode(0, rightSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + DataSchema resultSchema = + new DataSchema(new String[]{"result"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + return new JoinNode(0, resultSchema, PlanNode.NodeHint.EMPTY, List.of(left, right), JoinRelType.INNER, List.of(0), + List.of(0), List.of(), strategy); + } + + private static JoinNode asofJoin(DataSchema leftSchema, DataSchema rightSchema) { + ValueNode left = new ValueNode(0, leftSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + ValueNode right = new ValueNode(0, rightSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + DataSchema resultSchema = new DataSchema(new String[]{"result"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + RexExpression matchCondition = + new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN, "GREATER_THAN", + List.of(new RexExpression.InputRef(1), new RexExpression.InputRef(leftSchema.size() + 1))); + return new JoinNode(0, resultSchema, PlanNode.NodeHint.EMPTY, List.of(left, right), JoinRelType.ASOF, List.of(0), + List.of(0), List.of(), JoinNode.JoinStrategy.ASOF, matchCondition); + } + + private static AggregateNode aggregate(String functionName, boolean distinct, DataSchema inputSchema) { + return aggregate(functionName, distinct, inputSchema, new RexExpression.InputRef(0)); + } + + private static AggregateNode aggregate(String functionName, boolean distinct, DataSchema inputSchema, + RexExpression operand) { + ValueNode input = new ValueNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + RexExpression.FunctionCall aggCall = functionCall(functionName, distinct, operand); + DataSchema resultSchema = + new DataSchema(new String[]{"result"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}); + return new AggregateNode(0, resultSchema, PlanNode.NodeHint.EMPTY, List.of(input), List.of(aggCall), List.of(-1), + List.of(), AggregateNode.AggType.DIRECT, false, List.of(), 0); + } + + private static WindowNode window(String functionName) { + return window(functionName, VARIANT_SCHEMA, List.of(), List.of()); + } + + private static WindowNode window(String functionName, DataSchema inputSchema, List keys, + List collations) { + ValueNode input = new ValueNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + return new WindowNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(input), keys, collations, + List.of(functionCall(functionName, false, new RexExpression.InputRef(0))), WindowNode.WindowFrameType.ROWS, + Integer.MIN_VALUE, Integer.MAX_VALUE, WindowNode.WindowExclusion.NO_OTHERS, List.of()); + } + + private static RexExpression.FunctionCall functionCall(String functionName, boolean distinct, + RexExpression operand) { + return new RexExpression.FunctionCall(DataSchema.ColumnDataType.LONG, functionName, List.of(operand), distinct, + false); + } +} diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/type/TypeFactoryTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/type/TypeFactoryTest.java index 52a44627bd7a..639ba7a7f437 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/type/TypeFactoryTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/type/TypeFactoryTest.java @@ -112,6 +112,10 @@ public Iterator relDataTypeConversion() { basicType = TYPE_FACTORY.createSqlType(SqlTypeName.VARBINARY); break; } + case VARIANT: { + basicType = TYPE_FACTORY.createSqlType(SqlTypeName.VARIANT); + break; + } case BIG_DECIMAL: { basicType = TYPE_FACTORY.createSqlType(SqlTypeName.DECIMAL); break; @@ -184,7 +188,8 @@ private boolean isColNullable(Schema schema) { @Test(dataProvider = "relDataTypeConversion") public void testArrayTypes(FieldSpec.DataType dataType, RelDataType arrayType, boolean columnNullMode) { - if (dataType == FieldSpec.DataType.BIG_DECIMAL || dataType == FieldSpec.DataType.JSON) { + if (dataType == FieldSpec.DataType.BIG_DECIMAL || dataType == FieldSpec.DataType.JSON + || dataType == FieldSpec.DataType.VARIANT) { return; } TypeFactory typeFactory = new TypeFactory(); @@ -205,7 +210,8 @@ public void testArrayTypes(FieldSpec.DataType dataType, RelDataType arrayType, b @Test(dataProvider = "relDataTypeConversion") public void testNullableArrayTypes(FieldSpec.DataType dataType, RelDataType arrayType, boolean columnNullMode) { - if (dataType == FieldSpec.DataType.BIG_DECIMAL || dataType == FieldSpec.DataType.JSON) { + if (dataType == FieldSpec.DataType.BIG_DECIMAL || dataType == FieldSpec.DataType.JSON + || dataType == FieldSpec.DataType.VARIANT) { return; } TypeFactory typeFactory = new TypeFactory(); @@ -229,7 +235,8 @@ public void testNullableArrayTypes(FieldSpec.DataType dataType, RelDataType arra @Test(dataProvider = "relDataTypeConversion") public void testNotNullableArrayTypes(FieldSpec.DataType dataType, RelDataType arrayType, boolean columnNullMode) { - if (dataType == FieldSpec.DataType.BIG_DECIMAL || dataType == FieldSpec.DataType.JSON) { + if (dataType == FieldSpec.DataType.BIG_DECIMAL || dataType == FieldSpec.DataType.JSON + || dataType == FieldSpec.DataType.VARIANT) { return; } TypeFactory typeFactory = new TypeFactory(); @@ -259,6 +266,7 @@ public void testRelDataTypeConversion() { .addSingleValueDimension("STRING_COL", FieldSpec.DataType.STRING) .addSingleValueDimension("UUID_COL", FieldSpec.DataType.UUID) .addSingleValueDimension("BYTES_COL", FieldSpec.DataType.BYTES) + .addSingleValueDimension("VARIANT_COL", FieldSpec.DataType.VARIANT) .addSingleValueDimension("JSON_COL", FieldSpec.DataType.JSON) .addMultiValueDimension("INT_ARRAY_COL", FieldSpec.DataType.INT) .addMultiValueDimension("LONG_ARRAY_COL", FieldSpec.DataType.LONG) @@ -304,6 +312,9 @@ public void testRelDataTypeConversion() { case "BYTES_COL": Assert.assertEquals(field.getType(), new BasicSqlType(TypeSystem.INSTANCE, SqlTypeName.VARBINARY)); break; + case "VARIANT_COL": + Assert.assertEquals(field.getType(), new BasicSqlType(TypeSystem.INSTANCE, SqlTypeName.VARIANT)); + break; case "INT_ARRAY_COL": Assert.assertEquals(field.getType(), new ArraySqlType(new BasicSqlType(TypeSystem.INSTANCE, SqlTypeName.INTEGER), false)); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java index 0e7501c635ad..96c2a7bf397c 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java @@ -48,6 +48,7 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.validation.VariantTypeValidationVisitor; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; import org.apache.pinot.query.runtime.operator.utils.SortUtils; @@ -91,6 +92,9 @@ public class AggregateOperator extends MultiStageOperator { public AggregateOperator(OpChainExecutionContext context, MultiStageOperator input, AggregateNode node) { super(context); + if (node.getInputs().size() == 1) { + VariantTypeValidationVisitor.validateAggregateInputs(node, node.getInputs().get(0).getDataSchema()); + } _resultSchema = node.getDataSchema(); _aggFunctions = getAggFunctions(node.getAggCalls()); int numFunctions = _aggFunctions.length; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java index 1a6ea42c2f72..55cf3a6ab575 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java @@ -20,6 +20,7 @@ import com.google.common.base.Preconditions; import java.util.ArrayList; +import java.util.Arrays; import java.util.BitSet; import java.util.HashMap; import java.util.List; @@ -31,6 +32,7 @@ import org.apache.pinot.query.planner.partitioning.KeySelector; import org.apache.pinot.query.planner.partitioning.KeySelectorFactory; import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.operator.join.DoubleLookupTable; import org.apache.pinot.query.runtime.operator.join.FloatLookupTable; @@ -66,11 +68,30 @@ public class HashJoinOperator extends BaseJoinOperator { @Nullable private List _nullKeyRightRows; + /** + * Creates a hash join using schemas available on the join node. + * + *

For SEMI and ANTI joins whose node does not carry its inputs, the result schema contains only left columns, so + * this legacy constructor cannot validate the right key's logical type. New callers that need right-side VARIANT + * validation must use the overload that accepts {@code rightSchema}. + */ public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, MultiStageOperator rightInput, JoinNode node) { + this(context, leftInput, leftSchema, rightInput, tryInferRightSchema(leftSchema, node), node, false); + } + + public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, + MultiStageOperator rightInput, DataSchema rightSchema, JoinNode node) { + this(context, leftInput, leftSchema, rightInput, rightSchema, node, true); + } + + private HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, + MultiStageOperator rightInput, @Nullable DataSchema rightSchema, JoinNode node, boolean rightSchemaRequired) { super(context, leftInput, leftSchema, rightInput, node); List leftKeys = node.getLeftKeys(); Preconditions.checkState(!leftKeys.isEmpty(), "Hash join operator requires join keys"); + Preconditions.checkArgument(!rightSchemaRequired || rightSchema != null, "Right input schema must not be null"); + validateVariantJoinKeys(leftKeys, node.getRightKeys(), leftSchema, rightSchema); _leftKeySelector = KeySelectorFactory.getKeySelector(leftKeys); _rightKeySelector = KeySelectorFactory.getKeySelector(node.getRightKeys()); _rightTable = createLookupTable(leftKeys, leftSchema); @@ -79,18 +100,57 @@ public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator left _nullKeyRightRows = needUnmatchedRightRows() ? new ArrayList<>() : null; } - /// Constructor that takes the schema for NonEquiEvaluator as an argument + /** + * Constructor that takes the schema for NonEquiEvaluator as an argument. + * + *

For SEMI and ANTI joins whose node does not carry its inputs, the result schema contains only left columns, so + * this legacy constructor cannot validate the right key's logical type. New callers that need right-side VARIANT + * validation must use the overload that accepts {@code rightSchema}. + */ public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, MultiStageOperator rightInput, JoinNode node, DataSchema nonEquiEvaluationSchema) { + this(context, leftInput, leftSchema, rightInput, tryInferRightSchema(leftSchema, node), node, + nonEquiEvaluationSchema, false); + } + + /// Constructor that takes the schema for NonEquiEvaluator as an argument + public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, + MultiStageOperator rightInput, DataSchema rightSchema, JoinNode node, DataSchema nonEquiEvaluationSchema) { + this(context, leftInput, leftSchema, rightInput, rightSchema, node, nonEquiEvaluationSchema, true); + } + + private HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, + MultiStageOperator rightInput, @Nullable DataSchema rightSchema, JoinNode node, + DataSchema nonEquiEvaluationSchema, boolean rightSchemaRequired) { super(context, leftInput, leftSchema, rightInput, node, nonEquiEvaluationSchema); List leftKeys = node.getLeftKeys(); Preconditions.checkState(!leftKeys.isEmpty(), "Hash join operator requires join keys"); + Preconditions.checkArgument(!rightSchemaRequired || rightSchema != null, "Right input schema must not be null"); + validateVariantJoinKeys(leftKeys, node.getRightKeys(), leftSchema, rightSchema); _leftKeySelector = KeySelectorFactory.getKeySelector(leftKeys); _rightKeySelector = KeySelectorFactory.getKeySelector(node.getRightKeys()); _rightTable = createLookupTable(leftKeys, leftSchema); _matchedRightRows = needUnmatchedRightRows() ? new HashMap<>() : null; } + @Nullable + private static DataSchema tryInferRightSchema(DataSchema leftSchema, JoinNode node) { + List inputs = node.getInputs(); + if (inputs.size() > 1) { + return inputs.get(1).getDataSchema(); + } + + DataSchema resultSchema = node.getDataSchema(); + int rightColumnOffset = leftSchema.size(); + int rightColumnCount = resultSchema.size() - rightColumnOffset; + int maxRightKey = node.getRightKeys().stream().mapToInt(Integer::intValue).max().orElse(-1); + if (rightColumnCount <= maxRightKey) { + return null; + } + return new DataSchema(Arrays.copyOfRange(resultSchema.getColumnNames(), rightColumnOffset, resultSchema.size()), + Arrays.copyOfRange(resultSchema.getColumnDataTypes(), rightColumnOffset, resultSchema.size())); + } + private static LookupTable createLookupTable(List joinKeys, DataSchema schema) { if (joinKeys.size() > 1) { return new ObjectLookupTable(); @@ -109,6 +169,22 @@ private static LookupTable createLookupTable(List joinKeys, DataSchema } } + private static void validateVariantJoinKeys(List leftKeys, List rightKeys, DataSchema leftSchema, + @Nullable DataSchema rightSchema) { + for (int leftKey : leftKeys) { + DataSchema.ColumnDataType dataType = leftSchema.getColumnDataType(leftKey); + Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), + "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); + } + if (rightSchema != null) { + for (int rightKey : rightKeys) { + DataSchema.ColumnDataType dataType = rightSchema.getColumnDataType(rightKey); + Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), + "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); + } + } + } + @Override public String toExplainString() { return EXPLAIN_NAME; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java index 02fd6baf9e39..9ea55929b13c 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java @@ -82,6 +82,13 @@ public MultistageGroupByExecutor(int[] groupKeyIds, AggregationFunction[] aggFun _aggType = aggType; _leafReturnFinalResult = leafReturnFinalResult; _resultSchema = resultSchema; + for (int i = 0; i < groupKeyIds.length; i++) { + ColumnDataType dataType = resultSchema.getColumnDataType(i); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw new IllegalArgumentException( + "Raw VARIANT values do not support GROUP BY; extract a typed path with variantGet first"); + } + } int maxInitialResultHolderCapacity = getResolvedMaxInitialResultHolderCapacity(opChainMetadata, nodeHint); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java index 82db0410b46f..be69c8b60a6b 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java @@ -20,6 +20,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Joiner; +import com.google.common.base.Preconditions; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -73,6 +74,11 @@ public SortOperator(OpChainExecutionContext context, MultiStageOperator input, S // - There is no collation // - Input is already sorted List collations = node.getCollations(); + for (RelFieldCollation collation : collations) { + Preconditions.checkArgument( + _dataSchema.getColumnDataType(collation.getFieldIndex()).supportsOrdering(), + "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first"); + } if (collations.isEmpty() || input instanceof SortedMailboxReceiveOperator) { _priorityQueue = null; _rows = new ArrayList<>(Math.min(defaultHolderCapacity, _numRowsToKeep)); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java index bede521ad63b..a4b22e8c80eb 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java @@ -56,6 +56,11 @@ public SortedMailboxReceiveOperator(OpChainExecutionContext context, MailboxRece Preconditions.checkState(!CollectionUtils.isEmpty(node.getCollations()), "Field collations must be set"); _dataSchema = node.getDataSchema(); _collations = node.getCollations(); + for (RelFieldCollation collation : _collations) { + Preconditions.checkArgument( + _dataSchema.getColumnDataType(collation.getFieldIndex()).supportsOrdering(), + "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first"); + } } @Override diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java index c42ff15ddc7e..adb80f0df67c 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java @@ -34,6 +34,7 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.WindowNode; +import org.apache.pinot.query.planner.validation.VariantTypeValidationVisitor; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; import org.apache.pinot.query.runtime.operator.utils.AggregationUtils; @@ -108,6 +109,7 @@ public class WindowAggregateOperator extends MultiStageOperator { public WindowAggregateOperator(OpChainExecutionContext context, MultiStageOperator input, DataSchema inputSchema, WindowNode node) { super(context); + VariantTypeValidationVisitor.validateWindowInputs(node, inputSchema); _input = input; _resultSchema = node.getDataSchema(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java index 5e024a20fc8f..329a2176795f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java @@ -18,7 +18,9 @@ */ package org.apache.pinot.query.runtime.operator.factory; +import com.google.common.base.Preconditions; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.EnrichedJoinNode; import org.apache.pinot.query.planner.plannode.JoinNode; import org.apache.pinot.query.planner.plannode.PlanNode; @@ -37,13 +39,15 @@ public MultiStageOperator createJoinOperator(OpChainExecutionContext context, Mu PlanNode leftPlanNode, MultiStageOperator rightOperator, PlanNode rightPlanNode, JoinNode joinNode) { JoinNode.JoinStrategy joinStrategy = joinNode.getJoinStrategy(); DataSchema leftSchema = leftPlanNode.getDataSchema(); + DataSchema rightSchema = rightPlanNode.getDataSchema(); + validateJoinKeys(joinNode, leftSchema, rightSchema); switch (joinStrategy) { case HASH: if (joinNode.getLeftKeys().isEmpty()) { // TODO: Consider adding non-equi as a separate join strategy. return new NonEquiJoinOperator(context, leftOperator, leftSchema, rightOperator, joinNode); } else { - return new HashJoinOperator(context, leftOperator, leftSchema, rightOperator, joinNode); + return new HashJoinOperator(context, leftOperator, leftSchema, rightOperator, rightSchema, joinNode); } case LOOKUP: return new LookupJoinOperator(context, leftOperator, leftSchema, rightOperator, joinNode); @@ -54,6 +58,28 @@ public MultiStageOperator createJoinOperator(OpChainExecutionContext context, Mu } } + private static void validateJoinKeys(JoinNode joinNode, DataSchema leftSchema, DataSchema rightSchema) { + for (int leftKey : joinNode.getLeftKeys()) { + DataSchema.ColumnDataType dataType = leftSchema.getColumnDataType(leftKey); + Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), + "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); + } + for (int rightKey : joinNode.getRightKeys()) { + DataSchema.ColumnDataType dataType = rightSchema.getColumnDataType(rightKey); + Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), + "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); + } + if (joinNode.getJoinStrategy() == JoinNode.JoinStrategy.ASOF) { + RexExpression.FunctionCall matchCondition = (RexExpression.FunctionCall) joinNode.getMatchCondition(); + int leftMatchKey = ((RexExpression.InputRef) matchCondition.getFunctionOperands().get(0)).getIndex(); + int rightMatchKey = + ((RexExpression.InputRef) matchCondition.getFunctionOperands().get(1)).getIndex() - leftSchema.size(); + Preconditions.checkArgument(leftSchema.getColumnDataType(leftMatchKey).supportsOrdering() + && rightSchema.getColumnDataType(rightMatchKey).supportsOrdering(), + "Raw VARIANT values do not support ASOF JOIN match keys; extract a typed path with variantGet first"); + } + } + /// Enriched joins have been removed. This method is retained only for backward compatibility of the /// [JoinOperatorFactory] interface and always throws. A current broker never produces an /// [EnrichedJoinNode], so this is only reachable if a plan from an older-version broker is executed. diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java index ea4f79032038..5d867bf12693 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java @@ -115,7 +115,10 @@ public static class In extends FilterOperand { public In(List children, DataSchema dataSchema, boolean isNotIn) { _childOperands = new ArrayList<>(children.size()); for (RexExpression child : children) { - _childOperands.add(TransformOperandFactory.getTransformOperand(child, dataSchema)); + TransformOperand operand = TransformOperandFactory.getTransformOperand(child, dataSchema); + Preconditions.checkArgument(operand.getResultType().supportsEquality(), + "Raw VARIANT values do not support IN; extract a typed path with variantGet first"); + _childOperands.add(operand); } _isNotIn = isNotIn; } @@ -193,6 +196,8 @@ public Predicate(List operands, DataSchema dataSchema, IntPredica ColumnDataType lhsType = _lhs.getResultType(); ColumnDataType rhsType = _rhs.getResultType(); + Preconditions.checkArgument(lhsType.supportsOrdering() && rhsType.supportsOrdering(), + "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); if (lhsType == rhsType) { _requireCasting = false; _commonCastType = null; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java new file mode 100644 index 000000000000..709ce6112fbf --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java @@ -0,0 +1,95 @@ +/** + * 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.query.runtime.operator.operands; + +import com.google.common.base.Preconditions; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.query.planner.logical.RexExpression; + + +/** + * Query-local constant operand for parsing a JSON literal into Variant. + * + *

The literal is parsed exactly once while the expression tree is constructed. The cached internal value is + * immutable by convention and can therefore be reused for every input row. Instances are thread-safe after + * construction. + */ +final class LiteralParseJsonOperand implements TransformOperand { + private static final String PARSE_JSON = "parsejson"; + private static final String PARSE_JSON_TO_VARIANT = "parsejsontovariant"; + private static final String TRY_PARSE_JSON = "tryparsejson"; + private static final String TRY_PARSE_JSON_TO_VARIANT = "tryparsejsontovariant"; + + private final ColumnDataType _resultType; + @Nullable + private final Object _value; + + LiteralParseJsonOperand(RexExpression.FunctionCall functionCall, String canonicalName) { + Preconditions.checkArgument(isSupported(canonicalName), "Unsupported JSON-to-Variant function: %s", + functionCall.getFunctionName()); + List operands = functionCall.getFunctionOperands(); + Preconditions.checkArgument(operands.size() == 1 && operands.get(0) instanceof RexExpression.Literal, + "%s expects one literal argument", functionCall.getFunctionName()); + + RexExpression.Literal literal = (RexExpression.Literal) operands.get(0); + Object literalValue = literal.getValue(); + ColumnDataType literalType = literal.getDataType(); + boolean validLiteral = literalValue == null + || ((literalType == ColumnDataType.STRING || literalType == ColumnDataType.JSON) + && literalValue instanceof String); + Preconditions.checkArgument(validLiteral, + "%s argument must be a STRING or JSON literal", functionCall.getFunctionName()); + + _resultType = functionCall.getDataType(); + String json = (String) literalValue; + byte[] variant = isTolerant(canonicalName) + ? VariantUtils.tryParseJsonToVariant(json) : VariantUtils.parseJsonToVariant(json); + _value = variant != null ? _resultType.toInternal(variant) : null; + } + + static boolean isSupported(String canonicalName) { + switch (canonicalName) { + case PARSE_JSON: + case PARSE_JSON_TO_VARIANT: + case TRY_PARSE_JSON: + case TRY_PARSE_JSON_TO_VARIANT: + return true; + default: + return false; + } + } + + private static boolean isTolerant(String canonicalName) { + return canonicalName.equals(TRY_PARSE_JSON) || canonicalName.equals(TRY_PARSE_JSON_TO_VARIANT); + } + + @Override + public ColumnDataType getResultType() { + return _resultType; + } + + @Nullable + @Override + public Object apply(List row) { + return _value; + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java index 1c39351bce68..b7e0470f4795 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java @@ -20,6 +20,7 @@ import com.google.common.base.Preconditions; import java.util.List; +import org.apache.pinot.common.function.FunctionRegistry; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; @@ -43,6 +44,14 @@ public static TransformOperand getTransformOperand(RexExpression rexExpression, private static TransformOperand getTransformOperand(RexExpression.FunctionCall functionCall, DataSchema dataSchema) { List operands = functionCall.getFunctionOperands(); int numOperands = operands.size(); + String canonicalName = FunctionRegistry.canonicalize(functionCall.getFunctionName()); + if (VariantOperand.isSupported(canonicalName)) { + return new VariantOperand(functionCall, dataSchema, canonicalName); + } + if (LiteralParseJsonOperand.isSupported(canonicalName) && numOperands == 1 + && operands.get(0) instanceof RexExpression.Literal) { + return new LiteralParseJsonOperand(functionCall, canonicalName); + } switch (functionCall.getFunctionName()) { case "AND": Preconditions.checkState(numOperands >= 2, "AND takes >=2 arguments, got: %s", numOperands); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java new file mode 100644 index 000000000000..b76bfef52a81 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java @@ -0,0 +1,221 @@ +/** + * 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.query.runtime.operator.operands; + +import com.google.common.base.Preconditions; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.common.utils.VariantUtils.ResultType; +import org.apache.pinot.common.utils.VariantUtils.ReusableResult; +import org.apache.pinot.common.utils.VariantUtils.VariantPath; +import org.apache.pinot.query.planner.logical.RexExpression; + + +/** + * Query-local multi-stage operand for Variant scalar operations. + * + *

Literal paths and target types are compiled once at construction. Values enter and leave the operand in + * {@link DataSchema}'s internal representation, which {@link VariantUtils.ReusableResult} materializes directly after + * extraction. The reusable extraction result makes instances not thread-safe. + */ +final class VariantOperand implements TransformOperand { + private static final VariantPath ROOT_PATH = VariantUtils.compilePath("$"); + private static final String VARIANT_GET = "variantget"; + private static final String TRY_VARIANT_GET = "tryvariantget"; + private static final String VARIANT_EXISTS = "variantexists"; + private static final String IS_VARIANT_NULL = "isvariantnull"; + private static final String VARIANT_TYPE_OF = "varianttypeof"; + private static final String VARIANT_TO_JSON = "varianttojson"; + + private final ColumnDataType _resultType; + private final TransformOperand _variantOperand; + private final Operation _operation; + private final VariantPath _path; + @Nullable + private final ResultType _targetType; + private final ReusableResult _reusableResult = new ReusableResult(); + + VariantOperand(RexExpression.FunctionCall functionCall, DataSchema dataSchema, String canonicalName) { + _resultType = functionCall.getDataType(); + List operands = functionCall.getFunctionOperands(); + _operation = operation(canonicalName); + validateArgumentCount(_operation, operands.size(), functionCall.getFunctionName()); + _variantOperand = TransformOperandFactory.getTransformOperand(operands.get(0), dataSchema); + ColumnDataType inputType = _variantOperand.getResultType(); + Preconditions.checkArgument( + inputType == ColumnDataType.VARIANT || inputType == ColumnDataType.BYTES || inputType == ColumnDataType.UNKNOWN, + "%s first argument must be a VARIANT", functionCall.getFunctionName()); + + switch (_operation) { + case GET: + case TRY_GET: + _path = compilePathLiteral(operands.get(1), functionCall.getFunctionName()); + _targetType = operands.size() == 2 ? ResultType.VARIANT + : parseTypeLiteral(operands.get(2), functionCall.getFunctionName()); + break; + case EXISTS: + _path = compilePathLiteral(operands.get(1), functionCall.getFunctionName()); + _targetType = null; + break; + case IS_NULL: + case TYPE_OF: + _path = + operands.size() == 2 ? compilePathLiteral(operands.get(1), functionCall.getFunctionName()) : ROOT_PATH; + _targetType = null; + break; + case TO_JSON: + _path = ROOT_PATH; + _targetType = null; + break; + default: + throw new IllegalStateException("Unhandled Variant operation: " + _operation); + } + } + + static boolean isSupported(String canonicalName) { + switch (canonicalName) { + case VARIANT_GET: + case TRY_VARIANT_GET: + case VARIANT_EXISTS: + case IS_VARIANT_NULL: + case VARIANT_TYPE_OF: + case VARIANT_TO_JSON: + return true; + default: + return false; + } + } + + @Override + public ColumnDataType getResultType() { + return _resultType; + } + + @Nullable + @Override + public Object apply(List row) { + Object internalVariant = _variantOperand.apply(row); + byte[] variant = internalVariant != null + ? (byte[]) _variantOperand.getResultType().toExternal(internalVariant) : null; + if (_operation == Operation.GET) { + return extractInternal(variant, false); + } + if (_operation == Operation.TRY_GET) { + return extractInternal(variant, true); + } + Object externalResult; + switch (_operation) { + case EXISTS: + externalResult = VariantUtils.variantExists(variant, _path, _reusableResult); + break; + case IS_NULL: + externalResult = VariantUtils.isVariantNull(variant, _path, _reusableResult); + break; + case TYPE_OF: + externalResult = VariantUtils.variantTypeOf(variant, _path, _reusableResult); + break; + case TO_JSON: + externalResult = VariantUtils.variantToJson(variant); + break; + default: + throw new IllegalStateException("Unhandled Variant operation: " + _operation); + } + return externalResult != null ? _resultType.toInternal(externalResult) : null; + } + + @Nullable + private Object extractInternal(@Nullable byte[] variant, boolean tolerant) { + ResultType targetType = Preconditions.checkNotNull(_targetType, "Variant target type must be planned"); + boolean present = tolerant + ? VariantUtils.tryExtractInto(variant, _path, targetType, _reusableResult) + : VariantUtils.extractInto(variant, _path, targetType, _reusableResult); + return present ? _reusableResult.getInternalValue(targetType) : null; + } + + private static Operation operation(String canonicalName) { + switch (canonicalName) { + case VARIANT_GET: + return Operation.GET; + case TRY_VARIANT_GET: + return Operation.TRY_GET; + case VARIANT_EXISTS: + return Operation.EXISTS; + case IS_VARIANT_NULL: + return Operation.IS_NULL; + case VARIANT_TYPE_OF: + return Operation.TYPE_OF; + case VARIANT_TO_JSON: + return Operation.TO_JSON; + default: + throw new IllegalArgumentException("Unsupported Variant function: " + canonicalName); + } + } + + private static void validateArgumentCount(Operation operation, int numOperands, String functionName) { + boolean valid; + switch (operation) { + case GET: + case TRY_GET: + valid = numOperands == 2 || numOperands == 3; + break; + case EXISTS: + valid = numOperands == 2; + break; + case IS_NULL: + case TYPE_OF: + valid = numOperands == 1 || numOperands == 2; + break; + case TO_JSON: + valid = numOperands == 1; + break; + default: + throw new IllegalStateException("Unhandled Variant operation: " + operation); + } + Preconditions.checkArgument(valid, "Invalid number of arguments for %s: %s", functionName, numOperands); + } + + private static VariantPath compilePathLiteral(RexExpression operand, String functionName) { + return VariantUtils.compilePath(stringLiteral(operand, functionName, "path")); + } + + private static ResultType parseTypeLiteral(RexExpression operand, String functionName) { + return VariantUtils.parseResultType(stringLiteral(operand, functionName, "target type")); + } + + private static String stringLiteral(RexExpression operand, String functionName, String role) { + Preconditions.checkArgument(operand instanceof RexExpression.Literal, + "%s %s must be a string literal", functionName, role); + RexExpression.Literal literal = (RexExpression.Literal) operand; + Preconditions.checkArgument(literal.getDataType() == ColumnDataType.STRING && literal.getValue() instanceof String, + "%s %s must be a string literal", functionName, role); + return (String) literal.getValue(); + } + + private enum Operation { + GET, + TRY_GET, + EXISTS, + IS_NULL, + TYPE_OF, + TO_JSON + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java index 76f94b126880..d462cf080f2c 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java @@ -45,6 +45,7 @@ public BinarySetOperator(OpChainExecutionContext opChainExecutionContext, DataSchema dataSchema) { super(opChainExecutionContext, inputOperators, dataSchema); Preconditions.checkArgument(inputOperators.size() == 2, "Binary set operator should have 2 inputs"); + validateEqualitySupported(dataSchema, "INTERSECT/EXCEPT"); _leftChildOperator = inputOperators.get(0); _rightChildOperator = inputOperators.get(1); _rightRowSet = HashMultiset.create(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java index de49e3c53443..84f64d1fca35 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.query.runtime.operator.set; +import com.google.common.base.Preconditions; import java.util.List; import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.utils.DataSchema; @@ -40,6 +41,13 @@ public SetOperator(OpChainExecutionContext opChainExecutionContext, List inputOperators, DataSchema dataSchema) { super(opChainExecutionContext, inputOperators, dataSchema); + validateEqualitySupported(dataSchema, "UNION DISTINCT"); } @Override diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java index 60ab3e503564..1b298c967f99 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java @@ -32,6 +32,7 @@ import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.AggregateNode.AggType; import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.ValueNode; import org.apache.pinot.query.routing.VirtualServerAddress; import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; @@ -39,9 +40,11 @@ import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; import org.apache.pinot.spi.utils.CommonConstants.Server; import org.mockito.Mock; +import org.testng.Assert; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -96,6 +99,41 @@ public void shouldHandleUpstreamErrorBlocks() { assertTrue(block.isError(), "Input errors should propagate immediately"); } + @Test + public void testRejectsRawVariantAggregationAtRuntime() { + DataSchema inputSchema = + new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + ValueNode inputPlanNode = new ValueNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + RexExpression.FunctionCall sum = + new RexExpression.FunctionCall(ColumnDataType.DOUBLE, SqlKind.SUM.name(), + List.of(new RexExpression.InputRef(0))); + DataSchema resultSchema = new DataSchema(new String[]{"sum"}, new ColumnDataType[]{DOUBLE}); + AggregateNode aggregateNode = + new AggregateNode(0, resultSchema, PlanNode.NodeHint.EMPTY, List.of(inputPlanNode), List.of(sum), List.of(-1), + List.of(), AggType.DIRECT, false, List.of(), 0); + + QueryException exception = Assert.expectThrows(QueryException.class, + () -> new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, aggregateNode)); + assertTrue(exception.getMessage().contains("Aggregate function SUM does not support raw VARIANT")); + assertTrue(exception.getMessage().contains("variantGet")); + } + + @Test + public void testAllowsRawVariantCountAtRuntime() { + DataSchema inputSchema = + new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + ValueNode inputPlanNode = new ValueNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + RexExpression.FunctionCall count = + new RexExpression.FunctionCall(ColumnDataType.LONG, SqlKind.COUNT.name(), + List.of(new RexExpression.InputRef(0))); + DataSchema resultSchema = new DataSchema(new String[]{"count"}, new ColumnDataType[]{ColumnDataType.LONG}); + AggregateNode aggregateNode = + new AggregateNode(0, resultSchema, PlanNode.NodeHint.EMPTY, List.of(inputPlanNode), List.of(count), List.of(-1), + List.of(), AggType.DIRECT, false, List.of(), 0); + + Assert.assertNotNull(new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, aggregateNode)); + } + @Test public void shouldHandleEndOfStreamBlockWithNoOtherInputs() { // Given: diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java index dcd9c04ce05a..42402eefb762 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java @@ -42,7 +42,9 @@ import static org.mockito.MockitoAnnotations.openMocks; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; import static org.testng.internal.junit.ArrayAsserts.assertArrayEquals; @@ -284,6 +286,26 @@ public void shouldHandleAntiJoin() { assertTrue(operator.nextBlock().isSuccess()); } + @Test + public void shouldRejectRawVariantRightKeyForSemiJoin() { + assertRawVariantRightKeyRejected(JoinRelType.SEMI); + } + + @Test + public void shouldRejectRawVariantRightKeyForAntiJoin() { + assertRawVariantRightKeyRejected(JoinRelType.ANTI); + } + + @Test + public void shouldPreserveLegacyConstructorForSemiJoin() { + assertLegacyConstructorSupports(JoinRelType.SEMI); + } + + @Test + public void shouldPreserveLegacyConstructorForAntiJoin() { + assertLegacyConstructorSupports(JoinRelType.ANTI); + } + @Test public void shouldPropagateRightTableError() { _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA) @@ -735,8 +757,9 @@ public void shouldHandleCompositeKeySemiJoinWithNulls() { new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING, ColumnDataType.DOUBLE}); // Composite key join on columns 1 and 2 (string_col and double_col) - HashJoinOperator operator = getOperator(resultSchema, JoinRelType.SEMI, - List.of(1, 2), List.of(1, 2), List.of()); + HashJoinOperator operator = + getOperator(compositeSchema, compositeSchema, resultSchema, JoinRelType.SEMI, List.of(1, 2), List.of(1, 2), + List.of(), PlanNode.NodeHint.EMPTY); List resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); @@ -771,8 +794,9 @@ public void shouldHandleCompositeKeyAntiJoinWithNulls() { new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.STRING, ColumnDataType.DOUBLE}); // Composite key join on columns 1 and 2 (string_col and double_col) - HashJoinOperator operator = getOperator(resultSchema, JoinRelType.ANTI, - List.of(1, 2), List.of(1, 2), List.of()); + HashJoinOperator operator = + getOperator(compositeSchema, compositeSchema, resultSchema, JoinRelType.ANTI, List.of(1, 2), List.of(1, 2), + List.of(), PlanNode.NodeHint.EMPTY); List resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); @@ -872,7 +896,14 @@ public void shouldRecordJoinedOutputSizeWhenRightTableFitsButJoinedOutputExceeds private HashJoinOperator getOperator(DataSchema leftSchema, DataSchema resultSchema, JoinRelType joinType, List leftKeys, List rightKeys, List nonEquiConditions, PlanNode.NodeHint nodeHint) { - return new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, leftSchema, _rightInput, + return getOperator(leftSchema, leftSchema, resultSchema, joinType, leftKeys, rightKeys, nonEquiConditions, + nodeHint); + } + + private HashJoinOperator getOperator(DataSchema leftSchema, DataSchema rightSchema, DataSchema resultSchema, + JoinRelType joinType, List leftKeys, List rightKeys, + List nonEquiConditions, PlanNode.NodeHint nodeHint) { + return new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, leftSchema, _rightInput, rightSchema, new JoinNode(-1, resultSchema, nodeHint, List.of(), joinType, leftKeys, rightKeys, nonEquiConditions, JoinNode.JoinStrategy.HASH)); } @@ -880,9 +911,8 @@ private HashJoinOperator getOperator(DataSchema leftSchema, DataSchema resultSch private HashJoinOperator getOperator(DataSchema resultSchema, JoinRelType joinType, List leftKeys, List rightKeys, List nonEquiConditions, PlanNode.NodeHint nodeHint) { - return new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, DEFAULT_CHILD_SCHEMA, _rightInput, - new JoinNode(-1, resultSchema, nodeHint, List.of(), joinType, leftKeys, rightKeys, nonEquiConditions, - JoinNode.JoinStrategy.HASH)); + return getOperator(DEFAULT_CHILD_SCHEMA, DEFAULT_CHILD_SCHEMA, resultSchema, joinType, leftKeys, rightKeys, + nonEquiConditions, nodeHint); } private HashJoinOperator getOperator(DataSchema resultSchema, JoinRelType joinType, @@ -890,4 +920,27 @@ private HashJoinOperator getOperator(DataSchema resultSchema, JoinRelType joinTy return getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, joinType, leftKeys, rightKeys, nonEquiConditions, PlanNode.NodeHint.EMPTY); } + + private void assertRawVariantRightKeyRejected(JoinRelType joinType) { + DataSchema rightSchema = new DataSchema(new String[]{"variant_col"}, + new ColumnDataType[]{ColumnDataType.VARIANT}); + _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA).buildWithEos(); + _rightInput = new BlockListMultiStageOperator.Builder(rightSchema).buildWithEos(); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, + () -> getOperator(DEFAULT_CHILD_SCHEMA, rightSchema, DEFAULT_CHILD_SCHEMA, joinType, List.of(0), List.of(0), + List.of(), PlanNode.NodeHint.EMPTY)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support JOIN keys")); + } + + private void assertLegacyConstructorSupports(JoinRelType joinType) { + _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA).buildWithEos(); + _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA).buildWithEos(); + JoinNode joinNode = + new JoinNode(-1, DEFAULT_CHILD_SCHEMA, PlanNode.NodeHint.EMPTY, List.of(), joinType, List.of(0), List.of(0), + List.of(), JoinNode.JoinStrategy.HASH); + + assertNotNull(new HashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, DEFAULT_CHILD_SCHEMA, + _rightInput, joinNode)); + } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java index 0f7cff7c02e2..6ce55dc48e78 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java @@ -38,11 +38,13 @@ import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.LONG; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.STRING; +import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.VARIANT; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.mockito.MockitoAnnotations.openMocks; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; public class SortOperatorTest { @@ -499,6 +501,16 @@ public void shouldPreservePrecision() { assertTrue(operator.nextBlock().isSuccess(), "expected EOS block to propagate"); } + @Test + public void shouldRejectRawVariantCollation() { + DataSchema schema = new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{VARIANT}); + List collations = List.of(new RelFieldCollation(0)); + + IllegalArgumentException exception = + expectThrows(IllegalArgumentException.class, () -> getOperator(schema, collations)); + assertTrue(exception.getMessage().contains("ORDER BY does not support raw VARIANT")); + } + private SortOperator getOperator(DataSchema schema, List collations, int fetch, int offset) { return new SortOperator(OperatorTestUtil.getTracingContext(), _input, new SortNode(-1, schema, PlanNode.NodeHint.EMPTY, List.of(), collations, fetch, offset)); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java index 59aeccd6d5d5..83fd87cc6a00 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java @@ -50,11 +50,13 @@ import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT; import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.STRING; +import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.VARIANT; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; public class SortedMailboxReceiveOperatorTest { @@ -113,6 +115,17 @@ public void shouldThrowOnEmptyCollationKey() { getOperator(_stageMetadata1, RelDistribution.Type.SINGLETON, DATA_SCHEMA, List.of(), Long.MAX_VALUE); } + @Test + public void shouldRejectRawVariantCollation() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + DataSchema variantSchema = new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{VARIANT}); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, + () -> getOperator(_stageMetadata1, RelDistribution.Type.SINGLETON, variantSchema, FIELD_COLLATIONS, + Long.MAX_VALUE)); + assertTrue(exception.getMessage().contains("ORDER BY does not support raw VARIANT")); + } + @Test public void shouldTimeout() { when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 02fa104d6578..431accc81659 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -36,6 +36,7 @@ import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; import org.mockito.Mock; import org.testng.Assert; import org.testng.annotations.AfterMethod; @@ -115,6 +116,40 @@ public void testShouldHandleEndOfStreamBlockWithNoOtherInputs() { assertTrue(block.isSuccess(), "EOS blocks should propagate"); } + @Test + public void testRejectsRawVariantWindowKeys() { + DataSchema inputSchema = new DataSchema(new String[]{"payload"}, new ColumnDataType[]{VARIANT}); + DataSchema resultSchema = + new DataSchema(new String[]{"payload", "count"}, new ColumnDataType[]{VARIANT, LONG}); + MultiStageOperator input = new BlockListMultiStageOperator.Builder(inputSchema).buildWithEos(); + List aggCalls = List.of(getCount(new RexExpression.InputRef(0))); + + QueryException exception = Assert.expectThrows(QueryException.class, + () -> getOperator(inputSchema, resultSchema, List.of(0), List.of(), aggCalls, ROWS, + Integer.MIN_VALUE, Integer.MAX_VALUE, input)); + assertTrue(exception.getMessage().contains("Window PARTITION BY")); + + exception = Assert.expectThrows(QueryException.class, + () -> getOperator(inputSchema, resultSchema, List.of(), List.of(new RelFieldCollation(0)), aggCalls, ROWS, + Integer.MIN_VALUE, Integer.MAX_VALUE, input)); + assertTrue(exception.getMessage().contains("Window ORDER BY")); + } + + @Test + public void testAllowsWindowKeysOverTypedVariantExtraction() { + DataSchema inputSchema = new DataSchema(new String[]{"typedPayload"}, new ColumnDataType[]{STRING}); + DataSchema resultSchema = + new DataSchema(new String[]{"typedPayload", "count"}, new ColumnDataType[]{STRING, LONG}); + MultiStageOperator input = new BlockListMultiStageOperator.Builder(inputSchema).buildWithEos(); + List aggCalls = List.of(getCount(new RexExpression.InputRef(0))); + + WindowAggregateOperator operator = + getOperator(inputSchema, resultSchema, List.of(0), List.of(new RelFieldCollation(0)), aggCalls, ROWS, + Integer.MIN_VALUE, Integer.MAX_VALUE, input); + + assertTrue(operator.nextBlock().isSuccess()); + } + @Test public void testShouldWindowAggregateOverSingleInputBlock() { // Given: @@ -3524,6 +3559,10 @@ private static RexExpression.FunctionCall getSum(RexExpression arg) { return new RexExpression.FunctionCall(ColumnDataType.INT, SqlKind.SUM.name(), List.of(arg)); } + private static RexExpression.FunctionCall getCount(RexExpression arg) { + return new RexExpression.FunctionCall(ColumnDataType.LONG, SqlKind.COUNT.name(), List.of(arg)); + } + private static RexExpression.FunctionCall getMin(RexExpression arg) { return new RexExpression.FunctionCall(ColumnDataType.INT, SqlKind.MIN.name(), List.of(arg)); } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactoryTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactoryTest.java index 8276d8d7f0de..44dcad0f93d5 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactoryTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactoryTest.java @@ -18,9 +18,24 @@ */ package org.apache.pinot.query.runtime.operator.factory; +import java.util.List; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.runtime.operator.HashJoinOperator; +import org.apache.pinot.query.runtime.operator.MultiStageOperator; +import org.apache.pinot.query.runtime.operator.OperatorTestUtil; +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.assertThrows; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; public class DefaultJoinOperatorFactoryTest { @@ -34,4 +49,94 @@ public void createEnrichedJoinOperatorThrows() { assertThrows(UnsupportedOperationException.class, () -> factory.createEnrichedJoinOperator(null, null, null, null, null, null)); } + + @Test(dataProvider = "semiAndAntiJoinTypes") + public void createHashJoinUsesRightInputSchema(JoinRelType joinType) { + DataSchema leftSchema = new DataSchema(new String[]{"left_key"}, + new ColumnDataType[]{ColumnDataType.INT}); + DataSchema rightSchema = new DataSchema(new String[]{"right_key"}, + new ColumnDataType[]{ColumnDataType.INT}); + MultiStageOperator operator = createHashJoin(leftSchema, rightSchema, joinType); + + assertTrue(operator instanceof HashJoinOperator); + } + + @Test(dataProvider = "semiAndAntiJoinTypes") + public void createHashJoinRejectsRawVariantRightKey(JoinRelType joinType) { + DataSchema leftSchema = new DataSchema(new String[]{"left_key"}, + new ColumnDataType[]{ColumnDataType.INT}); + DataSchema rightSchema = new DataSchema(new String[]{"right_key"}, + new ColumnDataType[]{ColumnDataType.VARIANT}); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, + () -> createHashJoin(leftSchema, rightSchema, joinType)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support JOIN keys")); + } + + @Test + public void createLookupJoinRejectsRawVariantKeys() { + DataSchema typedSchema = new DataSchema(new String[]{"key"}, new ColumnDataType[]{ColumnDataType.INT}); + DataSchema variantSchema = + new DataSchema(new String[]{"key"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, + () -> createJoin(variantSchema, typedSchema, JoinRelType.INNER, JoinNode.JoinStrategy.LOOKUP)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support JOIN keys")); + + exception = expectThrows(IllegalArgumentException.class, + () -> createJoin(typedSchema, variantSchema, JoinRelType.INNER, JoinNode.JoinStrategy.LOOKUP)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support JOIN keys")); + } + + @Test + public void createAsofJoinRejectsRawVariantMatchKeys() { + DataSchema leftSchema = + new DataSchema(new String[]{"key", "match"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.VARIANT}); + DataSchema rightSchema = + new DataSchema(new String[]{"key", "match"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.LONG}); + RexExpression matchCondition = + new RexExpression.FunctionCall(ColumnDataType.BOOLEAN, "GREATER_THAN", + List.of(new RexExpression.InputRef(1), new RexExpression.InputRef(3))); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, + () -> createJoin(leftSchema, rightSchema, JoinRelType.ASOF, JoinNode.JoinStrategy.ASOF, matchCondition)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); + + DataSchema typedLeftSchema = + new DataSchema(new String[]{"key", "match"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.LONG}); + DataSchema variantRightSchema = + new DataSchema(new String[]{"key", "match"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.VARIANT}); + exception = expectThrows(IllegalArgumentException.class, + () -> createJoin(typedLeftSchema, variantRightSchema, JoinRelType.ASOF, JoinNode.JoinStrategy.ASOF, + matchCondition)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); + } + + @DataProvider + private static Object[][] semiAndAntiJoinTypes() { + return new Object[][]{{JoinRelType.SEMI}, {JoinRelType.ANTI}}; + } + + private static MultiStageOperator createHashJoin(DataSchema leftSchema, DataSchema rightSchema, + JoinRelType joinType) { + return createJoin(leftSchema, rightSchema, joinType, JoinNode.JoinStrategy.HASH); + } + + private static MultiStageOperator createJoin(DataSchema leftSchema, DataSchema rightSchema, JoinRelType joinType, + JoinNode.JoinStrategy joinStrategy) { + return createJoin(leftSchema, rightSchema, joinType, joinStrategy, null); + } + + private static MultiStageOperator createJoin(DataSchema leftSchema, DataSchema rightSchema, JoinRelType joinType, + JoinNode.JoinStrategy joinStrategy, RexExpression matchCondition) { + PlanNode leftPlanNode = mock(PlanNode.class); + when(leftPlanNode.getDataSchema()).thenReturn(leftSchema); + PlanNode rightPlanNode = mock(PlanNode.class); + when(rightPlanNode.getDataSchema()).thenReturn(rightSchema); + JoinNode joinNode = + new JoinNode(-1, leftSchema, PlanNode.NodeHint.EMPTY, List.of(), joinType, List.of(0), List.of(0), List.of(), + joinStrategy, matchCondition); + return new DefaultJoinOperatorFactory().createJoinOperator(OperatorTestUtil.getTracingContext(), + mock(MultiStageOperator.class), leftPlanNode, mock(MultiStageOperator.class), rightPlanNode, joinNode); + } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java new file mode 100644 index 000000000000..a9561ca709ba --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java @@ -0,0 +1,50 @@ +/** + * 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.query.runtime.operator.operands; + +import java.util.List; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.testng.Assert; +import org.testng.annotations.Test; + + +public class FilterOperandTest { + private static final DataSchema VARIANT_SCHEMA = + new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + private static final List VARIANT_OPERANDS = + List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(0)); + + @Test + public void testRawVariantInIsRejected() { + IllegalArgumentException exception = + Assert.expectThrows(IllegalArgumentException.class, + () -> new FilterOperand.In(VARIANT_OPERANDS, VARIANT_SCHEMA, false)); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support IN")); + } + + @Test + public void testRawVariantNotInIsRejected() { + IllegalArgumentException exception = + Assert.expectThrows(IllegalArgumentException.class, + () -> new FilterOperand.In(VARIANT_OPERANDS, VARIANT_SCHEMA, true)); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support IN")); + } +} diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/VariantOperandTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/VariantOperandTest.java new file mode 100644 index 000000000000..03a32d1c93a2 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/VariantOperandTest.java @@ -0,0 +1,271 @@ +/** + * 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.query.runtime.operator.operands; + +import java.math.BigDecimal; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.function.Consumer; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.VariantEnvelope; +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; + + +public class VariantOperandTest { + private static final DataSchema VARIANT_SCHEMA = + new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + + @Test + public void testFactoryUsesSpecializedOperandAndPreservesNullSemantics() { + List row = row(VariantUtils.parseJsonToVariant( + "{\"name\":\"alice\",\"number\":7,\"presentNull\":null,\"nested\":{\"enabled\":true}}")); + List missingRow = row(VariantUtils.parseJsonToVariant("{}")); + List variantNullRow = row(VariantUtils.parseJsonToVariant("null")); + List sqlNullRow = Collections.singletonList(null); + + TransformOperand strictGet = operand(ColumnDataType.STRING, "variantGet", + new RexExpression.InputRef(0), stringLiteral("$.name"), stringLiteral("STRING")); + assertTrue(strictGet instanceof VariantOperand); + assertEquals(strictGet.apply(row), "alice"); + assertNull(strictGet.apply(missingRow)); + assertEquals(strictGet.apply(row), "alice", "The reusable cursor must reset after a missing path"); + assertNull(strictGet.apply(sqlNullRow)); + + TransformOperand defaultGet = operand(ColumnDataType.VARIANT, "VARIANT_GET", + new RexExpression.InputRef(0), stringLiteral("$.nested")); + Object nested = defaultGet.apply(row); + assertTrue(nested instanceof ByteArray, "VARIANT results must retain the internal BYTES wrapper"); + assertEquals(VariantUtils.variantToJson((byte[]) ColumnDataType.VARIANT.toExternal(nested)), "{\"enabled\":true}"); + + TransformOperand tolerantGet = operand(ColumnDataType.DOUBLE, "tryVariantGet", + new RexExpression.InputRef(0), stringLiteral("$.name"), stringLiteral("DOUBLE")); + assertNull(tolerantGet.apply(row)); + TransformOperand mismatchedStrictGet = operand(ColumnDataType.DOUBLE, "variantGet", + new RexExpression.InputRef(0), stringLiteral("$.name"), stringLiteral("DOUBLE")); + assertThrows(IllegalArgumentException.class, () -> mismatchedStrictGet.apply(row)); + + TransformOperand exists = operand(ColumnDataType.BOOLEAN, "variantExists", + new RexExpression.InputRef(0), stringLiteral("$.presentNull")); + assertTrue(exists instanceof VariantOperand); + assertEquals(exists.apply(row), 1, "A present Variant null counts as present"); + assertEquals(exists.apply(missingRow), 0, "The reusable cursor must not retain a prior present result"); + assertNull(exists.apply(sqlNullRow), "variantExists preserves SQL null"); + assertEquals(operand(ColumnDataType.BOOLEAN, "variantExists", + new RexExpression.InputRef(0), stringLiteral("$.missing")).apply(row), 0); + + TransformOperand isNull = operand(ColumnDataType.BOOLEAN, "isVariantNull", + new RexExpression.InputRef(0), stringLiteral("$.presentNull")); + assertTrue(isNull instanceof VariantOperand); + assertEquals(isNull.apply(row), 1); + assertEquals(isNull.apply(missingRow), 0); + assertEquals(isNull.apply(sqlNullRow), 0, "isVariantNull(SQL NULL) is non-null false"); + TransformOperand rootIsNull = + operand(ColumnDataType.BOOLEAN, "isVariantNull", new RexExpression.InputRef(0)); + assertEquals(rootIsNull.apply(row), 0); + assertEquals(rootIsNull.apply(variantNullRow), 1); + assertEquals(rootIsNull.apply(sqlNullRow), 0); + + TransformOperand typeOf = operand(ColumnDataType.STRING, "variantTypeOf", + new RexExpression.InputRef(0), stringLiteral("$.presentNull")); + assertTrue(typeOf instanceof VariantOperand); + assertEquals(typeOf.apply(row), "NULL"); + assertNull(typeOf.apply(missingRow)); + assertNull(typeOf.apply(sqlNullRow)); + TransformOperand rootTypeOf = + operand(ColumnDataType.STRING, "variantTypeOf", new RexExpression.InputRef(0)); + assertEquals(rootTypeOf.apply(row), "OBJECT"); + assertEquals(rootTypeOf.apply(variantNullRow), "NULL"); + assertNull(rootTypeOf.apply(sqlNullRow)); + + TransformOperand toJson = + operand(ColumnDataType.STRING, "variantToJson", new RexExpression.InputRef(0)); + assertTrue(toJson instanceof VariantOperand); + assertEquals(toJson.apply(row), + "{\"name\":\"alice\",\"nested\":{\"enabled\":true},\"number\":7,\"presentNull\":null}"); + assertEquals(toJson.apply(variantNullRow), "null"); + assertNull(toJson.apply(sqlNullRow)); + } + + @Test + public void testEveryVariantGetTargetUsesInternalDataSchemaRepresentation() { + assertEquals(typedGet(variant(builder -> builder.appendBoolean(true)), "BOOLEAN", ColumnDataType.BOOLEAN), 1); + assertEquals(typedGet(variant(builder -> builder.appendInt(42)), "INT", ColumnDataType.INT), 42); + assertEquals(typedGet(variant(builder -> builder.appendLong(4_294_967_296L)), "LONG", ColumnDataType.LONG), + 4_294_967_296L); + assertEquals(typedGet(variant(builder -> builder.appendFloat(1.25F)), "FLOAT", ColumnDataType.FLOAT), 1.25F); + assertEquals(typedGet(variant(builder -> builder.appendDouble(12.5)), "DOUBLE", ColumnDataType.DOUBLE), 12.5); + + BigDecimal decimal = new BigDecimal("1234567890.12345"); + assertEquals(typedGet(variant(builder -> builder.appendDecimal(decimal)), "BIG_DECIMAL", + ColumnDataType.BIG_DECIMAL), decimal); + assertEquals(typedGet(variant(builder -> builder.appendString("click")), "STRING", ColumnDataType.STRING), "click"); + + byte[] binary = new byte[]{0, 1, (byte) 0xFF}; + Object binaryResult = + typedGet(variant(builder -> builder.appendBinary(ByteBuffer.wrap(binary))), "BYTES", ColumnDataType.BYTES); + assertEquals(binaryResult, new ByteArray(binary)); + + UUID uuid = UUID.fromString("12345678-1234-5678-9abc-def012345678"); + byte[] uuidVariant = variant(builder -> builder.appendUUID(uuid)); + Object uuidResult = typedGet(uuidVariant, "UUID", ColumnDataType.UUID); + assertTrue(uuidResult instanceof ByteArray, "UUID uses the internal BYTES wrapper"); + assertTrue(UuidUtils.equals(((ByteArray) uuidResult).getBytes(), UuidUtils.toBytes(uuid)), + "UUID extraction must copy the encoded 16-byte value directly"); + assertEquals(ColumnDataType.UUID.toExternal(uuidResult), uuid); + Object tolerantUuidResult = typedTryGet(uuidVariant, "UUID", ColumnDataType.UUID); + assertTrue(UuidUtils.equals(((ByteArray) tolerantUuidResult).getBytes(), UuidUtils.toBytes(uuid))); + + long timestampMicros = 1_700_000_000_123_000L; + byte[] timestampVariant = variant(builder -> builder.appendTimestampTz(timestampMicros)); + assertEquals(typedGet(timestampVariant, "TIMESTAMP", ColumnDataType.TIMESTAMP), 1_700_000_000_123L); + assertEquals(typedTryGet(timestampVariant, "TIMESTAMP", ColumnDataType.TIMESTAMP), 1_700_000_000_123L); + + Object variantResult = + typedGet(variant(builder -> builder.appendString("nested")), "VARIANT", ColumnDataType.VARIANT); + assertTrue(variantResult instanceof ByteArray, "VARIANT uses the internal BYTES wrapper"); + assertEquals(VariantUtils.variantToJson((byte[]) ColumnDataType.VARIANT.toExternal(variantResult)), "\"nested\""); + + assertEquals(typedGet(variant(builder -> builder.appendString("json")), "JSON", ColumnDataType.STRING), "\"json\""); + } + + @Test + public void testLiteralPathAndTargetTypeAreCompiledAtConstruction() { + assertThrows("Invalid paths must fail before any row is evaluated", IllegalArgumentException.class, + () -> operand(ColumnDataType.STRING, "variantGet", + new RexExpression.InputRef(0), stringLiteral("payload.name"), stringLiteral("STRING"))); + assertThrows("Invalid target types must fail before any row is evaluated", IllegalArgumentException.class, + () -> operand(ColumnDataType.STRING, "try_variant_get", + new RexExpression.InputRef(0), stringLiteral("$"), stringLiteral("UNSUPPORTED"))); + + DataSchema nonLiteralPathSchema = new DataSchema(new String[]{"payload", "path"}, + new ColumnDataType[]{ColumnDataType.VARIANT, ColumnDataType.STRING}); + RexExpression.FunctionCall nonLiteralPath = new RexExpression.FunctionCall(ColumnDataType.BOOLEAN, + "variant_exists", List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1))); + assertThrows(IllegalArgumentException.class, + () -> TransformOperandFactory.getTransformOperand(nonLiteralPath, nonLiteralPathSchema)); + + TransformOperand snakeCase = operand(ColumnDataType.BOOLEAN, "VARIANT_EXISTS", + new RexExpression.InputRef(0), stringLiteral("$")); + assertTrue(snakeCase instanceof VariantOperand, "Factory routing must use canonical function names"); + assertFalse(VariantOperand.isSupported("parsejson"), "Only path-sensitive Variant operations are specialized"); + } + + @Test + public void testLiteralParseJsonIsEvaluatedOnceAndPreservesNullSemantics() { + for (String functionName : List.of("parseJson", "parse_json", "parseJsonToVariant", + "parse_json_to_variant")) { + TransformOperand strict = operand(ColumnDataType.VARIANT, functionName, + stringLiteral("{\"name\":\"alice\"}")); + assertTrue(strict instanceof LiteralParseJsonOperand); + Object cachedValue = strict.apply(List.of()); + assertSame(strict.apply(List.of()), cachedValue, "The parsed Variant must be reused for every row"); + assertEquals(VariantUtils.variantToJson((byte[]) ColumnDataType.VARIANT.toExternal(cachedValue)), + "{\"name\":\"alice\"}"); + } + + assertThrows(IllegalArgumentException.class, + () -> operand(ColumnDataType.VARIANT, "parseJson", stringLiteral("{not-json"))); + + for (String functionName : List.of("tryParseJson", "try_parse_json", "tryParseJsonToVariant", + "try_parse_json_to_variant")) { + TransformOperand validTolerant = operand(ColumnDataType.VARIANT, functionName, stringLiteral("[1,2,3]")); + Object cachedValue = validTolerant.apply(List.of()); + assertSame(validTolerant.apply(List.of()), cachedValue, + "The tolerantly parsed Variant must be reused for every row"); + assertEquals(VariantUtils.variantToJson((byte[]) ColumnDataType.VARIANT.toExternal(cachedValue)), "[1,2,3]"); + + TransformOperand tolerant = operand(ColumnDataType.VARIANT, functionName, stringLiteral("{not-json")); + assertTrue(tolerant instanceof LiteralParseJsonOperand); + assertNull(tolerant.apply(List.of())); + assertNull(tolerant.apply(List.of())); + } + + TransformOperand sqlNull = operand(ColumnDataType.VARIANT, "parseJson", + new RexExpression.Literal(ColumnDataType.UNKNOWN, null)); + assertTrue(sqlNull instanceof LiteralParseJsonOperand); + assertNull(sqlNull.apply(List.of())); + + TransformOperand variantNull = + operand(ColumnDataType.VARIANT, "parseJson", stringLiteral("null")); + Object encodedVariantNull = variantNull.apply(List.of()); + assertTrue(VariantUtils.isVariantNull((byte[]) ColumnDataType.VARIANT.toExternal(encodedVariantNull)), + "JSON null must remain distinct from SQL null"); + + RexExpression.FunctionCall literalParse = new RexExpression.FunctionCall(ColumnDataType.VARIANT, "parseJson", + List.of(stringLiteral("{\"nested\":{\"name\":\"alice\"}}"))); + TransformOperand nestedGet = operand(ColumnDataType.STRING, "variantGet", + literalParse, stringLiteral("$.nested.name"), stringLiteral("STRING")); + assertEquals(nestedGet.apply(List.of()), "alice"); + + DataSchema stringSchema = + new DataSchema(new String[]{"json"}, new ColumnDataType[]{ColumnDataType.STRING}); + RexExpression.FunctionCall dynamicCall = new RexExpression.FunctionCall(ColumnDataType.VARIANT, "parseJson", + List.of(new RexExpression.InputRef(0))); + assertFalse(TransformOperandFactory.getTransformOperand(dynamicCall, stringSchema) + instanceof LiteralParseJsonOperand, "Non-literal parsing must keep the row-dependent function operand"); + } + + private static Object typedGet(byte[] variant, String targetType, ColumnDataType resultType) { + return operand(resultType, "variantGet", + new RexExpression.InputRef(0), stringLiteral("$"), stringLiteral(targetType)).apply(row(variant)); + } + + private static Object typedTryGet(byte[] variant, String targetType, ColumnDataType resultType) { + return operand(resultType, "tryVariantGet", + new RexExpression.InputRef(0), stringLiteral("$"), stringLiteral(targetType)).apply(row(variant)); + } + + private static TransformOperand operand(ColumnDataType resultType, String functionName, + RexExpression... operands) { + RexExpression.FunctionCall functionCall = + new RexExpression.FunctionCall(resultType, functionName, List.of(operands)); + return TransformOperandFactory.getTransformOperand(functionCall, VARIANT_SCHEMA); + } + + private static RexExpression.Literal stringLiteral(String value) { + return new RexExpression.Literal(ColumnDataType.STRING, value); + } + + private static List row(byte[] variant) { + return List.of(ColumnDataType.VARIANT.toInternal(variant)); + } + + private static byte[] variant(Consumer writer) { + VariantBuilder builder = new VariantBuilder(); + writer.accept(builder); + Variant variant = builder.build(); + return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); + } +} diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/IntersectOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/IntersectOperatorTest.java index 7f2b66b8b02d..90f1bd2fa930 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/IntersectOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/IntersectOperatorTest.java @@ -29,6 +29,8 @@ import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; + public class IntersectOperatorTest { @@ -147,4 +149,15 @@ public void testErrorBlockLeftChild() { } Assert.assertTrue(result.isError()); } + + @Test + public void testRejectsVariantSetOperations() { + DataSchema schema = new DataSchema(new String[]{"payload"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.VARIANT}); + List inputs = List.of(mock(MultiStageOperator.class), mock(MultiStageOperator.class)); + + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> new IntersectOperator(OperatorTestUtil.getTracingContext(), inputs, schema)); + Assert.assertTrue(exception.getMessage().contains("INTERSECT/EXCEPT does not support raw VARIANT")); + } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java index 1f64b4f6f6da..4f295b7754c7 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java @@ -30,6 +30,8 @@ import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; + public class UnionOperatorTest { @@ -112,4 +114,17 @@ public void testErrorBlockLeftChild() { } Assert.assertTrue(result.isError()); } + + @Test + public void testVariantSetOperationValidation() { + DataSchema schema = new DataSchema(new String[]{"payload"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.VARIANT}); + List inputs = List.of(mock(MultiStageOperator.class), mock(MultiStageOperator.class)); + + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> new UnionOperator(OperatorTestUtil.getTracingContext(), inputs, schema)); + Assert.assertTrue(exception.getMessage().contains("UNION DISTINCT does not support raw VARIANT")); + + new UnionAllOperator(OperatorTestUtil.getTracingContext(), inputs, schema); + } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java index 21791dee355c..b8f82bd55154 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/BaseSegmentCreator.java @@ -351,6 +351,9 @@ private boolean createDictionaryForColumn(ColumnStatistics stats, SegmentGenerat if (spec instanceof ComplexFieldSpec) { return false; } + if (spec.getDataType() == FieldSpec.DataType.VARIANT) { + return false; + } String column = spec.getName(); FieldIndexConfigs fieldIndexConfigs = config.getIndexConfigsByColName().get(column); @@ -663,7 +666,7 @@ public static void addColumnMetadataInfo(PropertiesConfiguration properties, Str } // Min/max value - if (fieldSpec.getFieldType() != FieldType.COMPLEX) { + if (fieldSpec.getFieldType() != FieldType.COMPLEX && fieldSpec.getDataType().supportsMinMax()) { // Regular (non-complex) field if (totalDocs > 0) { Object min = columnStatistics.getMinValue(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizer.java index 0f978d4ec7ab..25be6e669a6a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizer.java @@ -28,21 +28,21 @@ /// (`buildColumnar`) path. /// /// The row-major build runs each record through a `TransformPipeline` whose -/// `NullValueTransformer` substitutes [FieldSpec#getDefaultNullValue()] for `null` -/// and whose `DataTypeTransformer` coerces every value to the column's stored type (e.g. +/// `DataTypeTransformer` coerces every non-null value to the column's stored type (e.g. /// `Boolean` → `Integer` for a `BOOLEAN` column stored as `INT`, -/// `Timestamp` → `Long` for `TIMESTAMP`). The column-major driver deliberately runs -/// with no transform pipeline, so a non-segment source (e.g. Arrow) delivers values in the source's +/// `Timestamp` → `Long` for `TIMESTAMP`) and whose `NullValueTransformer` substitutes +/// [FieldSpec#getDefaultNullValue()] for values that remain `null`. The column-major driver deliberately runs with +/// no transform pipeline, so a non-segment source (e.g. Arrow) delivers values in the source's /// logical type with raw `null`s — which the typed collectors / index creators do not accept. /// /// This helper applies the equivalent of those two transformers to one value, in the same order: /// -/// 1. `NullValueTransformer`: a `null` value becomes the column default — the scalar -/// default for single-value columns, or a one-element `Object[]` of that scalar for -/// multi-value columns (matching `NullValueTransformerUtils.getDefaultNullValue`). -/// 2. `DataTypeTransformer`: [DataTypeTransformerUtils#transformValue] standardizes the +/// 1. `DataTypeTransformer`: [DataTypeTransformerUtils#transformValue] standardizes the /// value (collapsing single-element collections, dropping `null` elements from multi-value /// arrays) and converts it to `destDataType` via [PinotDataType]. +/// 2. `NullValueTransformer`: a `null` value becomes the column default — the scalar +/// default for single-value columns, or a one-element `Object[]` of that scalar for +/// multi-value columns (matching `NullValueTransformerUtils.getDefaultNullValue`). /// /// A whole-value `null` (or a value that standardizes to `null`, e.g. an empty array) /// therefore resolves to the column default rather than reaching the collector as `null`. The @@ -63,13 +63,11 @@ private ColumnarValueNormalizer() { /// @param value the raw value read from the `ColumnReader` (may be `null`) /// @return the normalized, never-`null` value to feed to the stats collector / index creator public static Object normalize(String column, FieldSpec fieldSpec, PinotDataType destDataType, Object value) { - if (value == null) { - value = defaultNullValue(fieldSpec); - } value = DataTypeTransformerUtils.transformValue(column, value, destDataType); if (value == null) { - // The value standardized to null (e.g. an empty multi-value array). Substitute the default so a - // null never reaches the typed collectors / index creators. + // The source was null or standardized to null (e.g. an empty multi-value array). FieldSpec stores its default + // in the canonical ingestion representation, so return it directly. In particular, VARIANT uses empty bytes as + // the reserved SQL-null sentinel, which must never be decoded as a PVAR envelope. value = defaultNullValue(fieldSpec); } return value; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java index 58c4951391e6..9c188ff61533 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/creator/impl/stats/AbstractColumnStatisticsCollector.java @@ -66,7 +66,9 @@ public AbstractColumnStatisticsCollector(FieldSpec fieldSpec, @Nullable FieldCon @Nullable PartitionFunction partitionFunction) { _fieldSpec = fieldSpec; _storedType = fieldSpec.getDataType().getStoredType(); - _sorted = fieldSpec.isSingleValueField(); + // PVAR envelope byte ordering is not semantic VARIANT ordering. Never advertise a VARIANT column as sorted even + // when its physical bytes happen to be monotonic. + _sorted = fieldSpec.isSingleValueField() && fieldSpec.getDataType().supportsOrdering(); _fieldConfig = fieldConfig; _partitionFunction = partitionFunction; _partitions = partitionFunction != null ? new HashSet<>() : null; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java index ae4b99f84790..a85a6aef2fac 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java @@ -105,21 +105,21 @@ private List getColumnsToAddMinMaxValue() { switch (_columnMinMaxValueGeneratorMode) { case ALL: for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!fieldSpec.isVirtualColumn()) { + if (supportsMinMax(fieldSpec)) { columnsToAddMinMaxValue.add(fieldSpec.getName()); } } break; case NON_METRIC: for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!fieldSpec.isVirtualColumn() && fieldSpec.getFieldType() != FieldSpec.FieldType.METRIC) { + if (supportsMinMax(fieldSpec) && fieldSpec.getFieldType() != FieldSpec.FieldType.METRIC) { columnsToAddMinMaxValue.add(fieldSpec.getName()); } } break; case TIME: for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { - if (!fieldSpec.isVirtualColumn() && (fieldSpec.getFieldType() == FieldSpec.FieldType.TIME + if (supportsMinMax(fieldSpec) && (fieldSpec.getFieldType() == FieldSpec.FieldType.TIME || fieldSpec.getFieldType() == FieldSpec.FieldType.DATE_TIME)) { columnsToAddMinMaxValue.add(fieldSpec.getName()); } @@ -132,6 +132,10 @@ private List getColumnsToAddMinMaxValue() { return columnsToAddMinMaxValue; } + private static boolean supportsMinMax(FieldSpec fieldSpec) { + return !fieldSpec.isVirtualColumn() && fieldSpec.getDataType().supportsMinMax(); + } + private boolean needAddColumnMinMaxValueForColumn(String columnName) { return needAddColumnMinMaxValueForColumn(_segmentMetadata.getColumnMetadataFor(columnName)); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SanitizationTransformerUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SanitizationTransformerUtils.java index 26e9bb0bbe33..42fc0ac2e258 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SanitizationTransformerUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SanitizationTransformerUtils.java @@ -50,10 +50,25 @@ private SanitizationTransformerUtils() { @Nullable public static SanitizedColumnInfo getSanitizedColumnInfo(FieldSpec fieldSpec) { FieldSpec.DataType dataType = fieldSpec.getDataType(); + MaxLengthExceedStrategy strategy = fieldSpec.getEffectiveMaxLengthExceedStrategy(); + + if (dataType == FieldSpec.DataType.VARIANT) { + if (strategy == MaxLengthExceedStrategy.TRIM_LENGTH + || strategy == MaxLengthExceedStrategy.SUBSTITUTE_DEFAULT_VALUE) { + throw new IllegalStateException( + "VARIANT envelope cannot use max length exceed strategy: " + strategy); + } + // ERROR can enforce a configured envelope-size bound without modifying the envelope. NO_ACTION deliberately + // bypasses generic byte sanitization because Variant payload validation belongs to VariantEnvelope/the producer. + if (strategy == MaxLengthExceedStrategy.ERROR) { + return new SanitizedColumnInfo(fieldSpec.getName(), fieldSpec.getEffectiveMaxLength(), strategy, + fieldSpec.getDefaultNullValue()); + } + return null; + } if (dataType == FieldSpec.DataType.STRING || dataType == FieldSpec.DataType.JSON || dataType == FieldSpec.DataType.BYTES) { - MaxLengthExceedStrategy strategy = fieldSpec.getEffectiveMaxLengthExceedStrategy(); // For STRING, always apply (to handle null characters even with NO_ACTION) // For JSON/BYTES, only apply if strategy is not NO_ACTION if (dataType == FieldSpec.DataType.STRING || strategy != MaxLengthExceedStrategy.NO_ACTION) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SchemaUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SchemaUtils.java index dafd3dad4c80..ff6e086dd82f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SchemaUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SchemaUtils.java @@ -19,6 +19,7 @@ package org.apache.pinot.segment.local.utils; import com.google.common.base.Preconditions; +import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -196,6 +197,9 @@ public static void validate(Schema schema, boolean isIgnoreCase) { .equals(FieldSpec.DataType.DOUBLE)) { validateDefaultIsNotNaN(fieldSpec); } + if (fieldSpec.getDataType() == FieldSpec.DataType.VARIANT) { + validateVariantFieldSpec(fieldSpec); + } if (!fieldSpec.isSingleValueField()) { validateMultiValueCompatibility(fieldSpec); } @@ -222,12 +226,27 @@ private static void validateDefaultIsNotNaN(FieldSpec fieldSpec) { fieldSpec.getName()); } + private static void validateVariantFieldSpec(FieldSpec fieldSpec) { + Preconditions.checkState(fieldSpec.getFieldType() == FieldSpec.FieldType.DIMENSION, + "VARIANT column must be a dimension: %s", fieldSpec.getName()); + Preconditions.checkState( + Arrays.equals((byte[]) fieldSpec.getDefaultNullValue(), FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_BYTES), + "VARIANT column cannot define a custom default null value: %s", fieldSpec.getName()); + FieldSpec.MaxLengthExceedStrategy strategy = fieldSpec.getEffectiveMaxLengthExceedStrategy(); + Preconditions.checkState(strategy != FieldSpec.MaxLengthExceedStrategy.TRIM_LENGTH + && strategy != FieldSpec.MaxLengthExceedStrategy.SUBSTITUTE_DEFAULT_VALUE, + "VARIANT column cannot use truncating or substituting max length strategy %s: %s", strategy, + fieldSpec.getName()); + } + /// Validations for MV type columns. Kept here (rather than in [Schema#validate()]) so that schema construction /// via `SchemaBuilder.build()` stays a pure DTO operation and only the controller-side ingest validation rejects /// MV JSON columns. private static void validateMultiValueCompatibility(FieldSpec fieldSpec) { Preconditions.checkState(!fieldSpec.getDataType().equals(FieldSpec.DataType.JSON), "JSON columns cannot be of multi-value type"); + Preconditions.checkState(fieldSpec.getDataType() != FieldSpec.DataType.VARIANT, + "VARIANT columns cannot be of multi-value type"); } /// Validates that the schema is compatible with the given table config diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java index 8ef3c48727e0..b10be0ac202d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java @@ -54,6 +54,7 @@ import org.apache.pinot.segment.spi.index.DictionaryIndexConfig; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.ForwardIndexConfig; import org.apache.pinot.segment.spi.index.IndexService; import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; @@ -65,6 +66,7 @@ import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.config.table.FieldConfig.EncodingType; import org.apache.pinot.spi.config.table.HashFunction; +import org.apache.pinot.spi.config.table.IndexConfig; import org.apache.pinot.spi.config.table.IndexingConfig; import org.apache.pinot.spi.config.table.MultiColumnTextIndexConfig; import org.apache.pinot.spi.config.table.QuotaConfig; @@ -769,8 +771,12 @@ private static Set validateMetricsAggregation(TableConfig tableConfig, S "Metrics aggregation cannot be enabled when the schema contains COMPLEX columns"); for (String dimension : schema.getDimensionNames()) { - Preconditions.checkState(schema.getFieldSpecFor(dimension).isSingleValueField(), + FieldSpec fieldSpec = schema.getFieldSpecFor(dimension); + Preconditions.checkState(fieldSpec.isSingleValueField(), "Metrics aggregation cannot be enabled with multi-value dimension column: %s", dimension); + Preconditions.checkState( + fieldSpec.getDataType().supportsEquality() && fieldSpec.getDataType().supportsHashing(), + "Metrics aggregation cannot use VARIANT dimension column as an aggregation key: %s", dimension); } Map dictConfigByCol = StandardIndexes.dictionary().getConfig(tableConfig, schema); @@ -1078,6 +1084,8 @@ static void validateUpsertAndDedupConfig(TableConfig tableConfig, Schema schema, if (comparisonColumns != null) { for (String column : comparisonColumns) { Preconditions.checkState(schema.hasColumn(column), "The comparison column does not exist on schema"); + Preconditions.checkState(schema.getFieldSpecFor(column).getDataType().supportsOrdering(), + "VARIANT column cannot be used as an upsert comparison column: %s", column); } } @@ -1452,6 +1460,10 @@ static void validatePartialUpsertStrategies(TableConfig tableConfig, Schema sche FieldSpec fieldSpec = schema.getFieldSpecFor(column); Preconditions.checkState(fieldSpec != null, "Merger cannot be applied to non-existing column: %s", column); + if (fieldSpec.getDataType() == DataType.VARIANT) { + Preconditions.checkState(columnStrategy == UpsertConfig.Strategy.OVERWRITE, + "VARIANT column supports only OVERWRITE partial-upsert strategy: %s", column); + } if (columnStrategy == UpsertConfig.Strategy.INCREMENT) { Preconditions.checkState(fieldSpec.getDataType().getStoredType().isNumeric(), @@ -1771,6 +1783,12 @@ private static void validateTierConfigList(@Nullable List tierConfig /// - Proper dependency between index types (e.g. inverted index columns must have dictionary) private static void validateIndexingConfigAndFieldConfigList(TableConfig tableConfig, Schema schema) { IndexingConfig indexingConfig = tableConfig.getIndexingConfig(); + if (schema.getAllFieldSpecs().stream().anyMatch(fieldSpec -> fieldSpec.getDataType() == DataType.VARIANT)) { + Preconditions.checkState( + schema.isEnableColumnBasedNullHandling() || indexingConfig.isNullHandlingEnabled(), + "Null handling must be enabled for tables containing VARIANT columns via " + + "schema.enableColumnBasedNullHandling or tableIndexConfig.nullHandlingEnabled"); + } List fieldConfigs = tableConfig.getFieldConfigList(); if (CollectionUtils.isNotEmpty(fieldConfigs)) { Set seenColumns = new HashSet<>(); @@ -1797,11 +1815,23 @@ private static void validateIndexingConfigAndFieldConfigList(TableConfig tableCo FieldSpec fieldSpec = schema.getFieldSpecFor(column); Preconditions.checkState(fieldSpec != null, "Failed to find column: %s in schema", column); FieldIndexConfigs indexConfigs = entry.getValue(); + if (fieldSpec.getDataType() == DataType.VARIANT) { + validateVariantIndexConfigs(indexConfigs, fieldSpec, allIndexes); + } for (IndexType indexType : allIndexes) { indexType.validate(indexConfigs, fieldSpec, tableConfig); } } + if (CollectionUtils.isNotEmpty(schema.getPrimaryKeyColumns())) { + for (String primaryKeyColumn : schema.getPrimaryKeyColumns()) { + FieldSpec fieldSpec = schema.getFieldSpecFor(primaryKeyColumn); + Preconditions.checkState(fieldSpec == null + || (fieldSpec.getDataType().supportsEquality() && fieldSpec.getDataType().supportsHashing()), + "VARIANT column cannot be used as a primary key: %s", primaryKeyColumn); + } + } + // Null value vector backfill on the time column needs the whole table config (for the time column name), so it // cannot be validated by NullValueIndexType.validate which only sees one field at a time. validateNullValueVectorBackfillForTimeColumn(tableConfig, schema, indexConfigsMap); @@ -1831,6 +1861,12 @@ private static void validateIndexingConfigAndFieldConfigList(TableConfig tableCo validateStarTreeIndexConfigs(starTreeIndexConfigs, indexConfigsMap, schema, TimestampIndexUtils.extractColumnsWithGranularity(tableConfig)); } + if (indexingConfig.isEnableDefaultStarTree()) { + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + Preconditions.checkState(fieldSpec.getDataType() != DataType.VARIANT, + "Default star-tree index cannot include VARIANT column: %s", fieldSpec.getName()); + } + } // TIMESTAMP index is not managed by FieldIndexConfigs, and we need to validate it separately. if (CollectionUtils.isNotEmpty(fieldConfigs)) { @@ -1855,6 +1891,8 @@ private static void validateIndexingConfigAndFieldConfigList(TableConfig tableCo FieldSpec fieldSpec = schema.getFieldSpecFor(column); Preconditions.checkState(fieldSpec != null, "Failed to find sorted column: %s in schema", column); Preconditions.checkState(fieldSpec.isSingleValueField(), "Cannot sort on multi-value column: %s", column); + Preconditions.checkState(fieldSpec.getDataType().supportsOrdering(), + "Cannot sort on VARIANT column: %s", column); } } @@ -1868,11 +1906,45 @@ private static void validateIndexingConfigAndFieldConfigList(TableConfig tableCo Preconditions.checkState(fieldSpec != null, "Failed to find partition column: %s in schema", column); Preconditions.checkState(fieldSpec.isSingleValueField(), "Cannot partition on multi-value column: %s", column); + Preconditions.checkState(fieldSpec.getDataType().supportsHashing(), + "Cannot partition on VARIANT column: %s", column); } } } } + private static void validateVariantIndexConfigs(FieldIndexConfigs indexConfigs, FieldSpec fieldSpec, + List> allIndexes) { + String column = fieldSpec.getName(); + Preconditions.checkState(fieldSpec.isSingleValueField(), + "VARIANT column must be single-value: %s", column); + + ForwardIndexConfig forwardIndexConfig = indexConfigs.getConfig(StandardIndexes.forward()); + Preconditions.checkState(forwardIndexConfig.isEnabled(), + "VARIANT column must have an enabled forward index: %s", column); + Preconditions.checkState(forwardIndexConfig.getEncodingType() == EncodingType.RAW, + "VARIANT column must use RAW forward index encoding: %s", column); + Preconditions.checkState(indexConfigs.getConfig(StandardIndexes.dictionary()).isDisabled(), + "VARIANT column cannot use a dictionary: %s", column); + + for (IndexType indexType : allIndexes) { + String indexId = indexType.getId(); + if (StandardIndexes.FORWARD_ID.equals(indexId) || StandardIndexes.NULL_VALUE_VECTOR_ID.equals(indexId) + || StandardIndexes.DICTIONARY_ID.equals(indexId) || StandardIndexes.OPEN_STRUCT_ID.equals(indexId)) { + // OPEN_STRUCT has an enabled default for every column but is applicable only to OPEN_STRUCT field specs. + continue; + } + Preconditions.checkState(!getIndexConfig(indexConfigs, indexType).isEnabled(), + "VARIANT column supports only a RAW forward index; index %s is not supported on column: %s", indexId, + column); + } + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private static IndexConfig getIndexConfig(FieldIndexConfigs indexConfigs, IndexType indexType) { + return indexConfigs.getConfig((IndexType) indexType); + } + /// Rejects a null value vector backfill opt-in on the time column when its default null value is outside the valid /// time range. /// @@ -2068,6 +2140,8 @@ private static void validateStarTreeIndexConfigs(List starT "Failed to find column: %s specified in star-tree index config in schema", column); Preconditions.checkState(fieldSpec.getDataType() != DataType.MAP, "Star-tree index cannot be created on MAP column: %s", column); + Preconditions.checkState(fieldSpec.getDataType() != DataType.VARIANT, + "Star-tree index cannot be created on VARIANT column: %s", column); } for (String column : dimensionColumns) { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ExpressionTransformerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ExpressionTransformerTest.java index 3fa66599b05d..636f62a67c53 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ExpressionTransformerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ExpressionTransformerTest.java @@ -25,6 +25,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; @@ -591,4 +592,39 @@ public void testJsonToArrayIngestionTransform() { Assert.assertTrue(transformedValue.getClass().isArray()); Assert.assertEquals(Arrays.asList((Object[]) transformedValue), Arrays.asList("a", "b", "c")); } + + @Test + public void testVariantScalarIngestionTransforms() { + Schema schema = new Schema.SchemaBuilder() + .addSingleValueDimension("payload", FieldSpec.DataType.VARIANT) + .addSingleValueDimension("eventType", FieldSpec.DataType.STRING) + .addSingleValueDimension("amount", FieldSpec.DataType.DOUBLE) + .addSingleValueDimension("hasCustomer", FieldSpec.DataType.BOOLEAN) + .addSingleValueDimension("nullValue", FieldSpec.DataType.BOOLEAN) + .addSingleValueDimension("payloadType", FieldSpec.DataType.STRING) + .build(); + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setTransformConfigs(List.of( + new TransformConfig("eventType", "variantGet(payload, '$.eventType', 'STRING')"), + new TransformConfig("amount", "tryVariantGet(payload, '$.amount', 'DOUBLE')"), + new TransformConfig("hasCustomer", "variantExists(payload, '$.customerId')"), + new TransformConfig("nullValue", "isVariantNull(payload, '$.deletedAt')"), + new TransformConfig("payloadType", "variantTypeOf(payload)"))); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE) + .setTableName("testVariantScalarIngestionTransforms") + .setIngestionConfig(ingestionConfig) + .build(); + ExpressionTransformer transformer = new ExpressionTransformer(tableConfig, schema); + + GenericRow row = new GenericRow(); + row.putValue("payload", VariantUtils.parseJsonToVariant( + "{\"eventType\":\"checkout\",\"amount\":42.5,\"customerId\":\"u-1\",\"deletedAt\":null}")); + transformer.transform(row); + + Assert.assertEquals(row.getValue("eventType"), "checkout"); + Assert.assertEquals(row.getValue("amount"), 42.5D); + Assert.assertEquals(row.getValue("hasCustomer"), true); + Assert.assertEquals(row.getValue("nullValue"), true); + Assert.assertEquals(row.getValue("payloadType"), "OBJECT"); + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizerTest.java new file mode 100644 index 000000000000..c9a7a9b698d0 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/ColumnarValueNormalizerTest.java @@ -0,0 +1,43 @@ +/** + * 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.segment.local.segment.creator.impl; + +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.PinotDataType; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; + + +public class ColumnarValueNormalizerTest { + private static final String COLUMN = "payload"; + + @Test + public void testNullVariantReturnsDefaultSentinelWithoutDecoding() { + FieldSpec fieldSpec = new DimensionFieldSpec(COLUMN, DataType.VARIANT, true); + + Object result = ColumnarValueNormalizer.normalize(COLUMN, fieldSpec, PinotDataType.VARIANT, null); + + assertSame(result, fieldSpec.getDefaultNullValue()); + assertEquals(((byte[]) result).length, 0); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/ColumnMetadataTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/ColumnMetadataTest.java index aa41f74b411c..3e0d478277bc 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/ColumnMetadataTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/ColumnMetadataTest.java @@ -50,6 +50,7 @@ import org.apache.pinot.spi.data.ComplexFieldSpec; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.FileFormat; @@ -374,6 +375,17 @@ public void testComplexFieldSpec() { assertEquals(intMapColumnMetadata.getFieldSpec(), intMapFieldSpec); } + @Test + public void testVariantFieldSpecRoundTrip() { + DimensionFieldSpec variantFieldSpec = new DimensionFieldSpec("payload", DataType.VARIANT, true); + PropertiesConfiguration config = new PropertiesConfiguration(); + BaseSegmentCreator.addFieldSpec(config, "payload", variantFieldSpec); + + FieldSpec restoredFieldSpec = ColumnMetadataImpl.extractFieldSpec("payload", config); + assertEquals(restoredFieldSpec, variantFieldSpec); + assertEquals((byte[]) restoredFieldSpec.getDefaultNullValue(), new byte[0]); + } + @Test public void testColumnMetadataEqualityIncludesForwardIndexEncoding() { DimensionFieldSpec fieldSpec = new DimensionFieldSpec("col", DataType.INT, true); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java index 1e38c4dfc698..17b5e4fb749e 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGeneratorTest.java @@ -24,6 +24,7 @@ import java.util.List; import org.apache.commons.configuration2.PropertiesConfiguration; import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.utils.VariantUtils; 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.segment.store.SegmentLocalFSDirectory; @@ -45,16 +46,18 @@ 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.assertTrue; -/// Regression tests for [ColumnMinMaxValueGenerator] on raw (no-dictionary) BYTES and UUID columns. +/// Regression tests for [ColumnMinMaxValueGenerator] on raw BYTES, UUID, and VARIANT columns. /// /// The generator reads the forward index when column metadata does not contain min/max values. These tests exercise /// the single-value, multi-value, and UUID (storedType=BYTES) forward index paths and verify their unsigned byte-wise /// ordering. The raw-BYTES loop once had its comparison directions inverted, silently persisting swapped min/max -/// metadata that value-based segment pruning then consumed. +/// metadata that value-based segment pruning then consumed. VARIANT is intentionally excluded because its physical +/// envelope bytes have no logical ordering. public class ColumnMinMaxValueGeneratorTest { private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), ColumnMinMaxValueGeneratorTest.class.getSimpleName()); @@ -62,6 +65,8 @@ public class ColumnMinMaxValueGeneratorTest { private static final String BYTES_COLUMN = "bytesCol"; private static final String BYTES_MV_COLUMN = "bytesMvCol"; private static final String UUID_COLUMN = "uuidCol"; + private static final String VARIANT_COLUMN = "variantCol"; + private static final String STRING_COLUMN = "stringCol"; // Ordered ascending by unsigned byte-wise comparison private static final byte[] BYTES_SMALL = new byte[]{0x00, 0x01}; @@ -117,6 +122,58 @@ public void testEmptyRawBytesMarksMinMaxInvalid() assertTrue(reloaded.getColumnMetadataFor(BYTES_COLUMN).isMinMaxValueInvalid()); } + @Test + public void testVariantExcludedFromMinMaxGeneration() + throws Exception { + for (ColumnMinMaxValueGeneratorMode mode + : List.of(ColumnMinMaxValueGeneratorMode.ALL, ColumnMinMaxValueGeneratorMode.NON_METRIC)) { + Schema schema = new Schema.SchemaBuilder().setSchemaName("variantSchema") + .setEnableColumnBasedNullHandling(true) + .addSingleValueDimension(VARIANT_COLUMN, DataType.VARIANT) + .addSingleValueDimension(STRING_COLUMN, DataType.STRING) + .build(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("variantTable") + .setNoDictionaryColumns(List.of(VARIANT_COLUMN)) + .setNullHandlingEnabled(true) + .build(); + GenericRow firstRow = new GenericRow(); + firstRow.putValue(VARIANT_COLUMN, VariantUtils.parseJsonToVariant("{\"value\":2}")); + firstRow.putValue(STRING_COLUMN, "z"); + GenericRow secondRow = new GenericRow(); + secondRow.putValue(VARIANT_COLUMN, VariantUtils.parseJsonToVariant("{\"value\":1}")); + secondRow.putValue(STRING_COLUMN, "a"); + + File indexDir = buildSegment(tableConfig, schema, List.of(firstRow, secondRow)); + removeMinMaxValuesFromMetadataFile(indexDir); + generateMinMaxValues(indexDir, mode); + + SegmentMetadataImpl reloaded = new SegmentMetadataImpl(indexDir); + assertNull(reloaded.getColumnMetadataFor(VARIANT_COLUMN).getMinValue()); + assertNull(reloaded.getColumnMetadataFor(VARIANT_COLUMN).getMaxValue()); + assertEquals(reloaded.getColumnMetadataFor(STRING_COLUMN).getMinValue(), "a"); + assertEquals(reloaded.getColumnMetadataFor(STRING_COLUMN).getMaxValue(), "z"); + } + } + + @Test + public void testSingleValueVariantIsNeverMarkedSorted() + throws Exception { + Schema schema = new Schema.SchemaBuilder().setSchemaName("variantSortedSchema") + .addSingleValueDimension(VARIANT_COLUMN, DataType.VARIANT) + .build(); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("variantSortedTable") + .setNoDictionaryColumns(List.of(VARIANT_COLUMN)) + .build(); + GenericRow row = new GenericRow(); + row.putValue(VARIANT_COLUMN, VariantUtils.parseJsonToVariant("{\"value\":1}")); + + File indexDir = buildSegment(tableConfig, schema, List.of(row)); + + SegmentMetadataImpl metadata = new SegmentMetadataImpl(indexDir); + assertFalse(metadata.getColumnMetadataFor(VARIANT_COLUMN).isSorted(), + "Physical envelope ordering must never be advertised as semantic VARIANT ordering"); + } + private static void assertMinMax(SegmentMetadataImpl segmentMetadata, String column, byte[] expectedMin, byte[] expectedMax) { ByteArray min = (ByteArray) segmentMetadata.getColumnMetadataFor(column).getMinValue(); @@ -128,11 +185,15 @@ private static void assertMinMax(SegmentMetadataImpl segmentMetadata, String col private static void generateMinMaxValues(File indexDir) throws Exception { + generateMinMaxValues(indexDir, ColumnMinMaxValueGeneratorMode.ALL); + } + + private static void generateMinMaxValues(File indexDir, ColumnMinMaxValueGeneratorMode mode) + throws Exception { SegmentMetadataImpl segmentMetadata = new SegmentMetadataImpl(indexDir); try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(indexDir, segmentMetadata, ReadMode.mmap); SegmentDirectory.Writer writer = segmentDirectory.createWriter()) { - ColumnMinMaxValueGenerator generator = - new ColumnMinMaxValueGenerator(segmentMetadata, writer, ColumnMinMaxValueGeneratorMode.ALL); + ColumnMinMaxValueGenerator generator = new ColumnMinMaxValueGenerator(segmentMetadata, writer, mode); generator.addColumnMinMaxValue(); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantSchemaValidationTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantSchemaValidationTest.java new file mode 100644 index 000000000000..b17c9a66772b --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantSchemaValidationTest.java @@ -0,0 +1,134 @@ +/** + * 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.segment.local.utils; + +import java.nio.ByteBuffer; +import org.apache.pinot.spi.data.DimensionFieldSpec; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.FieldSpec.MaxLengthExceedStrategy; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.PinotDataType; +import org.apache.pinot.spi.utils.VariantEnvelope; +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.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +public class VariantSchemaValidationTest { + private static final String COLUMN = "payload"; + + @Test + public void testValidSingleValueVariantDimension() { + Schema schema = schemaWith(new DimensionFieldSpec(COLUMN, DataType.VARIANT, true)); + SchemaUtils.validate(schema); + + FieldSpec fieldSpec = schema.getFieldSpecFor(COLUMN); + assertEquals((byte[]) fieldSpec.getDefaultNullValue(), new byte[0]); + assertEquals(PinotDataType.getPinotDataTypeForIngestion(fieldSpec), PinotDataType.VARIANT); + assertNull(SanitizationTransformerUtils.getSanitizedColumnInfo(fieldSpec)); + } + + @Test + public void testRejectsMultiValueVariant() { + Schema schema = schemaWith(new DimensionFieldSpec(COLUMN, DataType.VARIANT, false)); + IllegalStateException exception = + expectThrows(IllegalStateException.class, () -> SchemaUtils.validate(schema)); + assertTrue(exception.getMessage().contains("VARIANT columns cannot be of multi-value type")); + assertThrows(IllegalStateException.class, + () -> PinotDataType.getPinotDataTypeForIngestion(schema.getFieldSpecFor(COLUMN))); + } + + @Test + public void testRejectsCustomDefaultNullValue() { + byte[] encodedVariantNull = + VariantEnvelope.encode(ByteBuffer.wrap(new byte[]{1}), ByteBuffer.wrap(new byte[]{0})); + Schema schema = + schemaWith(new DimensionFieldSpec(COLUMN, DataType.VARIANT, true, encodedVariantNull)); + + IllegalStateException exception = + expectThrows(IllegalStateException.class, () -> SchemaUtils.validate(schema)); + assertTrue(exception.getMessage().contains("custom default null value")); + } + + @Test + public void testRejectsEnvelopeCorruptingMaxLengthStrategies() { + for (MaxLengthExceedStrategy strategy + : new MaxLengthExceedStrategy[]{MaxLengthExceedStrategy.TRIM_LENGTH, + MaxLengthExceedStrategy.SUBSTITUTE_DEFAULT_VALUE}) { + DimensionFieldSpec fieldSpec = new DimensionFieldSpec(COLUMN, DataType.VARIANT, true); + fieldSpec.setMaxLength(20); + fieldSpec.setMaxLengthExceedStrategy(strategy); + Schema schema = schemaWith(fieldSpec); + + IllegalStateException exception = + expectThrows(IllegalStateException.class, () -> SchemaUtils.validate(schema)); + assertTrue(exception.getMessage().contains("max length strategy")); + assertThrows(IllegalStateException.class, + () -> SanitizationTransformerUtils.getSanitizedColumnInfo(fieldSpec)); + } + } + + @Test + public void testErrorMaxLengthStrategyNeverMutatesEnvelope() { + DimensionFieldSpec fieldSpec = new DimensionFieldSpec(COLUMN, DataType.VARIANT, true); + fieldSpec.setMaxLength(32); + fieldSpec.setMaxLengthExceedStrategy(MaxLengthExceedStrategy.ERROR); + SchemaUtils.validate(schemaWith(fieldSpec)); + + SanitizationTransformerUtils.SanitizedColumnInfo columnInfo = + SanitizationTransformerUtils.getSanitizedColumnInfo(fieldSpec); + assertNotNull(columnInfo); + byte[] envelope = + VariantEnvelope.encode(ByteBuffer.wrap(new byte[]{1}), ByteBuffer.wrap(new byte[]{0})); + SanitizationTransformerUtils.SanitizationResult result = + SanitizationTransformerUtils.sanitizeValue(columnInfo, envelope); + assertNotNull(result); + assertSame(result.getValue(), envelope); + assertFalse(result.isSanitized()); + + fieldSpec.setMaxLength(VariantEnvelope.HEADER_SIZE); + columnInfo = SanitizationTransformerUtils.getSanitizedColumnInfo(fieldSpec); + SanitizationTransformerUtils.SanitizedColumnInfo finalColumnInfo = columnInfo; + assertThrows(IllegalStateException.class, + () -> SanitizationTransformerUtils.sanitizeValue(finalColumnInfo, envelope)); + } + + @Test + public void testIngestionConversionValidatesPvarEnvelope() { + byte[] envelope = + VariantEnvelope.encode(ByteBuffer.wrap(new byte[]{1}), ByteBuffer.wrap(new byte[]{0})); + assertSame(DataTypeTransformerUtils.transformValue(COLUMN, envelope, PinotDataType.VARIANT), envelope); + assertThrows(IllegalArgumentException.class, + () -> DataTypeTransformerUtils.transformValue(COLUMN, new byte[0], PinotDataType.VARIANT)); + assertThrows(IllegalArgumentException.class, + () -> DataTypeTransformerUtils.transformValue(COLUMN, new byte[]{1, 2}, PinotDataType.VARIANT)); + } + + private static Schema schemaWith(FieldSpec fieldSpec) { + return new Schema.SchemaBuilder().setSchemaName("variantSchema").addField(fieldSpec).build(); + } +} diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java new file mode 100644 index 000000000000..cab555d446be --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java @@ -0,0 +1,300 @@ +/** + * 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.segment.local.utils; + +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.util.List; +import java.util.Map; +import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.FieldConfig.CompressionCodec; +import org.apache.pinot.spi.config.table.FieldConfig.EncodingType; +import org.apache.pinot.spi.config.table.FieldConfig.IndexType; +import org.apache.pinot.spi.config.table.RoutingConfig; +import org.apache.pinot.spi.config.table.SegmentPartitionConfig; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.UpsertConfig; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.fail; + + +public class VariantTableConfigValidationTest { + private static final String TABLE_NAME = "variantTable"; + private static final String VARIANT_COLUMN = "payload"; + private static final String ID_COLUMN = "id"; + + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(ID_COLUMN, DataType.STRING) + .addSingleValueDimension(VARIANT_COLUMN, DataType.VARIANT) + .build(); + + private static final Schema UPSERT_SCHEMA = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(ID_COLUMN, DataType.STRING) + .addSingleValueDimension(VARIANT_COLUMN, DataType.VARIANT) + .setPrimaryKeyColumns(List.of(ID_COLUMN)) + .build(); + + @Test + public void testRawForwardIndexWithCompressionAndNullVectorIsValid() { + TableConfig tableConfig = tableWith(rawVariantFieldConfig()); + assertValid(tableConfig, SCHEMA); + } + + @Test + public void testVariantRequiresEffectiveStorageNullHandling() { + TableConfig nullHandlingDisabled = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .build(); + assertInvalid(nullHandlingDisabled, SCHEMA, "Null handling must be enabled for tables containing VARIANT columns"); + + Schema columnBasedNullHandlingSchema = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .setEnableColumnBasedNullHandling(true) + .addSingleValueDimension(ID_COLUMN, DataType.STRING) + .addSingleValueDimension(VARIANT_COLUMN, DataType.VARIANT) + .build(); + TableConfig columnBasedNullHandlingTable = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .build(); + assertValid(columnBasedNullHandlingTable, columnBasedNullHandlingSchema); + + Schema nonVariantSchema = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(ID_COLUMN, DataType.STRING) + .build(); + TableConfig nonVariantTable = + new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build(); + assertValid(nonVariantTable, nonVariantSchema); + } + + @Test + public void testDictionaryEncodingIsRejected() { + FieldConfig fieldConfig = new FieldConfig.Builder(VARIANT_COLUMN) + .withEncodingType(EncodingType.DICTIONARY) + .build(); + assertInvalid(tableWith(fieldConfig), SCHEMA, "RAW forward index encoding"); + } + + @Test + public void testImplicitDictionaryEncodingIsRejected() { + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .build(); + assertInvalid(tableConfig, SCHEMA, "RAW forward index encoding"); + } + + @Test + public void testExplicitDictionaryWithRawForwardIndexIsRejected() { + ObjectNode indexes = JsonUtils.newObjectNode(); + indexes.set("dictionary", JsonUtils.newObjectNode()); + FieldConfig fieldConfig = new FieldConfig.Builder(VARIANT_COLUMN) + .withEncodingType(EncodingType.RAW) + .withIndexes(indexes) + .build(); + assertInvalid(tableWith(fieldConfig), SCHEMA, "cannot use a dictionary"); + } + + @Test + public void testDisabledForwardIndexIsRejected() { + FieldConfig fieldConfig = new FieldConfig.Builder(VARIANT_COLUMN) + .withEncodingType(EncodingType.RAW) + .withProperties(Map.of(FieldConfig.FORWARD_INDEX_DISABLED, "true")) + .build(); + assertInvalid(tableWith(fieldConfig), SCHEMA, "enabled forward index"); + } + + @Test + public void testSecondaryIndexesAreRejected() { + for (IndexType indexType + : new IndexType[]{IndexType.INVERTED, IndexType.FST, IndexType.IFST, IndexType.TEXT, IndexType.JSON, + IndexType.RANGE}) { + FieldConfig fieldConfig = new FieldConfig.Builder(VARIANT_COLUMN) + .withEncodingType(EncodingType.RAW) + .withIndexTypes(List.of(indexType)) + .build(); + // Dictionary-backed index types can make the effective dictionary configuration fail first. + assertInvalid(tableWith(fieldConfig), SCHEMA, "VARIANT column"); + } + + TableConfig bloomTable = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .setBloomFilterColumns(List.of(VARIANT_COLUMN)) + .build(); + assertInvalid(bloomTable, SCHEMA, "supports only a RAW forward index"); + } + + @Test + public void testSortedAndPartitionKeysAreRejected() { + TableConfig sortedTable = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .setSortedColumn(VARIANT_COLUMN) + .build(); + assertInvalid(sortedTable, SCHEMA, "Cannot sort on VARIANT column"); + + SegmentPartitionConfig partitionConfig = new SegmentPartitionConfig( + Map.of(VARIANT_COLUMN, new ColumnPartitionConfig("Murmur", 4))); + TableConfig partitionedTable = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .setSegmentPartitionConfig(partitionConfig) + .build(); + assertInvalid(partitionedTable, SCHEMA, "Cannot partition on VARIANT column"); + } + + @Test + public void testPrimaryKeyIsRejected() { + Schema schema = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(ID_COLUMN, DataType.STRING) + .addSingleValueDimension(VARIANT_COLUMN, DataType.VARIANT) + .setPrimaryKeyColumns(List.of(VARIANT_COLUMN)) + .build(); + assertInvalid(tableWith(rawVariantFieldConfig()), schema, "cannot be used as a primary key"); + } + + @Test + public void testUpsertComparisonColumnIsRejected() { + UpsertConfig upsertConfig = new UpsertConfig(UpsertConfig.Mode.FULL); + upsertConfig.setComparisonColumn(VARIANT_COLUMN); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName(TABLE_NAME) + .setUpsertConfig(upsertConfig) + .setRoutingConfig( + new RoutingConfig(null, null, RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false)) + .build(); + + try { + TableConfigUtils.validateUpsertAndDedupConfig(tableConfig, UPSERT_SCHEMA); + fail("Expected VARIANT upsert comparison column validation to fail"); + } catch (IllegalStateException e) { + assertTrue(e.getMessage() != null + && e.getMessage().contains("VARIANT column cannot be used as an upsert comparison column"), + "Unexpected validation error: " + e.getMessage()); + } + } + + @Test + public void testNonOverwritePartialUpsertStrategyIsRejected() { + TableConfig tableConfig = partialUpsertTable(UpsertConfig.Strategy.IGNORE); + try { + TableConfigUtils.validatePartialUpsertStrategies(tableConfig, UPSERT_SCHEMA); + fail("Expected non-OVERWRITE VARIANT partial-upsert strategy validation to fail"); + } catch (IllegalStateException e) { + assertTrue(e.getMessage() != null + && e.getMessage().contains("VARIANT column supports only OVERWRITE partial-upsert strategy"), + "Unexpected validation error: " + e.getMessage()); + } + } + + @Test + public void testOverwritePartialUpsertStrategyIsValid() { + TableConfigUtils.validatePartialUpsertStrategies( + partialUpsertTable(UpsertConfig.Strategy.OVERWRITE), UPSERT_SCHEMA); + } + + @Test + public void testMetricsAggregationAndDefaultStarTreeAreRejected() { + TableConfig aggregateMetricsTable = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .setAggregateMetrics(true) + .build(); + assertInvalid(aggregateMetricsTable, SCHEMA, "VARIANT dimension column as an aggregation key"); + + TableConfig defaultStarTreeTable = tableWith(rawVariantFieldConfig()); + defaultStarTreeTable.getIndexingConfig().setEnableDefaultStarTree(true); + assertInvalid(defaultStarTreeTable, SCHEMA, "Default star-tree index cannot include VARIANT column"); + } + + @Test + public void testExplicitStarTreeIsRejected() { + StarTreeIndexConfig starTreeIndexConfig = + new StarTreeIndexConfig(List.of(ID_COLUMN), null, List.of("SUM__" + VARIANT_COLUMN), null, 1); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(rawVariantFieldConfig())) + .setStarTreeIndexConfigs(List.of(starTreeIndexConfig)) + .build(); + assertInvalid(tableConfig, SCHEMA, "Star-tree index cannot be created on VARIANT column"); + } + + private static FieldConfig rawVariantFieldConfig() { + return new FieldConfig.Builder(VARIANT_COLUMN) + .withEncodingType(EncodingType.RAW) + .withCompressionCodec(CompressionCodec.ZSTANDARD) + .build(); + } + + private static TableConfig tableWith(FieldConfig fieldConfig) { + return new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(fieldConfig)) + .build(); + } + + private static TableConfig partialUpsertTable(UpsertConfig.Strategy strategy) { + UpsertConfig upsertConfig = new UpsertConfig(UpsertConfig.Mode.PARTIAL); + upsertConfig.setComparisonColumn(ID_COLUMN); + upsertConfig.setPartialUpsertStrategies(Map.of(VARIANT_COLUMN, strategy)); + return new TableConfigBuilder(TableType.REALTIME) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setUpsertConfig(upsertConfig) + .build(); + } + + private static void assertValid(TableConfig tableConfig, Schema schema) { + try { + TableConfigUtils.validate(tableConfig, schema); + } catch (Exception e) { + fail("Expected validation to pass, but got: " + e.getMessage(), e); + } + } + + private static void assertInvalid(TableConfig tableConfig, Schema schema, String messageFragment) { + try { + TableConfigUtils.validate(tableConfig, schema); + fail("Expected validation failure containing: " + messageFragment); + } catch (Exception e) { + assertTrue(e.getMessage() != null && e.getMessage().contains(messageFragment), + "Expected '" + messageFragment + "' in error, but got: " + e.getMessage()); + } + } +} diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java index 71dc1003d760..01e3383519b7 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java @@ -487,6 +487,12 @@ public static FieldSpec extractFieldSpec(String column, PropertiesConfiguration DataType dataType = config.getEnum(Column.getKeyFor(column, Column.DATA_TYPE), DataType.class); boolean isSingleValue = config.getBoolean(Column.getKeyFor(column, Column.IS_SINGLE_VALUED), true); String defaultNullValueString = config.getString(Column.getKeyFor(column, Column.DEFAULT_NULL_VALUE), null); + if (dataType == DataType.VARIANT && "".equals(defaultNullValueString)) { + // An empty byte array is VARIANT's reserved default-null sentinel. Segment metadata persists byte arrays as hex, + // so the sentinel is written as an empty string. Treat it as the type default instead of asking VARIANT.convert() + // to decode it as a PVAR value. + defaultNullValueString = null; + } if (defaultNullValueString != null && dataType.getStoredType() == DataType.STRING) { defaultNullValueString = CommonsConfigurationUtils.recoverSpecialCharacterInPropertyValue(defaultNullValueString); } diff --git a/pinot-spi/VARIANT_DESIGN.md b/pinot-spi/VARIANT_DESIGN.md new file mode 100644 index 000000000000..22eab9a70d45 --- /dev/null +++ b/pinot-spi/VARIANT_DESIGN.md @@ -0,0 +1,330 @@ + +# Apache Pinot VARIANT — Design Document + +Status: draft for community and component-owner review. + +## 1. Motivation and goals + +Apache Parquet `VARIANT(1)` provides a portable binary representation for +semi-structured values. Pinot should ingest that representation without converting +every row through a JSON tree, retain the complete value for late-bound queries, and +let users materialize frequently queried paths into ordinary Pinot columns. + +This design has five goals: + +1. Define a first-class Pinot `VARIANT` logical type with an explicit persisted and + query-wire contract. +2. Ingest top-level, non-repeated Parquet `VARIANT(1)` values, including unshredded, + fully shredded, and partially shredded layouts. +3. Provide Spark-style parse, extraction, existence, type, null, and JSON-rendering + functions in both Pinot query engines. +4. Fail explicitly when a raw Variant value reaches an operation whose semantics are + undefined, while allowing users to extract a typed scalar for that operation. +5. Ship an executable table-definition, ingestion, and query story with integration + coverage. + +### 1.1 Non-goals + +The initial contract does not include: + +- nested or repeated Parquet Variant columns; +- Pinot multi-value `VARIANT` columns; +- streaming-format Variant ingestion; +- quoted object keys in Variant paths; +- semantic equality, hashing, or ordering of raw Variant values; +- mixed-version execution for queries that contain a Variant type; or +- automatic rollback after a Variant table has been activated. + +## 2. User model + +A Variant column is declared as a single-value dimension: + +```json +{ + "name": "payload", + "dataType": "VARIANT" +} +``` + +The table stores `payload` in a raw forward index without a dictionary. Storage null +handling must be enabled. A common layout retains the source value while materializing +an indexed scalar during ingestion: + +```json +{ + "columnName": "eventType", + "transformFunction": "variant_get(payload, '$.eventType', 'STRING')" +} +``` + +Late-bound paths remain queryable: + +```sql +SET enableNullHandling=true; + +SELECT + eventId, + variant_get(payload, '$.user.id', 'STRING') AS userId, + variant_get(payload, '$.amount', 'DOUBLE') AS amount +FROM variantEvents +WHERE eventType = 'checkout'; +``` + +## 3. Data flow and ownership + +```text +Parquet VARIANT(1) + | + | pinot-parquet: validate shape and reconstruct shredded values + v +metadata ByteBuffer + value ByteBuffer + | + | VariantEnvelope: add stable Pinot PVAR framing + v +single-value raw BYTES forward index in a Pinot segment + | + | DataSchema/protobuf: carry logical VARIANT, not generic BYTES + v +single-stage or multi-stage Variant functions + | + +--> typed scalar for filtering/grouping/joining/aggregation + | + +--> canonical JSON text for response clients +``` + +Ownership is deliberately split by module: + +- `pinot-spi` owns the public logical type, capability contract, and stable `PVAR` + envelope. +- `pinot-segment-local` owns schema, table, index, null, and segment-metadata + validation. +- `pinot-parquet` owns Parquet logical-type validation and reconstruction. +- `pinot-common`, `pinot-core`, `pinot-query-planner`, and `pinot-query-runtime` own + function semantics and defensive query guards. +- client, response, and DDL modules carry the logical type without redefining its + semantics. + +`pinot-common` depends only on the small `parquet-variant` codec. Parquet column, +schema, Hadoop, and reader dependencies remain isolated to the input-format plugin. + +## 4. Persisted `PVAR` envelope + +Pinot segments need one byte sequence that contains both buffers in a Parquet Variant +value. `VariantEnvelope` owns that framing. Version 1 is: + +```text +offset size field +0 4 ASCII magic "PVAR" +4 1 envelope version (1) +5 1 flags (0) +6 2 reserved (0) +8 4 metadata length, big-endian signed int restricted to >= 0 +12 4 value length, big-endian signed int restricted to >= 0 +16 M Parquet Variant metadata bytes +16 + M V Parquet Variant value bytes +``` + +The complete envelope must fit in a Java byte array. Decoders reject unknown versions, +non-zero flags or reserved bytes, negative lengths, bad magic, and length mismatches. +Metadata and value bytes are copied without interpretation by the envelope layer; the +Parquet Variant codec validates their internal representation. + +Version 1 is immutable and protected by a golden-byte test. A future incompatible +framing change must use a new envelope version and retain an explicit reader for every +supported old version. Unknown versions fail closed; they must never fall back to +generic `BYTES`. + +`decode(byte[])` returns zero-copy read-only views that alias the input array. The input +must therefore remain immutable for the lifetime of those views. Array, direct, and +read-only source buffers are supported without changing their positions. Parquet +ingestion copies payload bytes directly into the final envelope and must not allocate a +full intermediate payload array. + +## 5. Logical type and segment contract + +`FieldSpec.DataType.VARIANT` is logically distinct from `BYTES` but uses variable-width +bytes as its physical stored type. + +The initial schema and table constraints are: + +- dimension field only; +- single-value only; +- the reserved empty-byte default null value only; +- no truncating or default-substitution max-length policy; +- storage null handling enabled at schema or table level; +- enabled raw forward index; +- no dictionary or secondary index on the raw Variant column; +- not a primary key, partition column, sorted column, upsert comparison column, + metrics-aggregation key, or star-tree dimension; and +- partial upsert may use only the `OVERWRITE` strategy. + +Segment creation validates every non-null value as a complete `PVAR` envelope. Raw +Variant columns do not publish logical minimum or maximum values and are never marked +sorted, including single-row segments. + +### 5.1 Capability policy + +The Parquet representation is not a canonical semantic encoding: equivalent logical +objects can use different metadata dictionaries or shredded layouts. Byte equality, +byte hashing, and lexicographic byte order would therefore expose storage-layout +accidents as SQL semantics. + +The logical type reports: + +| Capability on raw value | Supported | +|---|---| +| equality / `IN` / pattern predicates | no | +| hashing / `GROUP BY` / `DISTINCT` | no | +| ordering / min-max / sort / ASOF match key | no | +| join, lookup, primary, or partition key | no | +| direct value-consuming aggregation | no | +| non-distinct `COUNT(raw_variant)` | yes | + +Planner validation provides the primary error. Runtime and single-stage guards remain +in place for mixed plans or future planner changes. Error messages direct users to +`variant_get` a typed path first. + +## 6. Parquet ingestion + +Reader initialization scans the file schema before publishing reader state. + +- A supported field is a top-level, non-repeated group annotated `VARIANT(1)`. +- `metadata` is required binary; at least one of `value` or `typed_value` is present. +- The parquet-java `VariantConverters` tree reconstructs fully or partially shredded + values. +- The unshredded path retains the exact metadata and value buffers and adds only the + `PVAR` envelope. +- Converter plans and field indexes are immutable and created once per reader. +- A failed reinitialization leaves the previous reader usable. A failure while closing + the old reader does not invalidate the successfully initialized replacement. + +Automatic reader selection preserves existing Avro-metadata precedence for backward +compatibility. A Parquet file with Avro metadata must explicitly select the native +reader to ingest Variant. + +Malformed rows follow the existing record-reader skip policy and do not corrupt reader +progress. Unsupported Variant spec versions and unsupported nested/repeated shapes fail +during initialization rather than appearing as ordinary structs. + +## 7. Query functions + +Both query engines expose: + +- `variant_get(value, path[, targetType])` +- `try_variant_get(value, path[, targetType])` +- `variant_exists(value, path)` +- `is_variant_null(value[, path])` +- `variant_typeof(value[, path])` +- `variant_to_json(value)` +- `parse_json(text)` / `parse_json_to_variant(text)` +- `try_parse_json(text)` / `try_parse_json_to_variant(text)` + +Function lookup remains case-insensitive under Pinot's existing registration rules. +The v1 path grammar supports root `$`, dot-separated object fields, and non-negative +array indexes. Supported extraction targets are `BOOLEAN`, `INT`, `LONG`, `FLOAT`, +`DOUBLE`, `BIG_DECIMAL`, `STRING`, `BYTES`, `UUID`, `TIMESTAMP`, `VARIANT`, and `JSON`. + +Strict functions reject malformed encodings, paths, incompatible conversions, numeric +overflow, and malformed JSON. `try_` functions return SQL null for those failures. +Literal paths, target types, and JSON inputs are compiled or parsed once per query. +Vectorized execution reuses thread-confined cursors and unboxed result holders. + +JSON parsing is bounded to 100 nesting levels. Parquet Variant decimal encoding is +bounded to 38 digits of precision and a scale of at most 38 after exact normalization. +Exponent bounds are checked before expansion so hostile inputs cannot request +pathological intermediate allocations. + +### 7.1 Null states + +Storage and query null handling are required because four states must remain distinct: + +| State | Stored bytes / null vector | `variant_exists` | `is_variant_null` | `variant_typeof` | scalar `variant_get` | +|---|---|---:|---:|---|---| +| SQL null | empty bytes, null bit set | SQL null | false | SQL null | SQL null | +| Variant null | non-empty `PVAR`, null bit clear | true | true | `NULL` | SQL null | +| missing path | valid non-empty `PVAR` | false | false | SQL null | SQL null | +| Variant string `"null"` | non-empty `PVAR` | true | false | `STRING` | `"null"` | + +Extracting a Variant null as `VARIANT` retains its non-empty envelope. +`variant_to_json` returns SQL null for SQL null, JSON text `null` for Variant null, and +JSON text `"null"` (including quotes) for the Variant string. + +## 8. Query wire and client compatibility + +The protobuf expression enum assigns `VARIANT` permanent number `24`. Numbers `22` and +`23` remain reserved for the separately allocated UUID contract. Existing numbers are +not renumbered. A frozen pre-Variant proto test verifies that an old peer reads value +`24` as `UNRECOGNIZED`, while a current peer continues to read every legacy value. + +Data blocks and `DataSchema` carry the logical `VARIANT` token, with a `ByteArray` +internal representation. JSON, Arrow, Java, HTTP JDBC, and gRPC JDBC response paths +preserve the distinction among SQL null, Variant null, and the Variant string `"null"`. + +There is no type negotiation or safe downgrade to `BYTES`. An old node can read +existing non-Variant traffic, but it cannot plan or execute an active Variant query. + +## 9. Deployment and rollback + +This feature has an activation gate: + +1. Upgrade controllers, brokers, servers, minions, clients, and every external segment + generation or ingestion job. +2. Verify all processes use the Variant-capable build. +3. Only then register a schema containing `VARIANT` or upload a segment containing a + `PVAR` column. + +Existing schemas, segments, and queries are unaffected during the rolling binary +upgrade. Variant columns must remain inactive until the fleet is homogeneous. + +Do not roll back to a pre-Variant binary while a Variant schema, segment, or in-flight +query is active. A rollback first requires draining Variant queries and removing or +migrating every active Variant table and segment. + +## 10. Review and delivery decomposition + +The implementation touches multiple ownership areas because no partial state is safe +to activate. The review dependency order is: + +1. public type, capability policy, `PVAR` framing, and query-wire number; +2. schema/table/index validation and Parquet ingestion; +3. function semantics plus single-stage and multi-stage guards; +4. DDL, response, Java, JDBC, gRPC, JSON, and Arrow propagation; and +5. quickstart, integration tests, compatibility tests, and benchmarks. + +The initial implementation is presented as one end-to-end draft so reviewers can +evaluate one activation contract and run one acceptance test. It is not merge-ready +until SPI/wire, Parquet, both query engines, and client/response owners approve their +areas. If maintainers prefer independent rollback units, the five groups above form a +dependency-ordered PR stack; every intermediate PR must compile, keep Variant +unactivatable until the safety guards land, and preserve the permanent wire allocation. + +## 11. Verification and future work + +The acceptance suite creates a table, ingests a real Parquet `VARIANT(1)` fixture, and +queries it through both engines. Unit suites cover envelope golden bytes, old/new proto +behavior, direct/read-only buffers, unshredded and shredded reconstruction, decimal +bounds, null states, every raw-value guard, clients, DDL, and reader lifecycle. + +Future proposals can independently address nested/repeated Variant, streaming formats, +quoted path keys, indexes over materialized subpaths, a canonical semantic equality +contract, and version negotiation. None may silently widen the v1 persisted or raw-query +semantics described here. diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java index de37cd58d3df..4fe26998482e 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/FieldSpec.java @@ -47,6 +47,7 @@ import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.TimestampUtils; import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.VariantEnvelope; /// The `FieldSpec` class contains all specs related to any field (column) in [Schema]. @@ -525,6 +526,8 @@ public static Object getDefaultNullValue(FieldType fieldType, DataType dataType, return DEFAULT_DIMENSION_NULL_VALUE_OF_JSON; case BYTES: return DEFAULT_DIMENSION_NULL_VALUE_OF_BYTES; + case VARIANT: + return DEFAULT_DIMENSION_NULL_VALUE_OF_BYTES; case UUID: // The nil UUID is the default null sentinel. Tables that ingest nil UUID as a real value should enable // column-based null handling to distinguish it from null rows. @@ -654,7 +657,11 @@ protected void appendFieldIdAndAliases(ObjectNode jsonObject) { protected void appendDefaultNullValue(ObjectNode jsonNode) { assert _defaultNullValue != null; String key = "defaultNullValue"; - if (!_defaultNullValue.equals(getDefaultNullValue(getFieldType(), _dataType, null))) { + Object typeDefaultNullValue = getDefaultNullValue(getFieldType(), _dataType, null); + boolean usesTypeDefault = _dataType == DataType.VARIANT + ? Arrays.equals((byte[]) _defaultNullValue, (byte[]) typeDefaultNullValue) + : _defaultNullValue.equals(typeDefaultNullValue); + if (!usesTypeDefault) { switch (_dataType) { case INT: jsonNode.put(key, (Integer) _defaultNullValue); @@ -684,6 +691,9 @@ protected void appendDefaultNullValue(ObjectNode jsonNode) { case BYTES: jsonNode.put(key, BytesUtils.toHexString((byte[]) _defaultNullValue)); break; + case VARIANT: + jsonNode.put(key, BytesUtils.toHexString((byte[]) _defaultNullValue)); + break; case UUID: jsonNode.put(key, UuidUtils.toString((byte[]) _defaultNullValue)); break; @@ -720,7 +730,7 @@ public boolean equals(Object o) { && Objects.equals(_maxLength, that._maxLength) && Objects.equals(_maxLengthExceedStrategy, that._maxLengthExceedStrategy) && _allowTrailingZeros == that._allowTrailingZeros - && _dataType.equals(_defaultNullValue, that._defaultNullValue) + && defaultNullValuesEqual(that) && Objects.equals(_transformFunction, that._transformFunction) && Objects.equals(_virtualColumnProvider, that._virtualColumnProvider) && Objects.equals(_description, that._description) @@ -733,10 +743,24 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(_name, _dataType, _singleValueField, _notNull, _maxLength, _maxLengthExceedStrategy, - _allowTrailingZeros, _dataType.hashCode(_defaultNullValue), _transformFunction, _virtualColumnProvider, + _allowTrailingZeros, defaultNullValueHashCode(), _transformFunction, _virtualColumnProvider, _description, _tags, _fieldId, _aliases, _metadata); } + private boolean defaultNullValuesEqual(FieldSpec that) { + // VARIANT deliberately rejects semantic value equality. FieldSpec equality only compares schema configuration, + // where the reserved byte-array null sentinel is structural metadata rather than a queryable Variant value. + return _dataType == DataType.VARIANT + ? Arrays.equals((byte[]) _defaultNullValue, (byte[]) that._defaultNullValue) + : _dataType.equals(_defaultNullValue, that._defaultNullValue); + } + + private int defaultNullValueHashCode() { + // Keep this paired with defaultNullValuesEqual() without enabling semantic hashing for VARIANT values. + return _dataType == DataType.VARIANT ? Arrays.hashCode((byte[]) _defaultNullValue) + : _dataType.hashCode(_defaultNullValue); + } + /// The `FieldType` enum is used to demonstrate the real world business logic for a column. /// /// `DIMENSION`: columns used to filter records. @@ -764,13 +788,13 @@ public enum FieldType { /// - `BOOLEAN` → [Integer] (`0` or `1`) /// - `TIMESTAMP` → [Long] (epoch millis) /// - `STRING` / `JSON` → [String] - /// - `BYTES` → `byte[]` + /// - `BYTES` / `VARIANT` → `byte[]` /// - `UUID` → `byte[]` (fixed 16-byte big-endian form) /// - `MAP` / `OPEN_STRUCT` → [Map] /// - `LIST` → [List] /// - /// `convertInternal` is the exception: it returns the internal storage form, which for `BYTES` and `UUID` is - /// [ByteArray] rather than `byte[]`. + /// `convertInternal` is the exception: it returns the internal storage form, which for `BYTES`, `VARIANT`, and + /// `UUID` is [ByteArray] rather than `byte[]`. @SuppressWarnings("rawtypes") public enum DataType { // LIST is for complex lists which is different from multi-value column of primitives @@ -791,7 +815,9 @@ public enum DataType { MAP(false), OPEN_STRUCT(false), LIST(false), - UNKNOWN(false); + UNKNOWN(false), + // VARIANT is a logical type stored as variable-width BYTES in a Pinot-owned PVAR envelope. + VARIANT(BYTES, false); private final DataType _storedType; private final int _size; @@ -857,7 +883,53 @@ public boolean isUnknown() { return _storedType == UNKNOWN; } - /// Converts the given string value to the data type. Returns byte\[\] for BYTES and UUID. + /// Returns whether generic operators may apply equality to raw values of this logical type. + public boolean supportsEquality() { + return this != VARIANT; + } + + /// Returns whether generic operators may hash raw values of this logical type. + public boolean supportsHashing() { + return this != VARIANT; + } + + /// Returns whether physical stored ordering is also a valid ordering for raw values of this logical type. + public boolean supportsOrdering() { + switch (this) { + case STRUCT: + case MAP: + case OPEN_STRUCT: + case LIST: + case UNKNOWN: + case VARIANT: + return false; + default: + return true; + } + } + + /// Returns whether segment metadata may derive logical minimum and maximum values from this type. + public boolean supportsMinMax() { + return supportsOrdering(); + } + + /// Returns whether generic aggregation functions may consume raw values of this logical type. + /// + /// Value-independent functions such as non-distinct `COUNT` do not consume the value and can still be + /// used when this returns `false`. + public boolean supportsDirectAggregation() { + return this != VARIANT; + } + + /// Returns whether generic pattern predicates may consume raw values of this logical type. + /// + /// This expresses logical-type restrictions only. Predicate type checking separately ensures that the stored + /// value is usable as text. + public boolean supportsPatternMatching() { + return this != VARIANT; + } + + /// Converts the given string value to the data type. Returns byte\[\] for BYTES, VARIANT and UUID. public Object convert(String value) { try { switch (this) { @@ -880,6 +952,10 @@ public Object convert(String value) { return value; case BYTES: return BytesUtils.toBytes(value); + case VARIANT: + byte[] envelope = BytesUtils.toBytes(value); + VariantEnvelope.decode(envelope); + return envelope; case UUID: return UuidUtils.toBytes(value); case MAP: @@ -896,10 +972,16 @@ public Object convert(String value) { } public boolean equals(Object value1, Object value2) { + if (!supportsEquality()) { + throw new UnsupportedOperationException(this + " does not support equality"); + } return this == BYTES || this == UUID ? Arrays.equals((byte[]) value1, (byte[]) value2) : value1.equals(value2); } public int hashCode(Object value) { + if (!supportsHashing()) { + throw new UnsupportedOperationException(this + " does not support hashing"); + } return this == BYTES || this == UUID ? Arrays.hashCode((byte[]) value) : value.hashCode(); } @@ -909,6 +991,9 @@ public int hashCode(Object value) { /// return -1 if value1 is less than value2 /// return 1 if value1 is greater than value2 public int compare(Object value1, Object value2) { + if (!supportsOrdering()) { + throw new UnsupportedOperationException(this + " does not support ordering"); + } switch (this) { case INT: case BOOLEAN: @@ -928,6 +1013,7 @@ public int compare(Object value1, Object value2) { case BYTES: case UUID: return ByteArray.compare((byte[]) value1, (byte[]) value2); + case VARIANT: case MAP: case OPEN_STRUCT: case LIST: @@ -945,6 +1031,9 @@ public String toString(Object value) { if (this == BYTES) { return BytesUtils.toHexString((byte[]) value); } + if (this == VARIANT) { + return BytesUtils.toHexString((byte[]) value); + } if (this == UUID) { return UuidUtils.toString((byte[]) value); } @@ -958,7 +1047,7 @@ public String toString(Object value) { return value.toString(); } - /// Converts the given string value to the data type. Returns ByteArray for BYTES and UUID. + /// Converts the given string value to the data type. Returns ByteArray for BYTES, VARIANT and UUID. public Comparable convertInternal(String value) { try { switch (this) { @@ -981,6 +1070,10 @@ public Comparable convertInternal(String value) { return value; case BYTES: return BytesUtils.toByteArray(value); + case VARIANT: + byte[] envelope = BytesUtils.toBytes(value); + VariantEnvelope.decode(envelope); + return new ByteArray(envelope); case UUID: return new ByteArray(UuidUtils.toBytes(value)); case MAP: diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java b/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java index 2d0c9bd88843..906a87f331ab 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/data/Schema.java @@ -123,6 +123,10 @@ public static Schema fromInputStream(InputStream schemaInputStream) public static void validate(FieldType fieldType, DataType dataType) { switch (fieldType) { case DIMENSION: + if (dataType == DataType.VARIANT) { + break; + } + // fall through case TIME: case DATE_TIME: switch (dataType) { @@ -591,7 +595,7 @@ public String toSingleLineJsonString() { /// The following validations are performed: /// /// - For dimension, time, date time fields, support [DataType]: INT, LONG, FLOAT, DOUBLE, BIG_DECIMAL, - /// BOOLEAN, TIMESTAMP, STRING, JSON, UUID, BYTES + /// BOOLEAN, TIMESTAMP, STRING, JSON, UUID, BYTES. VARIANT is supported only for dimension fields. /// - For metric fields, support [DataType]: INT, LONG, FLOAT, DOUBLE, BIG_DECIMAL, BYTES /// /// Note: multi-value compatibility checks (e.g. rejecting MV JSON columns) live in diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java index b571c38a9bc1..3cb3ee746b42 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java @@ -1019,6 +1019,77 @@ public Object convert(Object value, PinotDataType sourceType) { } }, + /** + * Pinot's external representation of a VARIANT value: a validated PVAR envelope in a {@code byte[]}. + */ + VARIANT { + @Override + public int toInt(Object value) { + throw unsupportedConversion("INT"); + } + + @Override + public long toLong(Object value) { + throw unsupportedConversion("LONG"); + } + + @Override + public float toFloat(Object value) { + throw unsupportedConversion("FLOAT"); + } + + @Override + public double toDouble(Object value) { + throw unsupportedConversion("DOUBLE"); + } + + @Override + public BigDecimal toBigDecimal(Object value) { + throw unsupportedConversion("BIG_DECIMAL"); + } + + @Override + public boolean toBoolean(Object value) { + throw unsupportedConversion("BOOLEAN"); + } + + @Override + public Timestamp toTimestamp(Object value) { + throw unsupportedConversion("TIMESTAMP"); + } + + @Override + public String toString(Object value) { + throw unsupportedConversion("STRING"); + } + + @Override + public byte[] toBytes(Object value) { + byte[] envelope = (byte[]) value; + VariantEnvelope.validateAndGetMetadataLength(envelope); + return envelope; + } + + @Override + public UUID toUUID(Object value) { + throw unsupportedConversion("UUID"); + } + + @Override + public byte[] convert(Object value, PinotDataType sourceType) { + if (sourceType == VARIANT) { + return toBytes(value); + } + byte[] envelope = sourceType.toBytes(value); + VariantEnvelope.validateAndGetMetadataLength(envelope); + return envelope; + } + + private UnsupportedOperationException unsupportedConversion(String destinationType) { + return new UnsupportedOperationException("Cannot convert value from VARIANT to " + destinationType); + } + }, + /// Wraps [UUID]. Internal representation is the 16-byte big-endian binary form. /// /// When converting from UUID to other types: @@ -1839,6 +1910,7 @@ public Object convert(Object value, PinotDataType sourceType) { /// Converts to the internal representation of the value. /// - `BOOLEAN` → `Integer` (0/1) /// - `TIMESTAMP` → `Long` (epoch millis) + /// - `VARIANT` → `byte[]` (PVAR envelope) /// - `UUID` → `byte[]` (16-byte big-endian) /// - `PRIMITIVE_BOOLEAN_ARRAY` / `BOOLEAN_ARRAY` → `Integer[]` (per-element 0/1) /// - `TIMESTAMP_ARRAY` → `Long[]` (per-element epoch millis) @@ -2046,6 +2118,11 @@ public static PinotDataType getPinotDataTypeForIngestion(FieldSpec fieldSpec) { throw new IllegalStateException("There is no multi-value type for JSON"); case BYTES: return fieldSpec.isSingleValueField() ? BYTES : BYTES_ARRAY; + case VARIANT: + if (fieldSpec.isSingleValueField()) { + return VARIANT; + } + throw new IllegalStateException("There is no multi-value type for VARIANT"); case UUID: return fieldSpec.isSingleValueField() ? UUID : UUID_ARRAY; case MAP: diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java new file mode 100644 index 000000000000..ab9579f91b12 --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java @@ -0,0 +1,263 @@ +/** + * 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.utils; + +import java.nio.ByteBuffer; +import java.util.Objects; +import javax.annotation.Nullable; + + +/** + * Pinot-owned framing for the two buffers that make up a Parquet Variant value. + * + *

The version-1 wire format is: + * 

+ *   0        4 bytes  ASCII magic "PVAR"
+ *   4        1 byte   envelope version (1)
+ *   5        1 byte   flags (0)
+ *   6        2 bytes  reserved (0)
+ *   8        4 bytes  metadata length, unsigned range restricted to Java array sizes
+ *   12       4 bytes  value length, unsigned range restricted to Java array sizes
+ *   16       M bytes  Parquet Variant metadata
+ *   16 + M   V bytes  Parquet Variant value
+ *
+ * + *

An empty byte array is deliberately not an envelope. Pinot reserves it as the default null value for a + * {@code VARIANT} field, allowing the null-value vector to distinguish SQL null from an encoded Variant null. + * + *

This class validates only Pinot's stable outer framing. Producers and consumers remain responsible for validating + * the Parquet Variant metadata and value payloads. + */ +public final class VariantEnvelope { + public static final int HEADER_SIZE = 16; + public static final byte VERSION = 1; + public static final byte FLAGS = 0; + + private static final int MAGIC = 0x50564152; // ASCII "PVAR" + + private VariantEnvelope() { + } + + /** + * Encodes the remaining bytes of the supplied metadata and value buffers without changing their positions or + * limits. + * + *

Array-backed buffers are copied directly from their backing arrays. Other buffers, including direct and + * read-only buffers, are read through independent duplicate views. + */ + public static byte[] encode(ByteBuffer metadata, ByteBuffer value) { + Objects.requireNonNull(metadata, "metadata must not be null"); + Objects.requireNonNull(value, "value must not be null"); + + int metadataLength = metadata.remaining(); + int valueLength = value.remaining(); + byte[] envelope = allocate(metadataLength, valueLength); + copyRemaining(metadata, envelope, HEADER_SIZE); + copyRemaining(value, envelope, HEADER_SIZE + metadataLength); + return envelope; + } + + /** + * Encodes slices of the supplied arrays without allocating intermediate buffer views. + */ + public static byte[] encode(byte[] metadata, int metadataOffset, int metadataLength, byte[] value, int valueOffset, + int valueLength) { + Objects.requireNonNull(metadata, "metadata must not be null"); + Objects.requireNonNull(value, "value must not be null"); + requireRange(metadata, metadataOffset, metadataLength, "metadata"); + requireRange(value, valueOffset, valueLength, "value"); + + byte[] envelope = allocate(metadataLength, valueLength); + System.arraycopy(metadata, metadataOffset, envelope, HEADER_SIZE, metadataLength); + System.arraycopy(value, valueOffset, envelope, HEADER_SIZE + metadataLength, valueLength); + return envelope; + } + + /** + * Decodes and validates an envelope, returning zero-copy, read-only views over its metadata and value buffers. + * + *

The returned views alias {@code envelope}; this method does not copy either payload. The decoded object and + * any views obtained from it keep the backing array alive, so the caller does not need to retain a separate + * reference to {@code envelope}. Mutations made to the input array after this method returns are visible through + * the views and can corrupt the decoded payload. Callers must therefore treat the input array as immutable for as + * long as the decoded object or any returned view may be used. + * + *

The decoded holder is safe for concurrent reads when the aliased input array is not mutated. Each accessor + * returns a read-only view with independent position and limit, so cursor movement by one reader does not affect + * another reader. + */ + public static Decoded decode(byte[] envelope) { + int metadataLength = validateAndGetMetadataLength(envelope); + int valueLength = envelope.length - HEADER_SIZE - metadataLength; + ByteBuffer metadata = + ByteBuffer.wrap(envelope, HEADER_SIZE, metadataLength).slice().asReadOnlyBuffer(); + ByteBuffer value = + ByteBuffer.wrap(envelope, HEADER_SIZE + metadataLength, valueLength).slice().asReadOnlyBuffer(); + return new Decoded(metadata, value); + } + + /** + * Validates the stable outer framing and returns the metadata length without allocating buffer views. + * + *

The value begins at {@code HEADER_SIZE + metadataLength}; its length is the remaining envelope length. + */ + public static int validateAndGetMetadataLength(byte[] envelope) { + Objects.requireNonNull(envelope, "envelope must not be null"); + if (envelope.length < HEADER_SIZE) { + throw new IllegalArgumentException( + "Variant envelope is too short: " + envelope.length + " bytes (minimum " + HEADER_SIZE + ')'); + } + + if (readInt(envelope, 0) != MAGIC) { + throw new IllegalArgumentException("Invalid Variant envelope magic; expected PVAR"); + } + int version = Byte.toUnsignedInt(envelope[4]); + if (version != Byte.toUnsignedInt(VERSION)) { + throw new IllegalArgumentException("Unsupported Variant envelope version: " + version); + } + int flags = Byte.toUnsignedInt(envelope[5]); + if (flags != Byte.toUnsignedInt(FLAGS)) { + throw new IllegalArgumentException("Unsupported Variant envelope flags: " + flags); + } + int reserved = (Byte.toUnsignedInt(envelope[6]) << Byte.SIZE) | Byte.toUnsignedInt(envelope[7]); + if (reserved != 0) { + throw new IllegalArgumentException("Variant envelope reserved field must be zero: " + reserved); + } + + int metadataLength = readInt(envelope, 8); + int valueLength = readInt(envelope, 12); + if (metadataLength < 0 || valueLength < 0) { + throw new IllegalArgumentException( + "Variant envelope lengths exceed supported Java buffer size: metadata=" + Integer.toUnsignedLong( + metadataLength) + ", value=" + Integer.toUnsignedLong(valueLength)); + } + long expectedLength = (long) HEADER_SIZE + metadataLength + valueLength; + if (expectedLength != envelope.length) { + throw new IllegalArgumentException( + "Variant envelope length mismatch: header declares " + expectedLength + " bytes but found " + + envelope.length); + } + return metadataLength; + } + + /** + * Returns whether the bytes form a complete, supported Variant envelope. + */ + public static boolean isEnvelope(@Nullable byte[] envelope) { + if (envelope == null) { + return false; + } + try { + validateAndGetMetadataLength(envelope); + return true; + } catch (IllegalArgumentException e) { + return false; + } + } + + /** + * Allocates an initialized envelope with writable, zero-filled metadata and value regions. + * + *

This is the zero-intermediate-copy producer API for sources that can write directly into a destination array. + * The metadata region begins at {@link #HEADER_SIZE}; the value region begins at + * {@code HEADER_SIZE + metadataLength}. Callers must completely fill both regions before publishing the envelope. + */ + public static byte[] allocate(int metadataLength, int valueLength) { + if (metadataLength < 0 || valueLength < 0) { + throw new IllegalArgumentException( + "Variant envelope lengths must be non-negative: metadata=" + metadataLength + ", value=" + valueLength); + } + long envelopeLength = (long) HEADER_SIZE + metadataLength + valueLength; + if (envelopeLength > Integer.MAX_VALUE) { + throw new IllegalArgumentException("Variant envelope exceeds maximum Java byte array size: " + envelopeLength); + } + byte[] envelope = new byte[(int) envelopeLength]; + writeInt(envelope, 0, MAGIC); + envelope[4] = VERSION; + envelope[5] = FLAGS; + writeInt(envelope, 8, metadataLength); + writeInt(envelope, 12, valueLength); + return envelope; + } + + private static void copyRemaining(ByteBuffer source, byte[] target, int targetOffset) { + int length = source.remaining(); + if (length == 0) { + return; + } + if (source.hasArray()) { + System.arraycopy(source.array(), source.arrayOffset() + source.position(), target, targetOffset, length); + } else { + ByteBuffer sourceView = source.duplicate(); + sourceView.get(target, targetOffset, length); + } + } + + private static void requireRange(byte[] bytes, int offset, int length, String description) { + if (offset < 0 || length < 0 || (long) offset + length > bytes.length) { + throw new IllegalArgumentException( + "Variant " + description + " range exceeds source bounds: offset=" + offset + ", length=" + length + + ", sourceLength=" + bytes.length); + } + } + + private static int readInt(byte[] bytes, int offset) { + return Byte.toUnsignedInt(bytes[offset]) << 24 + | Byte.toUnsignedInt(bytes[offset + 1]) << 16 + | Byte.toUnsignedInt(bytes[offset + 2]) << 8 + | Byte.toUnsignedInt(bytes[offset + 3]); + } + + private static void writeInt(byte[] bytes, int offset, int value) { + bytes[offset] = (byte) (value >>> 24); + bytes[offset + 1] = (byte) (value >>> 16); + bytes[offset + 2] = (byte) (value >>> 8); + bytes[offset + 3] = (byte) value; + } + + /** + * Read-only views of the two Parquet Variant buffers stored in an envelope. + * + *

Instances retain and alias the envelope array supplied to {@link VariantEnvelope#decode(byte[])}. They are + * safe for concurrent reads only while that array remains unmodified. + */ + public static final class Decoded { + private final ByteBuffer _metadata; + private final ByteBuffer _value; + + private Decoded(ByteBuffer metadata, ByteBuffer value) { + _metadata = metadata; + _value = value; + } + + /** + * Returns a read-only metadata view with independent position and limit. + */ + public ByteBuffer getMetadata() { + return _metadata.asReadOnlyBuffer(); + } + + /** + * Returns a read-only value view with independent position and limit. + */ + public ByteBuffer getValue() { + return _value.asReadOnlyBuffer(); + } + } +} diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/data/VariantFieldSpecTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/data/VariantFieldSpecTest.java new file mode 100644 index 000000000000..4518229a8c58 --- /dev/null +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/data/VariantFieldSpecTest.java @@ -0,0 +1,115 @@ +/** + * 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.data; + +import java.nio.ByteBuffer; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.data.FieldSpec.FieldType; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.VariantEnvelope; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +public class VariantFieldSpecTest { + private static final byte[] METADATA = new byte[]{1, 2}; + private static final byte[] VALUE = new byte[]{3, 4, 5}; + + @Test + public void testStoredTypeAndDefaultNullValue() { + assertEquals(DataType.VARIANT.getStoredType(), DataType.BYTES); + assertFalse(DataType.VARIANT.isFixedWidth()); + assertFalse(DataType.VARIANT.supportsEquality()); + assertFalse(DataType.VARIANT.supportsHashing()); + assertFalse(DataType.VARIANT.supportsOrdering()); + assertFalse(DataType.VARIANT.supportsMinMax()); + assertFalse(DataType.VARIANT.supportsDirectAggregation()); + assertFalse(DataType.VARIANT.supportsPatternMatching()); + assertTrue(DataType.BYTES.supportsEquality()); + assertTrue(DataType.BYTES.supportsHashing()); + assertTrue(DataType.BYTES.supportsOrdering()); + assertTrue(DataType.BYTES.supportsMinMax()); + assertTrue(DataType.BYTES.supportsDirectAggregation()); + assertTrue(DataType.BYTES.supportsPatternMatching()); + assertFalse(DataType.MAP.supportsOrdering()); + assertFalse(DataType.MAP.supportsMinMax()); + + DimensionFieldSpec fieldSpec = new DimensionFieldSpec("payload", DataType.VARIANT, true); + assertEquals((byte[]) fieldSpec.getDefaultNullValue(), new byte[0]); + assertEquals(fieldSpec.getDefaultNullValueString(), ""); + } + + @Test + public void testExternalAndInternalConversionValidateEnvelope() { + byte[] envelope = + VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE)); + String envelopeHex = BytesUtils.toHexString(envelope); + + assertEquals((byte[]) DataType.VARIANT.convert(envelopeHex), envelope); + assertEquals(DataType.VARIANT.convertInternal(envelopeHex), new ByteArray(envelope)); + assertEquals(DataType.VARIANT.toString(envelope), envelopeHex); + + assertThrows(IllegalArgumentException.class, () -> DataType.VARIANT.convert("")); + assertThrows(IllegalArgumentException.class, () -> DataType.VARIANT.convert("00")); + } + + @Test + public void testSemanticByteOperationsAreUnsupported() { + byte[] envelope = + VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE)); + + assertThrows(UnsupportedOperationException.class, () -> DataType.VARIANT.equals(envelope, envelope)); + assertThrows(UnsupportedOperationException.class, () -> DataType.VARIANT.hashCode(envelope)); + assertThrows(UnsupportedOperationException.class, () -> DataType.VARIANT.compare(envelope, envelope)); + } + + @Test + public void testSchemaStructuralEqualityAndHashCode() + throws Exception { + DimensionFieldSpec first = new DimensionFieldSpec("payload", DataType.VARIANT, true); + DimensionFieldSpec equivalent = new DimensionFieldSpec("payload", DataType.VARIANT, true); + DimensionFieldSpec different = new DimensionFieldSpec("otherPayload", DataType.VARIANT, true); + + assertEquals(first, equivalent); + assertEquals(first.hashCode(), equivalent.hashCode()); + assertNotEquals(first, different); + + Schema firstSchema = new Schema.SchemaBuilder().setSchemaName("events").addField(first).build(); + Schema equivalentSchema = + new Schema.SchemaBuilder().setSchemaName("events").addField(equivalent).build(); + assertEquals(firstSchema, equivalentSchema); + assertEquals(firstSchema.hashCode(), equivalentSchema.hashCode()); + assertEquals(Schema.fromString(firstSchema.toString()), firstSchema); + } + + @Test + public void testVariantIsDimensionOnly() { + Schema.validate(FieldType.DIMENSION, DataType.VARIANT); + assertThrows(IllegalStateException.class, () -> Schema.validate(FieldType.METRIC, DataType.VARIANT)); + assertThrows(IllegalStateException.class, () -> Schema.validate(FieldType.TIME, DataType.VARIANT)); + assertThrows(IllegalStateException.class, () -> Schema.validate(FieldType.DATE_TIME, DataType.VARIANT)); + assertThrows(IllegalStateException.class, () -> Schema.validate(FieldType.COMPLEX, DataType.VARIANT)); + } +} diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/PinotDataTypeTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/PinotDataTypeTest.java index 9d8b65697cae..620ae9a394b0 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/PinotDataTypeTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/PinotDataTypeTest.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import java.math.BigDecimal; +import java.nio.ByteBuffer; import java.sql.Timestamp; import java.time.Instant; import java.time.LocalDate; @@ -221,6 +222,20 @@ public void testBytes() { assertEquals(BYTES.convert(new String[]{"0001"}, STRING_ARRAY), new byte[]{0, 1}); } + @Test + public void testVariant() { + byte[] envelope = + VariantEnvelope.encode(ByteBuffer.wrap(new byte[]{1, 2}), ByteBuffer.wrap(new byte[]{3, 4})); + + assertEquals(VARIANT.convert(envelope, BYTES), envelope); + assertEquals(VARIANT.toInternal(envelope), envelope); + assertEquals(BYTES.convert(envelope, VARIANT), envelope); + assertThrows(IllegalArgumentException.class, () -> VARIANT.convert(new byte[0], BYTES)); + assertThrows(IllegalArgumentException.class, () -> VARIANT.convert(new byte[]{1, 2, 3}, BYTES)); + assertThrows(UnsupportedOperationException.class, () -> VARIANT.toString(envelope)); + assertThrows(UnsupportedOperationException.class, () -> VARIANT.toInt(envelope)); + } + @Test public void testUUID() { java.util.UUID uuid = java.util.UUID.fromString(UUID_VALUE); diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/VariantEnvelopeTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/VariantEnvelopeTest.java new file mode 100644 index 000000000000..690ad44da638 --- /dev/null +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/VariantEnvelopeTest.java @@ -0,0 +1,263 @@ +/** + * 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.utils; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.ReadOnlyBufferException; +import java.util.Arrays; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + + +public class VariantEnvelopeTest { + private static final byte[] METADATA = new byte[]{1, 2, 3}; + private static final byte[] VALUE = new byte[]{0, 10, 20, 30}; + private static final byte[] VERSION_ONE_FROZEN = new byte[]{ + 'P', 'V', 'A', 'R', 1, 0, 0, 0, + 0, 0, 0, 3, 0, 0, 0, 4, + 1, 2, 3, 0, 10, 20, 30 + }; + + @Test + public void testRoundTripUsesRemainingBytesWithoutMutatingInputs() { + ByteBuffer metadata = ByteBuffer.wrap(new byte[]{99, 1, 2, 3, 98}); + metadata.position(1); + metadata.limit(4); + ByteBuffer value = ByteBuffer.wrap(new byte[]{97, 0, 10, 20, 30, 96}); + value.position(1); + value.limit(5); + + byte[] envelope = VariantEnvelope.encode(metadata, value); + + assertEquals(metadata.position(), 1); + assertEquals(value.position(), 1); + assertEquals(envelope.length, VariantEnvelope.HEADER_SIZE + METADATA.length + VALUE.length); + assertEquals(Arrays.copyOfRange(envelope, 0, 4), new byte[]{'P', 'V', 'A', 'R'}); + assertEquals(Byte.toUnsignedInt(envelope[4]), Byte.toUnsignedInt(VariantEnvelope.VERSION)); + assertEquals(envelope[5], VariantEnvelope.FLAGS); + assertEquals(ByteBuffer.wrap(envelope).order(ByteOrder.BIG_ENDIAN).getInt(8), METADATA.length); + assertEquals(ByteBuffer.wrap(envelope).order(ByteOrder.BIG_ENDIAN).getInt(12), VALUE.length); + + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + assertEquals(readBytes(decoded.getMetadata()), METADATA); + assertEquals(readBytes(decoded.getValue()), VALUE); + assertEquals(VariantEnvelope.validateAndGetMetadataLength(envelope), METADATA.length); + assertTrue(VariantEnvelope.isEnvelope(envelope)); + } + + @Test + public void testRoundTripFromArraySlices() { + byte[] metadata = new byte[]{99, 1, 2, 3, 98}; + byte[] value = new byte[]{97, 0, 10, 20, 30, 96}; + + byte[] envelope = VariantEnvelope.encode(metadata, 1, 3, value, 1, 4); + + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + assertEquals(readBytes(decoded.getMetadata()), METADATA); + assertEquals(readBytes(decoded.getValue()), VALUE); + assertEquals(VariantEnvelope.validateAndGetMetadataLength(envelope), METADATA.length); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.encode(metadata, -1, 3, value, 1, 4)); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.encode(metadata, 1, 5, value, 1, 4)); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.encode(metadata, 1, 3, value, 1, 6)); + } + + @Test + public void testArrayBackedBuffersHonorArrayOffsetPositionAndLimit() { + ByteBuffer metadata = ByteBuffer.wrap(new byte[]{88, 99, 1, 2, 3, 98, 87}); + metadata.position(1); + metadata.limit(6); + metadata = metadata.slice(); + metadata.position(1); + metadata.limit(4); + + ByteBuffer value = ByteBuffer.wrap(new byte[]{86, 97, 0, 10, 20, 30, 96, 85}); + value.position(1); + value.limit(7); + value = value.slice(); + value.position(1); + value.limit(5); + + assertTrue(metadata.hasArray()); + assertTrue(value.hasArray()); + assertEquals(metadata.arrayOffset(), 1); + assertEquals(value.arrayOffset(), 1); + assertFrozenEncodingPreservesState(metadata, value); + } + + @Test + public void testDirectBuffersUseFallbackWithoutChangingState() { + ByteBuffer metadata = ByteBuffer.allocateDirect(5); + metadata.put(new byte[]{99, 1, 2, 3, 98}); + metadata.position(1); + metadata.limit(4); + ByteBuffer value = ByteBuffer.allocateDirect(6); + value.put(new byte[]{97, 0, 10, 20, 30, 96}); + value.position(1); + value.limit(5); + + assertFalse(metadata.hasArray()); + assertFalse(value.hasArray()); + assertFrozenEncodingPreservesState(metadata, value); + } + + @Test + public void testReadOnlyBuffersUseFallbackWithoutChangingState() { + ByteBuffer metadata = ByteBuffer.wrap(new byte[]{99, 1, 2, 3, 98}).asReadOnlyBuffer(); + metadata.position(1); + metadata.limit(4); + ByteBuffer value = ByteBuffer.wrap(new byte[]{97, 0, 10, 20, 30, 96}).asReadOnlyBuffer(); + value.position(1); + value.limit(5); + + assertFalse(metadata.hasArray()); + assertFalse(value.hasArray()); + assertFrozenEncodingPreservesState(metadata, value); + } + + @Test + public void testVersionOneFrozenBytes() { + assertEquals(VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE)), VERSION_ONE_FROZEN); + assertEquals(VariantEnvelope.encode(METADATA, 0, METADATA.length, VALUE, 0, VALUE.length), VERSION_ONE_FROZEN); + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(VERSION_ONE_FROZEN); + assertEquals(readBytes(decoded.getMetadata()), METADATA); + assertEquals(readBytes(decoded.getValue()), VALUE); + } + + @Test + public void testDirectProducerAllocation() { + byte[] envelope = VariantEnvelope.allocate(METADATA.length, VALUE.length); + System.arraycopy(METADATA, 0, envelope, VariantEnvelope.HEADER_SIZE, METADATA.length); + System.arraycopy(VALUE, 0, envelope, VariantEnvelope.HEADER_SIZE + METADATA.length, VALUE.length); + + assertEquals(envelope, VERSION_ONE_FROZEN); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.allocate(-1, 0)); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.allocate(0, -1)); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.allocate(Integer.MAX_VALUE, 0)); + } + + @Test + public void testDecodedBuffersAreReadOnlyAndHaveIndependentPositions() { + VariantEnvelope.Decoded decoded = + VariantEnvelope.decode(VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE))); + + ByteBuffer firstMetadata = decoded.getMetadata(); + assertThrows(ReadOnlyBufferException.class, () -> firstMetadata.put((byte) 0)); + firstMetadata.get(); + assertEquals(decoded.getMetadata().position(), 0); + + ByteBuffer firstValue = decoded.getValue(); + assertThrows(ReadOnlyBufferException.class, () -> firstValue.put((byte) 0)); + firstValue.get(); + assertEquals(decoded.getValue().position(), 0); + } + + @Test + public void testDecodedBuffersAliasInputEnvelope() { + byte[] envelope = VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE)); + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + + envelope[VariantEnvelope.HEADER_SIZE] = 42; + envelope[VariantEnvelope.HEADER_SIZE + METADATA.length] = 43; + + assertEquals(decoded.getMetadata().get(0), (byte) 42); + assertEquals(decoded.getValue().get(0), (byte) 43); + } + + @Test + public void testEmptyPayloadBuffersStillProduceEnvelope() { + byte[] envelope = VariantEnvelope.encode(ByteBuffer.allocate(0), ByteBuffer.allocate(0)); + assertEquals(envelope.length, VariantEnvelope.HEADER_SIZE); + assertTrue(VariantEnvelope.isEnvelope(envelope)); + assertEquals(VariantEnvelope.decode(envelope).getMetadata().remaining(), 0); + assertEquals(VariantEnvelope.decode(envelope).getValue().remaining(), 0); + + assertFalse(VariantEnvelope.isEnvelope(new byte[0])); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.decode(new byte[0])); + } + + @Test + public void testRejectsInvalidHeaderFields() { + byte[] valid = VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE)); + + byte[] invalidMagic = valid.clone(); + invalidMagic[0] = 'X'; + assertInvalid(invalidMagic, "magic"); + + byte[] invalidVersion = valid.clone(); + invalidVersion[4] = 2; + assertInvalid(invalidVersion, "version"); + + byte[] invalidFlags = valid.clone(); + invalidFlags[5] = 1; + assertInvalid(invalidFlags, "flags"); + + byte[] invalidReserved = valid.clone(); + invalidReserved[7] = 1; + assertInvalid(invalidReserved, "reserved"); + } + + @Test + public void testRejectsInvalidLengths() { + byte[] valid = VariantEnvelope.encode(ByteBuffer.wrap(METADATA), ByteBuffer.wrap(VALUE)); + + byte[] negativeMetadataLength = valid.clone(); + ByteBuffer.wrap(negativeMetadataLength).order(ByteOrder.BIG_ENDIAN).putInt(8, -1); + assertInvalid(negativeMetadataLength, "lengths"); + + byte[] negativeValueLength = valid.clone(); + ByteBuffer.wrap(negativeValueLength).order(ByteOrder.BIG_ENDIAN).putInt(12, -1); + assertInvalid(negativeValueLength, "lengths"); + + assertInvalid(Arrays.copyOf(valid, valid.length - 1), "length mismatch"); + assertInvalid(Arrays.copyOf(valid, valid.length + 1), "length mismatch"); + } + + private static void assertInvalid(byte[] envelope, String messageFragment) { + IllegalArgumentException exception = + expectThrows(IllegalArgumentException.class, () -> VariantEnvelope.decode(envelope)); + assertTrue(exception.getMessage().contains(messageFragment), exception.getMessage()); + assertThrows(IllegalArgumentException.class, () -> VariantEnvelope.validateAndGetMetadataLength(envelope)); + assertFalse(VariantEnvelope.isEnvelope(envelope)); + } + + private static void assertFrozenEncodingPreservesState(ByteBuffer metadata, ByteBuffer value) { + int metadataPosition = metadata.position(); + int metadataLimit = metadata.limit(); + int valuePosition = value.position(); + int valueLimit = value.limit(); + + assertEquals(VariantEnvelope.encode(metadata, value), VERSION_ONE_FROZEN); + assertEquals(metadata.position(), metadataPosition); + assertEquals(metadata.limit(), metadataLimit); + assertEquals(value.position(), valuePosition); + assertEquals(value.limit(), valueLimit); + } + + private static byte[] readBytes(ByteBuffer buffer) { + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + return bytes; + } +} diff --git a/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/compile/DataTypeMapper.java b/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/compile/DataTypeMapper.java index 4fbdf4ae79e2..76b9c4edd16f 100644 --- a/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/compile/DataTypeMapper.java +++ b/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/compile/DataTypeMapper.java @@ -51,6 +51,7 @@ public final class DataTypeMapper { map.put("BINARY", DataType.BYTES); map.put("BYTES", DataType.BYTES); map.put("JSON", DataType.JSON); + map.put("VARIANT", DataType.VARIANT); NAME_TO_DATATYPE = map; } diff --git a/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SchemaEmitter.java b/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SchemaEmitter.java index 24a6e1e436e9..10cf10426452 100644 --- a/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SchemaEmitter.java +++ b/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SchemaEmitter.java @@ -21,6 +21,7 @@ import java.math.BigDecimal; import java.time.Instant; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -134,6 +135,9 @@ private static String emitColumn(FieldSpec spec) { boolean atNaturalDefault; if (defaultValue == null || naturalDefault == null) { atNaturalDefault = (defaultValue == null && naturalDefault == null); + } else if (spec.getDataType() == DataType.VARIANT) { + // This compares schema metadata (the reserved SQL-null sentinel), not Variant values. + atNaturalDefault = Arrays.equals((byte[]) defaultValue, (byte[]) naturalDefault); } else if (spec.getDataType() == DataType.BIG_DECIMAL && defaultValue instanceof BigDecimal && naturalDefault instanceof BigDecimal) { atNaturalDefault = ((BigDecimal) defaultValue).compareTo((BigDecimal) naturalDefault) == 0; @@ -185,6 +189,8 @@ private static String emitDataType(DataType dt) { return "JSON"; case BYTES: return "BYTES"; + case VARIANT: + return "VARIANT"; default: // Fall back to the enum name for types we don't have a canonical short name for // (LIST, MAP, STRUCT, UNKNOWN). These are not yet expressible in DDL but emitting the diff --git a/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SqlIdentifiers.java b/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SqlIdentifiers.java index e8cacf01062b..2695445e8b44 100644 --- a/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SqlIdentifiers.java +++ b/pinot-sql-ddl/src/main/java/org/apache/pinot/sql/ddl/reverse/SqlIdentifiers.java @@ -51,7 +51,7 @@ final class SqlIdentifiers { // rather than an identifier and the column declaration would fail. "INT", "INTEGER", "SMALLINT", "TINYINT", "BIGINT", "LONG", "FLOAT", "REAL", "DOUBLE", "DECIMAL", "NUMERIC", "BIG_DECIMAL", "BOOLEAN", "TIMESTAMP", "VARCHAR", "CHAR", "STRING", - "VARBINARY", "BINARY", "BYTES", "JSON"); + "VARBINARY", "BINARY", "BYTES", "JSON", "VARIANT"); private static final Pattern BARE_IDENTIFIER = Pattern.compile("[A-Za-z_][A-Za-z0-9_]*"); diff --git a/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/compile/DdlCompilerTest.java b/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/compile/DdlCompilerTest.java index 77721726b930..6d999985915e 100644 --- a/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/compile/DdlCompilerTest.java +++ b/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/compile/DdlCompilerTest.java @@ -57,6 +57,17 @@ public void offlineMinimal() { assertEquals(c.getTableConfig().getValidationConfig().getReplication(), "1"); } + @Test + public void variantColumnMapsToSingleValueDimension() { + CompiledCreateTable compiled = + compileCreate("CREATE TABLE events (payload VARIANT) TABLE_TYPE = OFFLINE"); + + FieldSpec payload = compiled.getSchema().getFieldSpecFor("payload"); + assertTrue(payload instanceof DimensionFieldSpec); + assertEquals(payload.getDataType(), DataType.VARIANT); + assertTrue(payload.isSingleValueField()); + } + @Test public void databaseQualifiedNamePreservedInTableConfig() { CompiledCreateTable c = compileCreate( diff --git a/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/reverse/CanonicalDdlEmitterTest.java b/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/reverse/CanonicalDdlEmitterTest.java index 83f07057f176..19e8ce2a291b 100644 --- a/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/reverse/CanonicalDdlEmitterTest.java +++ b/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/reverse/CanonicalDdlEmitterTest.java @@ -170,6 +170,20 @@ public void noDefaultEmittedForBytesAtNaturalDefault() { "BYTES column at natural default must not emit a DEFAULT clause; got:\n" + emitted); } + @Test + public void variantTypeEmitsWithoutReservedNullSentinelDefault() { + Schema schema = new Schema.SchemaBuilder() + .setSchemaName("events") + .addSingleValueDimension("payload", DataType.VARIANT) + .build(); + TableConfig config = new TableConfigBuilder(TableType.OFFLINE).setTableName("events").build(); + + String emitted = CanonicalDdlEmitter.emit(schema, config); + assertTrue(emitted.contains("payload VARIANT DIMENSION"), emitted); + assertFalse(emitted.contains("DEFAULT"), + "VARIANT's reserved SQL-null sentinel must not be emitted as a user default:\n" + emitted); + } + /// Regression: BigDecimal.toString() can emit scientific notation (1E+30) which Calcite's /// Literal() rule does not accept. The fix routes BIG_DECIMAL defaults through toPlainString(). @Test diff --git a/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/roundtrip/RoundTripTest.java b/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/roundtrip/RoundTripTest.java index cfa240a16674..b7316d468163 100644 --- a/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/roundtrip/RoundTripTest.java +++ b/pinot-sql-ddl/src/test/java/org/apache/pinot/sql/ddl/roundtrip/RoundTripTest.java @@ -84,6 +84,18 @@ public void minimalOfflineTable() { assertRoundTrip(schema, config); } + @Test + public void variantOfflineTable() { + Schema schema = new Schema.SchemaBuilder() + .setSchemaName("events") + .addSingleValueDimension("payload", DataType.VARIANT) + .build(); + TableConfig config = new TableConfigBuilder(TableType.OFFLINE) + .setTableName("events") + .build(); + assertRoundTrip(schema, config); + } + @Test public void offlineTableWithRetentionAndTenants() { Schema schema = new Schema.SchemaBuilder() diff --git a/pinot-tools/pom.xml b/pinot-tools/pom.xml index 8959e1dd7d30..1797ca13d1b8 100644 --- a/pinot-tools/pom.xml +++ b/pinot-tools/pom.xml @@ -323,6 +323,16 @@ + + org.apache.pinot.tools.VariantQuickStart + quick-start-variant-batch + + 4G + + -Dlog4j2.configurationFile=conf/quickstart-log4j2.xml + + + org.apache.pinot.tools.AuthQuickstart quick-start-auth diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java new file mode 100644 index 000000000000..5b60836fdd76 --- /dev/null +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java @@ -0,0 +1,115 @@ +/** + * 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.tools; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.TimeUnit; +import org.apache.pinot.tools.admin.PinotAdministrator; +import org.apache.pinot.tools.admin.command.QuickstartRunner; + + +/** + * Batch quickstart for ingesting and querying an Apache Parquet VARIANT column. + */ +public class VariantQuickStart extends Quickstart { + private static final String[] VARIANT_TABLE_DIRECTORIES = {"examples/batch/variantEvents"}; + private static final int EXPECTED_NUM_ROWS = 5; + private static final long BOOTSTRAP_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(5); + private static final long BOOTSTRAP_POLL_INTERVAL_MS = TimeUnit.SECONDS.toMillis(1); + + @Override + public List types() { + return List.of("VARIANT", "OFFLINE_VARIANT", "OFFLINE-VARIANT", "BATCH_VARIANT", "BATCH-VARIANT"); + } + + @Override + protected String[] getDefaultBatchTableDirectories() { + return VARIANT_TABLE_DIRECTORIES; + } + + @Override + protected void waitForBootstrapToComplete(QuickstartRunner runner) + throws Exception { + printStatus(Color.CYAN, + "***** Waiting up to 5 minutes for the Parquet VARIANT segment to become queryable *****"); + long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(BOOTSTRAP_TIMEOUT_MS); + JsonNode lastResponse = null; + Exception lastException = null; + while (System.nanoTime() < deadline) { + try { + lastResponse = runner.runQuery("SELECT COUNT(*) FROM variantEvents"); + JsonNode rows = lastResponse.path("resultTable").path("rows"); + if (lastResponse.path("exceptions").isEmpty() && rows.isArray() && rows.size() == 1 + && rows.get(0).isArray() && rows.get(0).size() == 1 + && rows.get(0).get(0).asInt(-1) == EXPECTED_NUM_ROWS) { + printStatus(Color.GREEN, + "***** Parquet VARIANT segment is ready with " + EXPECTED_NUM_ROWS + " rows *****"); + return; + } + } catch (Exception e) { + lastException = e; + } + Thread.sleep(BOOTSTRAP_POLL_INTERVAL_MS); + } + + String detail = lastResponse != null ? lastResponse.toString() + : lastException != null ? lastException.getMessage() : "no query response"; + throw new IllegalStateException( + "Timed out waiting for variantEvents to contain exactly " + EXPECTED_NUM_ROWS + " rows; last result: " + + detail, lastException); + } + + @Override + public void runSampleQueries(QuickstartRunner runner) + throws Exception { + printStatus(Color.YELLOW, "***** Parquet VARIANT quickstart setup complete *****"); + + runQuery(runner, "Materialized hot field", + "SELECT eventType, COUNT(*) FROM variantEvents GROUP BY eventType ORDER BY eventType"); + runQuery(runner, "Extract nested values directly from VARIANT", + "SELECT eventId, variant_get(payload, '$.user.id', 'STRING') AS userId, " + + "variant_get(payload, '$.amount', 'DOUBLE') AS amount " + + "FROM variantEvents WHERE eventType = 'checkout' ORDER BY eventId"); + runQuery(runner, "Render the retained Parquet VARIANT value as JSON", + "SELECT eventId, variantToJson(payload) FROM variantEvents ORDER BY eventId LIMIT 5"); + } + + private static void runQuery(QuickstartRunner runner, String description, String query) + throws Exception { + printStatus(Color.YELLOW, description); + printStatus(Color.CYAN, "Query : " + query); + JsonNode response = runner.runQuery("SET enableNullHandling=true; " + query); + if (!response.path("exceptions").isEmpty()) { + throw new IllegalStateException("VARIANT quickstart query failed: " + response.path("exceptions")); + } + printStatus(Color.YELLOW, prettyPrintResponse(response)); + printStatus(Color.GREEN, "***************************************************"); + } + + public static void main(String[] args) + throws Exception { + List arguments = new ArrayList<>(); + arguments.addAll(Arrays.asList("QuickStart", "-type", "VARIANT")); + arguments.addAll(Arrays.asList(args)); + PinotAdministrator.main(arguments.toArray(new String[0])); + } +} diff --git a/pinot-tools/src/main/resources/examples/batch/variantEvents/README.md b/pinot-tools/src/main/resources/examples/batch/variantEvents/README.md new file mode 100644 index 000000000000..cb7e75d54847 --- /dev/null +++ b/pinot-tools/src/main/resources/examples/batch/variantEvents/README.md @@ -0,0 +1,88 @@ + + +# Parquet VARIANT quickstart table + +`variantEvents_data.parquet` stores `payload` using the Apache Parquet +`VARIANT(1)` logical type. This sample verifies Pinot's current reader for a +top-level, non-repeated Variant column. Pinot retains the encoded value in a +raw `VARIANT` column and materializes `$.eventType` into an indexed STRING +column while ingesting the file. The ingestion spec explicitly selects +`ParquetNativeRecordReader`; use the native reader for VARIANT files that also +contain Avro schema metadata so existing Parquet reader selection remains +backward compatible. + +The persisted format, query semantics, compatibility rules, and rollout gate are +documented in the [VARIANT design document](../../../../../../../pinot-spi/VARIANT_DESIGN.md). + +Build and start the dedicated quickstart: + +```shell +./mvnw clean install -DskipTests -Pbin-dist -Pbuild-shaded-jar +build/bin/quick-start-variant-batch.sh +``` + +The quickstart creates the schema and offline table, runs the standalone +Parquet ingestion job, uploads the segment, and executes these representative +queries: + +```sql +SELECT eventType, COUNT(*) +FROM variantEvents +GROUP BY eventType +ORDER BY eventType; + +SELECT + eventId, + variant_get(payload, '$.user.id', 'STRING') AS userId, + variant_get(payload, '$.amount', 'DOUBLE') AS amount +FROM variantEvents +WHERE eventType = 'checkout' +ORDER BY eventId; + +SELECT eventId, variantToJson(payload) +FROM variantEvents +ORDER BY eventId; +``` + +The sample schema and table both enable storage null handling. VARIANT tables +must enable either schema column-based null handling or table-level null +handling so that SQL null, Variant null, and missing paths stay distinct. +All sample queries also use `SET enableNullHandling=true`; Pinot rejects +VARIANT SQL functions without that query option instead of returning ambiguous +null results. + +## Production rollout + +`VARIANT` introduces schema and query-wire type names that older Pinot +processes do not recognize. Upgrade every controller, broker, server, minion, +and external ingestion job before registering a schema containing `VARIANT`. +Keep VARIANT columns out of queries until every broker and server has been +upgraded, and do not roll back to a pre-VARIANT build while a VARIANT table is +active. Existing tables and queries are unaffected during the rolling upgrade. + +This follows the same user workflow and the supported subset of the +[Spark VARIANT function contract](https://spark.apache.org/docs/latest/api/sql/variant-functions/): +the source file carries a typed Variant value, frequently filtered paths can +be materialized during ingestion, and other paths remain available for +late-bound query-time extraction. The committed fixture was written with +parquet-java. A Spark-produced interoperability golden, nested or repeated +Variant columns, quoted path keys, streaming ingestion, and rolling +mixed-version queries over VARIANT columns are outside this sample's verified +scope. diff --git a/pinot-tools/src/main/resources/examples/batch/variantEvents/ingestionJobSpec.yaml b/pinot-tools/src/main/resources/examples/batch/variantEvents/ingestionJobSpec.yaml new file mode 100644 index 000000000000..17bad7689112 --- /dev/null +++ b/pinot-tools/src/main/resources/examples/batch/variantEvents/ingestionJobSpec.yaml @@ -0,0 +1,51 @@ +# +# 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. +# + +executionFrameworkSpec: + name: 'standalone' + segmentGenerationJobRunnerClassName: 'org.apache.pinot.plugin.ingestion.batch.standalone.SegmentGenerationJobRunner' + segmentTarPushJobRunnerClassName: 'org.apache.pinot.plugin.ingestion.batch.standalone.SegmentTarPushJobRunner' + segmentUriPushJobRunnerClassName: 'org.apache.pinot.plugin.ingestion.batch.standalone.SegmentUriPushJobRunner' + segmentMetadataPushJobRunnerClassName: 'org.apache.pinot.plugin.ingestion.batch.standalone.SegmentMetadataPushJobRunner' + +jobType: SegmentCreationAndTarPush +inputDirURI: 'examples/batch/variantEvents/rawdata' +includeFileNamePattern: 'glob:**/*.parquet' +outputDirURI: 'examples/batch/variantEvents/segments' +overwriteOutput: true + +pinotFSSpecs: + - scheme: file + className: org.apache.pinot.spi.filesystem.LocalPinotFS + +recordReaderSpec: + dataFormat: 'parquet' + className: 'org.apache.pinot.plugin.inputformat.parquet.ParquetNativeRecordReader' + +tableSpec: + tableName: 'variantEvents' + schemaURI: 'http://localhost:9000/tables/variantEvents/schema' + tableConfigURI: 'http://localhost:9000/tables/variantEvents' + +pinotClusterSpecs: + - controllerURI: 'http://localhost:9000' + +pushJobSpec: + pushAttempts: 2 + pushRetryIntervalMillis: 1000 diff --git a/pinot-tools/src/main/resources/examples/batch/variantEvents/rawdata/variantEvents_data.parquet b/pinot-tools/src/main/resources/examples/batch/variantEvents/rawdata/variantEvents_data.parquet new file mode 100644 index 0000000000000000000000000000000000000000..0d901fde2b28ca8d5f6dc7c235d68301b2219a62 GIT binary patch literal 1442 zcmbVMO=uHQ5T1RzO*Uz4ZS_5NZP2JRB+!;5+a#?7sX|4ig;Mn<$Ts_m!TgzQQVUWG z9z6I%LGk3FNWDl;r3Vilqz4ZbL3&o{(TgH@@#wtWX150qRo)}>X5Y+wGxJT}%(>Y# z1zN@Rargb#NR)Ow>yt}3N%0!fK!-?)CHkKz4RDw^QwJyvCdmf-Y8jS zrPiw2C7V~8NnsUmtOjC#H7nLy^;*?2mW43@W_uLZ&!->0cZ+Nl1jO12-paj1ekKfd ziP$^8c8Ob$VFLQ?A5VgS17viZ z4J;Et?<2Yv+~dzsdGH`q9z$2L} za(qiEAol;yk|u`OZ}SFm=L*<*H;u7(AsjAX+d&qNt0;~NW36wYfiFg5q*AkZd8BGq zI8t`UtklaKw99p~al6Is;YuUc=8a~lR*j`c()p2eEHT$Am91EQ%p9|fjFIAVqxsyp mH8*ZpJUgG6&+zn^ku!>^(W1%6(n(#+pSNKX-=Gparquet-hadoop ${parquet.version} + + org.apache.parquet + parquet-variant + ${parquet.version} + org.apache.orc orc-core From 630ad9298cd580ab6a69e2eb5d552e82b68946bc Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 28 Jul 2026 15:34:31 -0700 Subject: [PATCH 2/8] Preserve SQL null comparison semantics Allow SQL NULL literals to participate in comparison validation without being classified as raw VARIANT values. Keep equality and DISTINCT FROM behavior aligned in the single-stage transform layer and the multi-stage filter operand, with regression coverage for both engines. --- .../BinaryOperatorTransformFunction.java | 7 +++++- .../BinaryOperatorTransformFunctionTest.java | 20 +++++++++++++++++ .../DistinctFromTransformFunctionTest.java | 12 ++++++++++ .../operator/operands/FilterOperand.java | 5 +++-- .../operator/operands/FilterOperandTest.java | 22 +++++++++++++++++++ 5 files changed, 63 insertions(+), 3 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java index a7fed5b82b89..ad2809ed8728 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java @@ -108,8 +108,13 @@ public void init(List arguments, Map c _rightTransformFunction = arguments.get(1); DataType leftDataType = _leftTransformFunction.getResultMetadata().getDataType(); DataType rightDataType = _rightTransformFunction.getResultMetadata().getDataType(); - Preconditions.checkArgument(leftDataType.supportsOrdering() && rightDataType.supportsOrdering(), + Preconditions.checkArgument((leftDataType == DataType.UNKNOWN || leftDataType.supportsOrdering()) + && (rightDataType == DataType.UNKNOWN || rightDataType.supportsOrdering()), "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); + if (leftDataType == DataType.UNKNOWN || rightDataType == DataType.UNKNOWN) { + _alwaysNull = true; + return; + } _leftStoredType = leftDataType.getStoredType(); _rightStoredType = rightDataType.getStoredType(); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunctionTest.java index 8ddfa1d6db11..339f610040f2 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunctionTest.java @@ -31,6 +31,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; /// BinaryOperatorTransformFunctionTest abstracts common test methods for EqualsTransformFunctionTest, @@ -286,6 +287,25 @@ public void testBinaryOperatorTransformFunctionUUIDNoDict() { testTransformFunction(transformFunction, expectedValues); } + @Test + public void testRawVariantComparisonRejected() { + ExpressionContext expression = RequestContextUtils.getExpression( + String.format("%s(parseJson('{\"value\":1}'), 1)", getFunctionName())); + BadQueryRequestException exception = expectThrows(BadQueryRequestException.class, + () -> TransformFunctionFactory.getNullHandlingEnabled(expression, _dataSourceMap)); + assertTrue(exception.getMessage().contains("Raw VARIANT values do not support comparison")); + } + + @Test + public void testLeftNullRightLiteral() { + ExpressionContext expression = + RequestContextUtils.getExpression(String.format("%s(null, 1)", getFunctionName())); + TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap); + RoaringBitmap bitmap = new RoaringBitmap(); + bitmap.add(0L, NUM_ROWS); + testTransformFunctionWithNull(transformFunction, new boolean[NUM_ROWS], bitmap); + } + @Test(dataProvider = "testIllegalArguments", expectedExceptions = {BadQueryRequestException.class}) public void testIllegalArguments(String expressionStr) { ExpressionContext expression = RequestContextUtils.getExpression(expressionStr); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/DistinctFromTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/DistinctFromTransformFunctionTest.java index 6536fd5f8b50..dcd4f4cba978 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/DistinctFromTransformFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/DistinctFromTransformFunctionTest.java @@ -20,6 +20,7 @@ import java.io.File; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -260,6 +261,17 @@ public void testDistinctFromLeftLiteralRightIdentifier() testTransformFunction(expression, expectedIntValues, _projectionBlock, _dataSourceMap); } + @Test + public void testDistinctFromLeftNullRightLiteral() + throws Exception { + ExpressionContext expression = RequestContextUtils.getExpression(String.format(_expression, "NULL", "1")); + boolean[] expectedIntValues = new boolean[NUM_ROWS]; + if (_isDistinctFrom) { + Arrays.fill(expectedIntValues, true); + } + testTransformFunction(expression, expectedIntValues, _projectionBlock, _dataSourceMap); + } + @Test public void testDistinctFromLeftFunctionRightIdentifier() throws Exception { diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java index 5d867bf12693..808f117b83f4 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java @@ -196,9 +196,10 @@ public Predicate(List operands, DataSchema dataSchema, IntPredica ColumnDataType lhsType = _lhs.getResultType(); ColumnDataType rhsType = _rhs.getResultType(); - Preconditions.checkArgument(lhsType.supportsOrdering() && rhsType.supportsOrdering(), + Preconditions.checkArgument((lhsType == ColumnDataType.UNKNOWN || lhsType.supportsOrdering()) + && (rhsType == ColumnDataType.UNKNOWN || rhsType.supportsOrdering()), "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); - if (lhsType == rhsType) { + if (lhsType == ColumnDataType.UNKNOWN || rhsType == ColumnDataType.UNKNOWN || lhsType == rhsType) { _requireCasting = false; _commonCastType = null; } else { diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java index a9561ca709ba..62ed1d5f50f4 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java @@ -27,11 +27,33 @@ public class FilterOperandTest { + private static final DataSchema INT_SCHEMA = + new DataSchema(new String[]{"value"}, new ColumnDataType[]{ColumnDataType.INT}); private static final DataSchema VARIANT_SCHEMA = new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + private static final RexExpression NULL_LITERAL = new RexExpression.Literal(ColumnDataType.UNKNOWN, null); private static final List VARIANT_OPERANDS = List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(0)); + @Test + public void testComparisonWithNullLiteral() { + FilterOperand.Predicate leftNull = new FilterOperand.Predicate( + List.of(NULL_LITERAL, new RexExpression.InputRef(0)), INT_SCHEMA, value -> value == 0); + Assert.assertNull(leftNull.apply(List.of(1))); + + FilterOperand.Predicate rightNull = new FilterOperand.Predicate( + List.of(new RexExpression.InputRef(0), NULL_LITERAL), INT_SCHEMA, value -> value == 0); + Assert.assertNull(rightNull.apply(List.of(1))); + } + + @Test + public void testRawVariantComparisonIsRejected() { + IllegalArgumentException exception = + Assert.expectThrows(IllegalArgumentException.class, + () -> new FilterOperand.Predicate(VARIANT_OPERANDS, VARIANT_SCHEMA, value -> value == 0)); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support comparison")); + } + @Test public void testRawVariantInIsRejected() { IllegalArgumentException exception = From 0eb7b2ce5f41e83aff29a7772b786faf86880150 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 28 Jul 2026 15:33:36 -0700 Subject: [PATCH 3/8] Harden VARIANT mixed-version and execution behavior Exercise new-broker/old-server, mixed-server, and old-broker/new-server compatibility, including deterministic handling of the VARIANT wire type before server upgrade and rejection of partial mixed-fleet results. Centralize join-key validation, clarify reusable extraction ownership, and strengthen scalar, null-placeholder, window, and end-to-end coverage. --- compatibility-verifier/compCheck.sh | 39 +++- .../config/queries/routing-ready.queries | 21 ++ .../config/queries/variant-wire-mixed.queries | 21 ++ .../config/queries/variant-wire.queries | 21 ++ .../query-results/routing-ready.results | 21 ++ .../query-results/variant-wire-mixed.results | 20 ++ .../config/query-results/variant-wire.results | 20 ++ .../old-broker-new-servers.yaml | 33 +++ .../post-server-2-upgrade.yaml | 6 + .../post-server-upgrade.yaml | 16 ++ .../pre-server-upgrade.yaml | 23 ++ .../config/queries/routing-ready.queries | 20 ++ .../query-results/routing-ready.results | 20 ++ .../old-broker-new-servers.yaml | 31 +++ .../apache/pinot/common/utils/DataSchema.java | 4 +- .../pinot/common/utils/VariantUtils.java | 15 +- .../function/scalar/VariantFunctionsTest.java | 57 +++++ .../java/org/apache/pinot/compat/BaseOp.java | 52 ++++- .../pinot/compat/CompatibilityOpsRunner.java | 3 + .../apache/pinot/compat/FileContainsOp.java | 131 +++++++++++ .../java/org/apache/pinot/compat/QueryOp.java | 215 +++++++++++++++--- .../java/org/apache/pinot/compat/Utils.java | 24 +- .../org/apache/pinot/compat/BaseOpTest.java | 80 +++++++ .../pinot/compat/FileContainsOpTest.java | 57 +++++ .../org/apache/pinot/compat/QueryOpTest.java | 195 ++++++++++++++++ .../VariantGetTransformFunctionTest.java | 41 +++- .../tests/custom/VariantTypeTest.java | 24 +- .../validation/JoinKeyTypeValidator.java | 76 +++++++ .../VariantTypeValidationVisitor.java | 29 +-- .../VariantTypeValidationVisitorTest.java | 4 +- .../runtime/operator/HashJoinOperator.java | 21 +- .../factory/DefaultJoinOperatorFactory.java | 27 +-- .../operator/WindowAggregateOperatorTest.java | 34 ++- pinot-spi/VARIANT_DESIGN.md | 14 +- 34 files changed, 1297 insertions(+), 118 deletions(-) create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/routing-ready.queries create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire-mixed.queries create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire.queries create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/routing-ready.results create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire-mixed.results create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire.results create mode 100644 compatibility-verifier/multi-stage-query-engine-test-suite/old-broker-new-servers.yaml create mode 100644 compatibility-verifier/sample-test-suite/config/queries/routing-ready.queries create mode 100644 compatibility-verifier/sample-test-suite/config/query-results/routing-ready.results create mode 100644 compatibility-verifier/sample-test-suite/old-broker-new-servers.yaml create mode 100644 pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java create mode 100644 pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java create mode 100644 pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java create mode 100644 pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/FileContainsOpTest.java create mode 100644 pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/QueryOpTest.java create mode 100644 pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java diff --git a/compatibility-verifier/compCheck.sh b/compatibility-verifier/compCheck.sh index 7ab06270fb80..ca8f832b2290 100755 --- a/compatibility-verifier/compCheck.sh +++ b/compatibility-verifier/compCheck.sh @@ -22,6 +22,8 @@ # from one version to the other given 2 commit hashes. It first builds # Pinot in the 2 given directories and then upgrades in the following order: # Controller -> Broker -> Server +# An optional old-broker-new-servers.yaml phase temporarily restores the old +# broker after all servers are upgraded to verify the opposite wire direction. # # TODO Some ideas to explore: # It will be nice to have the script take arguments about what is to be done. @@ -496,14 +498,24 @@ setupControllerVariables setupBrokerVariables setupServerVariables -export JAVA_OPTS="-DControllerPort=${CONTROLLER_PORT} -DBrokerQueryPort=${BROKER_QUERY_PORT} -DServerAdminPort=${SERVER_ADMIN_PORT}" - mkdir ${PID_DIR} mkdir ${LOG_DIR} oldTargetDir="$workingDir"/oldTargetDir newTargetDir="$workingDir"/newTargetDir +oldServerSupportsVariant=false +oldExpressionsProto="${oldTargetDir}/pinot-common/src/main/proto/expressions.proto" +if [ -f "${oldExpressionsProto}" ] \ + && grep -Eq '^[[:space:]]*VARIANT[[:space:]]*=[[:space:]]*24[[:space:]]*;' "${oldExpressionsProto}"; then + oldServerSupportsVariant=true +fi +echo "Old server supports VARIANT query-wire type: ${oldServerSupportsVariant}" + +export JAVA_OPTS="-DControllerPort=${CONTROLLER_PORT} -DBrokerQueryPort=${BROKER_QUERY_PORT} \ +-DServerAdminPort=${SERVER_ADMIN_PORT} -Dpinot.compat.oldServerSupportsVariant=${oldServerSupportsVariant} \ +-Dpinot.compat.logDir=${LOG_DIR}" + setupCompatTester # check that the default ports are open @@ -604,7 +616,30 @@ if [ -f "${SERVER_CONF_2}" ]; then exit 1 fi fi +fi + +if [ -f "$testSuiteDir/old-broker-new-servers.yaml" ]; then + echo "Temporarily downgrading broker to test the old broker with upgraded servers" + stopService broker + startService broker "$oldTargetDir" "$BROKER_CONF" + waitForBrokerReady + + "$COMPAT_TESTER" "$testSuiteDir/old-broker-new-servers.yaml" "$genNum" + if [ $? -ne 0 ]; then + if [ $keepClusterOnFailure == "false" ]; then + stopServices + fi + echo "Failed with old broker and upgraded servers" + exit 1 + fi + echo "Restoring upgraded broker before server rollback" + stopService broker + startService broker "$newTargetDir" "$BROKER_CONF" + waitForBrokerReady +fi + +if [ -f "${SERVER_CONF_2}" ]; then echo "Downgrading server 2" # Upgrade completed, now do a rollback stopService server2 diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/routing-ready.queries b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/routing-ready.queries new file mode 100644 index 000000000000..88a2c45b1808 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/routing-ready.queries @@ -0,0 +1,21 @@ +# +# 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. +# + +SELECT COUNT(*) FROM FeatureTest1 WHERE generationNumber = __GENERATION_NUMBER__ +SELECT COUNT(*) FROM FeatureTest2 WHERE generationNumber = __GENERATION_NUMBER__ diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire-mixed.queries b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire-mixed.queries new file mode 100644 index 000000000000..81f0b914dcb3 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire-mixed.queries @@ -0,0 +1,21 @@ +# +# 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. +# + +# Scan all four uploaded segments so the leaf stage reaches both the upgraded and old server. +SELECT COUNT(*) FROM FeatureTest1 WHERE variant_get(parse_json(concat(concat('"', stringDimSV1), '"')), '$', 'STRING') = stringDimSV1 diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire.queries b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire.queries new file mode 100644 index 000000000000..81d019ae145b --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/queries/variant-wire.queries @@ -0,0 +1,21 @@ +# +# 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. +# + +# Keep the VARIANT-producing call in a leaf filter so the server is part of the wire-compatibility check. +SELECT COUNT(*) FROM FeatureTest1 WHERE generationNumber = __GENERATION_NUMBER__ AND variant_get(parse_json(concat(concat('"', stringDimSV1), '"')), '$', 'STRING') = stringDimSV1 diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/routing-ready.results b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/routing-ready.results new file mode 100644 index 000000000000..53d8cdb48f04 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/routing-ready.results @@ -0,0 +1,21 @@ +# +# 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. +# + +{"resultTable":{"dataSchema":{"columnNames":["EXPR$0"],"columnDataTypes":["LONG"]},"rows":[[10]]},"exceptions":[],"numDocsScanned":10} +{"resultTable":{"dataSchema":{"columnNames":["EXPR$0"],"columnDataTypes":["LONG"]},"rows":[[66]]},"exceptions":[],"numDocsScanned":66} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire-mixed.results b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire-mixed.results new file mode 100644 index 000000000000..8f6183fa0d5c --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire-mixed.results @@ -0,0 +1,20 @@ +# +# 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. +# + +{"resultTable":{"dataSchema":{"columnNames":["EXPR$0"],"columnDataTypes":["LONG"]},"rows":[[40]]},"exceptions":[],"numDocsScanned":40} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire.results b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire.results new file mode 100644 index 000000000000..d573b9e73d9e --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/config/query-results/variant-wire.results @@ -0,0 +1,20 @@ +# +# 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. +# + +{"resultTable":{"dataSchema":{"columnNames":["EXPR$0"],"columnDataTypes":["LONG"]},"rows":[[10]]},"exceptions":[],"numDocsScanned":10} diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/old-broker-new-servers.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/old-broker-new-servers.yaml new file mode 100644 index 000000000000..5924c42c15e1 --- /dev/null +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/old-broker-new-servers.yaml @@ -0,0 +1,33 @@ +# +# 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. +# + +description: Operations to be run with the old broker and upgraded servers +operations: + - type: queryOp + description: Wait for the restarted old broker to rebuild its routing tables + useMultiStageQueryEngine: true + queryFileName: queries/routing-ready.queries + expectedResultsFileName: query-results/routing-ready.results + maxAttempts: 60 + retryDelayMs: 1000 + - type: queryOp + description: Verify the old broker's pre-VARIANT plans remain compatible with upgraded servers + useMultiStageQueryEngine: true + queryFileName: queries/feature-test-multi-stage.queries + expectedResultsFileName: query-results/feature-test-multi-stage.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-2-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-2-upgrade.yaml index 6dbb916845cd..ef0aa5000588 100644 --- a/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-2-upgrade.yaml +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-2-upgrade.yaml @@ -41,3 +41,9 @@ operations: useMultiStageQueryEngine: true queryFileName: queries/feature-test-multi-stage.queries expectedResultsFileName: query-results/feature-test-multi-stage.results + - type: queryOp + description: Verify the VARIANT wire query after both servers are upgraded + useMultiStageQueryEngine: true + nullHandlingEnabled: true + queryFileName: queries/variant-wire.queries + expectedResultsFileName: query-results/variant-wire.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml index 786cb3fb3981..86b02dfc13ac 100644 --- a/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/post-server-upgrade.yaml @@ -41,3 +41,19 @@ operations: useMultiStageQueryEngine: true queryFileName: queries/feature-test-multi-stage.queries expectedResultsFileName: query-results/feature-test-multi-stage.results + - type: queryOp + description: Reject the VARIANT wire query without partial rows while the server fleet is mixed + runIfSystemProperty: pinot.compat.oldServerSupportsVariant + runIfSystemPropertyValue: "false" + useMultiStageQueryEngine: true + nullHandlingEnabled: true + queryFileName: queries/variant-wire-mixed.queries + expectedErrorMessageContains: "Caught exception while deserializing stage plan" + - type: queryOp + description: Run the VARIANT wire query across the mixed fleet when the old server supports the type + runIfSystemProperty: pinot.compat.oldServerSupportsVariant + runIfSystemPropertyValue: "true" + useMultiStageQueryEngine: true + nullHandlingEnabled: true + queryFileName: queries/variant-wire-mixed.queries + expectedResultsFileName: query-results/variant-wire-mixed.results diff --git a/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml index 167f71fadf2f..2b2157033906 100644 --- a/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml +++ b/compatibility-verifier/multi-stage-query-engine-test-suite/pre-server-upgrade.yaml @@ -41,3 +41,26 @@ operations: useMultiStageQueryEngine: true queryFileName: queries/feature-test-multi-stage.queries expectedResultsFileName: query-results/feature-test-multi-stage.results + - type: queryOp + description: Report the deterministic VARIANT wire-type error while servers are still on the old version + runIfSystemProperty: pinot.compat.oldServerSupportsVariant + runIfSystemPropertyValue: "false" + useMultiStageQueryEngine: true + nullHandlingEnabled: true + queryFileName: queries/variant-wire.queries + expectedErrorMessageContains: "Caught exception while deserializing stage plan" + - type: queryOp + description: Run the VARIANT wire query when the old servers already advertise the type + runIfSystemProperty: pinot.compat.oldServerSupportsVariant + runIfSystemPropertyValue: "true" + useMultiStageQueryEngine: true + nullHandlingEnabled: true + queryFileName: queries/variant-wire.queries + expectedResultsFileName: query-results/variant-wire.results + - type: fileContainsOp + description: Verify the exact unsupported VARIANT wire type in the old server log + runIfSystemProperty: pinot.compat.oldServerSupportsVariant + runIfSystemPropertyValue: "false" + directorySystemProperty: pinot.compat.logDir + fileNameGlob: "server*.log" + expectedTextContains: "Unsupported proto ColumnDataType: UNRECOGNIZED" diff --git a/compatibility-verifier/sample-test-suite/config/queries/routing-ready.queries b/compatibility-verifier/sample-test-suite/config/queries/routing-ready.queries new file mode 100644 index 000000000000..7bac823c16ab --- /dev/null +++ b/compatibility-verifier/sample-test-suite/config/queries/routing-ready.queries @@ -0,0 +1,20 @@ +# +# 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. +# + +SELECT COUNT(*) FROM FeatureTest1 WHERE generationNumber = __GENERATION_NUMBER__ diff --git a/compatibility-verifier/sample-test-suite/config/query-results/routing-ready.results b/compatibility-verifier/sample-test-suite/config/query-results/routing-ready.results new file mode 100644 index 000000000000..60667c7ec5f3 --- /dev/null +++ b/compatibility-verifier/sample-test-suite/config/query-results/routing-ready.results @@ -0,0 +1,20 @@ +# +# 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. +# + +{"resultTable":{"dataSchema":{"columnDataTypes":["LONG"],"columnNames":["count(*)"]},"rows":[[10]]},"exceptions":[],"numServersQueried":1,"numServersResponded":1,"numSegmentsQueried":1,"numSegmentsProcessed":1,"numSegmentsMatched":1,"numDocsScanned":10,"numEntriesScannedPostFilter":0,"numGroupsLimitReached":false,"totalDocs":10,"timeUsedMs":4,"segmentStatistics":[],"traceInfo":{},"minConsumingFreshnessTimeMs":0} diff --git a/compatibility-verifier/sample-test-suite/old-broker-new-servers.yaml b/compatibility-verifier/sample-test-suite/old-broker-new-servers.yaml new file mode 100644 index 000000000000..5b6819ee3414 --- /dev/null +++ b/compatibility-verifier/sample-test-suite/old-broker-new-servers.yaml @@ -0,0 +1,31 @@ +# +# 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. +# + +description: Operations to be run with the old broker and upgraded servers +operations: + - type: queryOp + description: Wait for the restarted old broker to rebuild the FeatureTest1 routing table + queryFileName: queries/routing-ready.queries + expectedResultsFileName: query-results/routing-ready.results + maxAttempts: 60 + retryDelayMs: 1000 + - type: queryOp + description: Verify single-stage queries through the old broker and upgraded servers + queryFileName: queries/feature-test-1-sql.queries + expectedResultsFileName: query-results/feature-test-1-rest-sql.results 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 8464d0329ea7..082bc7625a90 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 @@ -337,8 +337,8 @@ public RelDataType toType(RelDataTypeFactory typeFactory) { return typeFactory.createSqlType(SqlTypeName.VARIANT); } }, - // UUID is a logical type backed by BYTES; keep it directly after BYTES. ColumnDataType is serialized by name (not - // ordinal) via name()/valueOf(), so enum order does not affect wire compatibility. + // UUID is a logical type backed by BYTES. ColumnDataType is serialized by name (not ordinal) via name()/valueOf(), + // so enum order does not affect wire compatibility. UUID(BYTES, null) { @Override public RelDataType toType(RelDataTypeFactory typeFactory) { diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java index 38fb47ac2b13..0239a9f13bff 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java @@ -157,7 +157,10 @@ private VariantPath(PathElement[] elements) { * Reusable, unboxed destination for vectorized Variant extraction. * *

Only the getter corresponding to the requested {@link ResultType} is defined after a successful extraction. - * The instance is mutable and not thread-safe; callers should retain one per transform-function instance. + * The instance is mutable and not thread-safe; callers should retain one per transform-function instance. Every + * extraction may replace its state. Each successful byte-valued extraction installs a newly materialized array. + * Values returned as {@code byte[]} or as a {@link ByteArray} may be retained after this result is reused, but they + * are read-only by contract and must be copied before mutation. */ public static final class ReusableResult { private final Cursor _cursor = new Cursor(); @@ -195,6 +198,9 @@ public String getStringValue() { /** * Returns the extracted BYTES, VARIANT, or direct 16-byte UUID representation. + * + *

The returned array is replaced, but not mutated, by the next byte-valued extraction. It may be retained after + * this result is reused, but must be treated as immutable and copied before mutation. */ public byte[] getBytesValue() { return _bytesValue; @@ -206,6 +212,9 @@ public UUID getUuidValue() { /** * Materializes the extracted value in the external representation used by scalar functions and ingestion. + * + *

For BYTES and VARIANT, the returned {@code byte[]} may be retained after this result is reused. It must be + * treated as immutable and copied before mutation. */ public Object getExternalValue(ResultType resultType) { switch (resultType) { @@ -240,7 +249,9 @@ public Object getExternalValue(ResultType resultType) { * Materializes the extracted value in {@link DataSchema}'s internal representation. * *

TIMESTAMP remains epoch milliseconds and UUID wraps the directly copied 16-byte value, avoiding an - * external-object round trip in the multi-stage engine. + * external-object round trip in the multi-stage engine. For BYTES, UUID, and VARIANT, the returned + * {@link ByteArray} wraps a newly materialized array that may be retained after this result is reused. Neither the + * wrapper nor its array may be mutated; callers must copy the array before mutation. */ public Object getInternalValue(ResultType resultType) { switch (resultType) { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java new file mode 100644 index 000000000000..7ad55045abe1 --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java @@ -0,0 +1,57 @@ +/** + * 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.function.scalar; + +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.assertTrue; + + +/** + * Tests strict and tolerant behavior exposed by the public scalar VARIANT function facade. + */ +public class VariantFunctionsTest { + @Test + public void testScalarFunctionFacade() { + byte[] variant = VariantFunctions.parseJsonToVariant( + "{\"eventType\":\"click\",\"payload\":{\"count\":7},\"variantNull\":null}"); + + assertEquals(VariantFunctions.variantGet(variant, "$.eventType", "STRING"), "click"); + assertEquals(VariantFunctions.variantToJson(VariantFunctions.variantGet(variant, "$.payload")), "{\"count\":7}"); + assertEquals(VariantFunctions.tryVariantGet(variant, "$.payload", "JSON"), "{\"count\":7}"); + assertEquals(VariantFunctions.variantToJson(VariantFunctions.tryVariantGet(variant, "$.eventType")), "\"click\""); + assertTrue(VariantFunctions.variantExists(variant, "$.variantNull")); + assertFalse(VariantFunctions.variantExists(variant, "$.missing")); + assertFalse(VariantFunctions.isVariantNull(variant)); + assertTrue(VariantFunctions.isVariantNull(variant, "$.variantNull")); + assertEquals(VariantFunctions.variantTypeOf(variant), "OBJECT"); + assertEquals(VariantFunctions.variantTypeOf(variant, "$.payload"), "OBJECT"); + assertEquals(VariantFunctions.variantToJson(variant), + "{\"eventType\":\"click\",\"payload\":{\"count\":7},\"variantNull\":null}"); + + byte[] tolerant = VariantFunctions.tryParseJsonToVariant("{\"value\":11}"); + assertEquals(VariantFunctions.variantGet(tolerant, "$.value", "INT"), 11); + assertNull(VariantFunctions.tryParseJsonToVariant("{not-json")); + assertNull(VariantFunctions.tryVariantGet(variant, "$.missing")); + assertNull(VariantFunctions.tryVariantGet(variant, "$.missing", "STRING")); + } +} diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java index 2869adda0dcb..193dd7df1c23 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.annotation.JsonTypeInfo; import java.net.URI; import java.util.Properties; +import javax.annotation.Nullable; import org.apache.pinot.client.PinotClientException; import org.apache.pinot.client.admin.PinotAdminClient; import org.apache.pinot.client.admin.PinotAdminTransport; @@ -34,11 +35,12 @@ @JsonSubTypes.Type(value = SegmentOp.class, name = "segmentOp"), @JsonSubTypes.Type(value = TableOp.class, name = "tableOp"), @JsonSubTypes.Type(value = QueryOp.class, name = "queryOp"), - @JsonSubTypes.Type(value = StreamOp.class, name = "streamOp") + @JsonSubTypes.Type(value = StreamOp.class, name = "streamOp"), + @JsonSubTypes.Type(value = FileContainsOp.class, name = "fileContainsOp") }) public abstract class BaseOp { enum OpType { - TABLE_OP, SEGMENT_OP, QUERY_OP, STREAM_OP, + TABLE_OP, SEGMENT_OP, QUERY_OP, STREAM_OP, FILE_CONTAINS_OP, } private String _name; @@ -48,6 +50,10 @@ enum OpType { protected static final String GENERATION_NUMBER_PLACEHOLDER = "__GENERATION_NUMBER__"; protected static final String CONFIG_PLACEHOLDER = "/config/"; private String _parentDir; + @Nullable + private String _runIfSystemProperty; + @Nullable + private String _runIfSystemPropertyValue; protected BaseOp(OpType opType) { _opType = opType; @@ -81,7 +87,49 @@ public String getAbsoluteFileName(String fileName) { return _parentDir + CONFIG_PLACEHOLDER + fileName; } + /** + * Returns the system property that gates this operation, or {@code null} when the operation is unconditional. + */ + @Nullable + public String getRunIfSystemProperty() { + return _runIfSystemProperty; + } + + public void setRunIfSystemProperty(@Nullable String runIfSystemProperty) { + _runIfSystemProperty = runIfSystemProperty; + } + + /** + * Returns the exact property value required to run this operation, or {@code null} when the operation is + * unconditional. + */ + @Nullable + public String getRunIfSystemPropertyValue() { + return _runIfSystemPropertyValue; + } + + public void setRunIfSystemPropertyValue(@Nullable String runIfSystemPropertyValue) { + _runIfSystemPropertyValue = runIfSystemPropertyValue; + } + public boolean run(int generationNumber) { + if ((_runIfSystemProperty == null) != (_runIfSystemPropertyValue == null) + || _runIfSystemProperty != null && (_runIfSystemProperty.isBlank() || _runIfSystemPropertyValue.isBlank())) { + LOGGER.error("Both runIfSystemProperty and a non-blank runIfSystemPropertyValue must be configured together"); + return false; + } + if (_runIfSystemProperty != null) { + String actualValue = System.getProperty(_runIfSystemProperty); + if (actualValue == null) { + LOGGER.error("Required system property {} is not configured", _runIfSystemProperty); + return false; + } + if (!_runIfSystemPropertyValue.equals(actualValue)) { + LOGGER.info("Skipping OpType {} because system property {} is '{}', expected '{}'", _opType, + _runIfSystemProperty, actualValue, _runIfSystemPropertyValue); + return true; + } + } LOGGER.info("Running OpType {} : {}", _opType.toString(), getDescription()); return runOp(generationNumber); } diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/CompatibilityOpsRunner.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/CompatibilityOpsRunner.java index eba60c62c0f0..6d6c2af132cb 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/CompatibilityOpsRunner.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/CompatibilityOpsRunner.java @@ -115,6 +115,9 @@ protected Object constructObject(Node node) { case "streamOp": node.setType(StreamOp.class); break; + case "fileContainsOp": + node.setType(FileContainsOp.class); + break; default: throw new RuntimeException("Unknown type: " + type); } diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java new file mode 100644 index 000000000000..0448da66f5b8 --- /dev/null +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java @@ -0,0 +1,131 @@ +/** + * 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.compat; + +import java.io.BufferedReader; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.DirectoryStream; +import java.nio.file.Files; +import java.nio.file.Path; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Verifies that at least one file matching a glob in a system-property directory contains an expected text fragment. + * Instances are configured and invoked serially by the compatibility runner and are not thread-safe. + */ +public class FileContainsOp extends BaseOp { + private static final Logger LOGGER = LoggerFactory.getLogger(FileContainsOp.class); + private static final int MAX_ATTEMPTS = 10; + private static final long RETRY_DELAY_MS = 100L; + + private String _directorySystemProperty; + private String _fileNameGlob; + private String _expectedTextContains; + + public FileContainsOp() { + super(OpType.FILE_CONTAINS_OP); + } + + public String getDirectorySystemProperty() { + return _directorySystemProperty; + } + + public void setDirectorySystemProperty(String directorySystemProperty) { + _directorySystemProperty = directorySystemProperty; + } + + public String getFileNameGlob() { + return _fileNameGlob; + } + + public void setFileNameGlob(String fileNameGlob) { + _fileNameGlob = fileNameGlob; + } + + public String getExpectedTextContains() { + return _expectedTextContains; + } + + public void setExpectedTextContains(String expectedTextContains) { + _expectedTextContains = expectedTextContains; + } + + @Override + boolean runOp(int generationNumber) { + if (isBlank(_directorySystemProperty) || isBlank(_fileNameGlob) || isBlank(_expectedTextContains)) { + LOGGER.error( + "directorySystemProperty, fileNameGlob, and expectedTextContains must all be configured and non-blank"); + return false; + } + String directoryName = System.getProperty(_directorySystemProperty); + if (isBlank(directoryName)) { + LOGGER.error("Required log-directory system property {} is not configured", _directorySystemProperty); + return false; + } + Path directory = Path.of(directoryName); + if (!Files.isDirectory(directory)) { + LOGGER.error("Configured log directory does not exist: {}", directory); + return false; + } + + try { + for (int attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) { + if (containsExpectedText(directory, _fileNameGlob, _expectedTextContains)) { + return true; + } + if (attempt < MAX_ATTEMPTS) { + Thread.sleep(RETRY_DELAY_MS); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.error("Interrupted while waiting for expected text in {} files under {}", _fileNameGlob, directory); + return false; + } catch (IOException e) { + LOGGER.error("Failed to inspect {} files under {}", _fileNameGlob, directory, e); + return false; + } + LOGGER.error("No {} file under {} contained expected text: {}", _fileNameGlob, directory, _expectedTextContains); + return false; + } + + static boolean containsExpectedText(Path directory, String fileNameGlob, String expectedText) + throws IOException { + try (DirectoryStream paths = Files.newDirectoryStream(directory, fileNameGlob)) { + for (Path path : paths) { + try (BufferedReader reader = Files.newBufferedReader(path, StandardCharsets.UTF_8)) { + String line; + while ((line = reader.readLine()) != null) { + if (line.contains(expectedText)) { + return true; + } + } + } + } + } + return false; + } + + private static boolean isBlank(String value) { + return value == null || value.isBlank(); + } +} diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java index 536ab1024e70..94c5bb7fcb42 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java @@ -22,8 +22,10 @@ import com.fasterxml.jackson.databind.JsonNode; import java.io.BufferedReader; import java.io.FileInputStream; +import java.io.IOException; import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; +import javax.annotation.Nullable; import org.apache.pinot.common.utils.SqlResultComparator; import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.tools.utils.ExplainPlanUtils; @@ -31,8 +33,8 @@ import org.slf4j.LoggerFactory; -/// Executes queries in the query file, and compares the results with the ones given in the -/// expected results file +/// Executes queries in the query file and verifies either result rows or an expected error-message substring. +/// Exactly one of `expectedResultsFileName` and `expectedErrorMessageContains` must be configured. /// /// TODO: /// - If we use current timestamp for realtime tables, we may not be able to use pre-canned queries. @@ -42,8 +44,14 @@ public class QueryOp extends BaseOp { private static final String COMMENT_DELIMITER = "#"; private String _queryFileName; + @Nullable private String _expectedResultsFileName; + @Nullable + private String _expectedErrorMessageContains; private boolean _useMultiStageQueryEngine = false; + private boolean _enableNullHandling = false; + private int _maxAttempts = 1; + private long _retryDelayMs = 1000L; public QueryOp() { super(OpType.QUERY_OP); @@ -62,14 +70,36 @@ public void setQueryFileName(String queryFileName) { _queryFileName = queryFileName; } + /** + * Returns the expected-results file, or {@code null} when expected-error mode is configured. + */ + @Nullable public String getExpectedResultsFileName() { return _expectedResultsFileName; } - public void setExpectedResultsFileName(String expectedResultsFileName) { + /** + * Configures result-comparison mode. Pass {@code null} to clear it before configuring expected-error mode. + */ + public void setExpectedResultsFileName(@Nullable String expectedResultsFileName) { _expectedResultsFileName = expectedResultsFileName; } + /** + * Returns the required error-message substring, or {@code null} when result-comparison mode is configured. + */ + @Nullable + public String getExpectedErrorMessageContains() { + return _expectedErrorMessageContains; + } + + /** + * Configures expected-error mode. Pass {@code null} to clear it before configuring result-comparison mode. + */ + public void setExpectedErrorMessageContains(@Nullable String expectedErrorMessageContains) { + _expectedErrorMessageContains = expectedErrorMessageContains; + } + public boolean getUseMultiStageQueryEngine() { return _useMultiStageQueryEngine; } @@ -78,31 +108,99 @@ public void setUseMultiStageQueryEngine(boolean useMultiStageQueryEngine) { _useMultiStageQueryEngine = useMultiStageQueryEngine; } + public boolean isNullHandlingEnabled() { + return _enableNullHandling; + } + + public void setNullHandlingEnabled(boolean nullHandlingEnabled) { + _enableNullHandling = nullHandlingEnabled; + } + + public int getMaxAttempts() { + return _maxAttempts; + } + + public void setMaxAttempts(int maxAttempts) { + _maxAttempts = maxAttempts; + } + + public long getRetryDelayMs() { + return _retryDelayMs; + } + + public void setRetryDelayMs(long retryDelayMs) { + _retryDelayMs = retryDelayMs; + } + @Override boolean runOp(int generationNumber) { - LOGGER.info("Verifying queries in {} against results in {}", _queryFileName, _expectedResultsFileName); + if (!hasValidExpectedOutcomeConfiguration()) { + LOGGER.error( + "Exactly one of expectedResultsFileName or expectedErrorMessageContains must be configured for queries in {}", + _queryFileName); + return false; + } + if (_maxAttempts < 1 || _retryDelayMs < 0) { + LOGGER.error("maxAttempts must be positive and retryDelayMs must be non-negative for queries in {}", + _queryFileName); + return false; + } + if (_expectedErrorMessageContains != null) { + LOGGER.info("Verifying queries in {} fail with an error containing '{}'", _queryFileName, + _expectedErrorMessageContains); + } else { + LOGGER.info("Verifying queries in {} against results in {}", _queryFileName, _expectedResultsFileName); + } try { for (int i = 1; i <= generationNumber; i++) { - if (!verifyQueries(i)) { + if (!verifyQueriesWithRetry(i)) { return false; } } return true; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.error("Interrupted while retrying queries in {}", _queryFileName); + return false; } catch (Exception e) { LOGGER.error("FAILED to verify queries in {}: {}", _queryFileName, e); return false; } } + private boolean verifyQueriesWithRetry(int generationNumber) + throws Exception { + for (int attempt = 1; attempt <= _maxAttempts; attempt++) { + if (verifyQueries(generationNumber)) { + return true; + } + if (attempt < _maxAttempts) { + LOGGER.info("Retrying queries in {} after failed attempt {} of {}", _queryFileName, attempt, _maxAttempts); + Thread.sleep(_retryDelayMs); + } + } + return false; + } + + boolean hasValidExpectedOutcomeConfiguration() { + if (_expectedResultsFileName != null && _expectedErrorMessageContains != null) { + return false; + } + boolean hasExpectedResults = _expectedResultsFileName != null && !_expectedResultsFileName.isBlank(); + boolean hasExpectedError = + _expectedErrorMessageContains != null && !_expectedErrorMessageContains.isBlank(); + return hasExpectedResults != hasExpectedError; + } + boolean verifyQueries(int generationNumber) throws Exception { boolean testPassed = false; try (BufferedReader queryReader = new BufferedReader( new InputStreamReader(new FileInputStream(getAbsoluteFileName(_queryFileName)), StandardCharsets.UTF_8)); - BufferedReader expectedResultReader = new BufferedReader( - new InputStreamReader(new FileInputStream(getAbsoluteFileName(_expectedResultsFileName)), - StandardCharsets.UTF_8))) { + BufferedReader expectedResultReader = _expectedErrorMessageContains != null ? null + : new BufferedReader(new InputStreamReader( + new FileInputStream(getAbsoluteFileName(_expectedResultsFileName)), StandardCharsets.UTF_8))) { int succeededQueryCount = 0; int totalQueryCount = 0; @@ -116,38 +214,47 @@ boolean verifyQueries(int generationNumber) } query = query.replaceAll(GENERATION_NUMBER_PLACEHOLDER, String.valueOf(generationNumber)); JsonNode expectedJson = null; - try { - String expectedResultLine = expectedResultReader.readLine(); - while (shouldIgnore(expectedResultLine)) { - expectedResultLine = expectedResultReader.readLine(); + if (_expectedErrorMessageContains == null) { + try { + String expectedResultLine = readNextExpectedResult(expectedResultReader, queryLineNum); + expectedJson = JsonUtils.stringToJsonNode(expectedResultLine); + } catch (Exception e) { + LOGGER.error("Comparison FAILED: Line: {} Exception caught while getting expected response for query: '{}'", + queryLineNum, query, e); } - expectedJson = JsonUtils.stringToJsonNode(expectedResultLine); - } catch (Exception e) { - LOGGER.error("Comparison FAILED: Line: {} Exception caught while getting expected response for query: '{}'", - queryLineNum, query, e); } JsonNode actualJson = null; - if (expectedJson != null) { + if (_expectedErrorMessageContains != null || expectedJson != null) { try { actualJson = _useMultiStageQueryEngine - ? Utils.postMultiStageSqlQuery(query, ClusterDescriptor.getInstance().getBrokerUrl()) - : Utils.postSqlQuery(query, ClusterDescriptor.getInstance().getBrokerUrl()); + ? Utils.postMultiStageSqlQuery(query, ClusterDescriptor.getInstance().getBrokerUrl(), + _enableNullHandling) + : Utils.postSqlQuery(query, ClusterDescriptor.getInstance().getBrokerUrl(), _enableNullHandling); } catch (Exception e) { LOGGER.error("Comparison FAILED: Line: {} Exception caught while running query: '{}', explain plan: {}", queryLineNum, query, getExplainPlan(query), e); } } - if (expectedJson != null && actualJson != null) { + if (actualJson != null && (_expectedErrorMessageContains != null || expectedJson != null)) { try { - boolean passed = _useMultiStageQueryEngine - ? SqlResultComparator.areMultiStageQueriesEqual(actualJson, expectedJson, query) - : SqlResultComparator.areEqual(actualJson, expectedJson, query); + boolean passed = matchesExpectedResponse(actualJson, expectedJson, _expectedErrorMessageContains, + _useMultiStageQueryEngine, query); if (passed) { succeededQueryCount++; - LOGGER.debug("Comparison PASSED: Line: {}, query: '{}', actual response: {}, expected response: {}", - queryLineNum, query, actualJson, expectedJson); + if (_expectedErrorMessageContains != null) { + LOGGER.debug("Comparison PASSED: Line: {}, query: '{}', actual response contains expected error: '{}'", + queryLineNum, query, _expectedErrorMessageContains); + } else { + LOGGER.debug("Comparison PASSED: Line: {}, query: '{}', actual response: {}, expected response: {}", + queryLineNum, query, actualJson, expectedJson); + } + } else if (_expectedErrorMessageContains != null) { + LOGGER.error( + "Comparison FAILED: Line: {}, query: '{}', actual response: {}, expected an exception containing: " + + "'{}'", + queryLineNum, query, actualJson, _expectedErrorMessageContains); } else { LOGGER.error( "Comparison FAILED: Line: {}, query: '{}', actual response: {}, expected response: {}, explain " @@ -172,16 +279,70 @@ boolean verifyQueries(int generationNumber) return testPassed; } + static String readNextExpectedResult(BufferedReader expectedResultReader, int queryLineNum) + throws IOException { + String expectedResultLine; + while ((expectedResultLine = expectedResultReader.readLine()) != null) { + if (!expectedResultLine.trim().isEmpty() && !expectedResultLine.trim().startsWith(COMMENT_DELIMITER)) { + return expectedResultLine; + } + } + throw new IOException("Expected results file ended before query at line " + queryLineNum); + } + + static boolean matchesExpectedResponse(JsonNode actual, @Nullable JsonNode expected, + @Nullable String expectedErrorMessageContains, boolean useMultiStageQueryEngine, String query) + throws IOException { + if (expectedErrorMessageContains != null) { + return hasExpectedError(actual, expectedErrorMessageContains); + } + return useMultiStageQueryEngine + ? SqlResultComparator.areMultiStageQueriesEqual(actual, expected, query) + : SqlResultComparator.areEqual(actual, expected, query); + } + + /** + * Returns whether a response has no result rows and contains only errors matching the required substring. + */ + static boolean hasExpectedError(@Nullable JsonNode response, @Nullable String expectedErrorMessageContains) { + if (response == null || expectedErrorMessageContains == null || expectedErrorMessageContains.isEmpty()) { + return false; + } + JsonNode resultRows = response.path("resultTable").path("rows"); + if ((resultRows.isArray() && !resultRows.isEmpty()) || response.path("numRowsResultSet").asInt() > 0) { + return false; + } + JsonNode exceptions = response.path("exceptions"); + if (exceptions.isArray()) { + if (exceptions.isEmpty()) { + return false; + } + for (JsonNode exception : exceptions) { + if (!hasExpectedErrorMessage(exception, expectedErrorMessageContains)) { + return false; + } + } + return true; + } + return hasExpectedErrorMessage(exceptions, expectedErrorMessageContains); + } + + private static boolean hasExpectedErrorMessage(JsonNode exception, String expectedErrorMessageContains) { + JsonNode message = exception.path("message"); + return message.isTextual() && message.asText().contains(expectedErrorMessageContains); + } + private String getExplainPlan(String query) { try { if (!_useMultiStageQueryEngine) { JsonNode explainPlanResponse = - Utils.postSqlQuery("explain plan for " + query, ClusterDescriptor.getInstance().getBrokerUrl()); + Utils.postSqlQuery("explain plan for " + query, ClusterDescriptor.getInstance().getBrokerUrl(), + _enableNullHandling); return ExplainPlanUtils.formatExplainPlan(explainPlanResponse); } else { JsonNode explainPlanResponse = Utils.postMultiStageSqlQuery("explain plan for " + query, - ClusterDescriptor.getInstance().getBrokerUrl()); + ClusterDescriptor.getInstance().getBrokerUrl(), _enableNullHandling); return ExplainPlanUtils.formatMultiStageExplainPlan(explainPlanResponse); } } catch (Throwable error) { diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/Utils.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/Utils.java index edae1b17795c..2713c8a9fe39 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/Utils.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/Utils.java @@ -50,9 +50,14 @@ public static void replaceContent(File originalDataFile, File replacedDataFile, public static JsonNode postSqlQuery(String query, String brokerBaseApiUrl) throws IOException { + return postSqlQuery(query, brokerBaseApiUrl, false); + } + + public static JsonNode postSqlQuery(String query, String brokerBaseApiUrl, boolean enableNullHandling) + throws IOException { ObjectNode payload = JsonUtils.newObjectNode(); payload.put("sql", query); - payload.put("queryOptions", "groupByMode=sql;responseFormat=sql"); + payload.put("queryOptions", getSingleStageQueryOptions(enableNullHandling)); return JsonUtils.stringToJsonNode( ControllerTest.sendPostRequest(brokerBaseApiUrl + "/query/sql", payload.toString())); @@ -60,11 +65,26 @@ public static JsonNode postSqlQuery(String query, String brokerBaseApiUrl) public static JsonNode postMultiStageSqlQuery(String query, String brokerBaseApiUrl) throws IOException { + return postMultiStageSqlQuery(query, brokerBaseApiUrl, false); + } + + public static JsonNode postMultiStageSqlQuery(String query, String brokerBaseApiUrl, boolean enableNullHandling) + throws IOException { ObjectNode payload = JsonUtils.newObjectNode(); payload.put("sql", query); - payload.put("queryOptions", "useMultistageEngine=true"); + payload.put("queryOptions", getMultiStageQueryOptions(enableNullHandling)); return JsonUtils.stringToJsonNode( ControllerTest.sendPostRequest(brokerBaseApiUrl + "/query/sql", payload.toString())); } + + static String getSingleStageQueryOptions(boolean enableNullHandling) { + String queryOptions = "groupByMode=sql;responseFormat=sql"; + return enableNullHandling ? queryOptions + ";enableNullHandling=true" : queryOptions; + } + + static String getMultiStageQueryOptions(boolean enableNullHandling) { + String queryOptions = "useMultistageEngine=true"; + return enableNullHandling ? queryOptions + ";enableNullHandling=true" : queryOptions; + } } diff --git a/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java new file mode 100644 index 000000000000..6a6dc6c55723 --- /dev/null +++ b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java @@ -0,0 +1,80 @@ +/** + * 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.compat; + +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/** + * Tests capability-gated compatibility operations. + */ +public class BaseOpTest { + private static final String TEST_PROPERTY = BaseOpTest.class.getName() + ".enabled"; + + @Test + public void testSystemPropertyRunCondition() { + String previousValue = System.getProperty(TEST_PROPERTY); + try { + CountingOp op = new CountingOp(); + assertTrue(op.run(1)); + assertEquals(op._runCount, 1); + + op.setRunIfSystemProperty(TEST_PROPERTY); + assertFalse(op.run(1)); + assertEquals(op._runCount, 1); + + op.setRunIfSystemPropertyValue("true"); + System.clearProperty(TEST_PROPERTY); + assertFalse(op.run(1)); + assertEquals(op._runCount, 1); + + System.setProperty(TEST_PROPERTY, "false"); + assertTrue(op.run(1)); + assertEquals(op._runCount, 1); + + System.setProperty(TEST_PROPERTY, "true"); + assertTrue(op.run(1)); + assertEquals(op._runCount, 2); + } finally { + if (previousValue == null) { + System.clearProperty(TEST_PROPERTY); + } else { + System.setProperty(TEST_PROPERTY, previousValue); + } + } + } + + private static final class CountingOp extends BaseOp { + private int _runCount; + + private CountingOp() { + super(OpType.QUERY_OP); + } + + @Override + boolean runOp(int generationNumber) { + _runCount++; + return true; + } + } +} diff --git a/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/FileContainsOpTest.java b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/FileContainsOpTest.java new file mode 100644 index 000000000000..7c85172f5706 --- /dev/null +++ b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/FileContainsOpTest.java @@ -0,0 +1,57 @@ +/** + * 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.compat; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +/// Tests exact root-cause matching across compatibility-runner log files. +public class FileContainsOpTest { + + @Test + public void testContainsExpectedTextRequiresMatchingFileAndText() + throws Exception { + Path directory = Files.createTempDirectory("pinot-file-contains-op"); + Path serverLog = directory.resolve("server.4.log"); + Path brokerLog = directory.resolve("broker.3.log"); + try { + Files.writeString(serverLog, + "Caused by: java.lang.IllegalStateException: Unsupported proto ColumnDataType: UNRECOGNIZED\n", + StandardCharsets.UTF_8); + Files.writeString(brokerLog, "Caught exception while deserializing stage plan\n", StandardCharsets.UTF_8); + + assertTrue(FileContainsOp.containsExpectedText(directory, "server*.log", + "Unsupported proto ColumnDataType: UNRECOGNIZED")); + assertFalse(FileContainsOp.containsExpectedText(directory, "server*.log", + "Caught exception while deserializing stage plan")); + assertFalse(FileContainsOp.containsExpectedText(directory, "missing.*.log", + "Unsupported proto ColumnDataType: UNRECOGNIZED")); + } finally { + Files.deleteIfExists(serverLog); + Files.deleteIfExists(brokerLog); + Files.deleteIfExists(directory); + } + } +} diff --git a/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/QueryOpTest.java b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/QueryOpTest.java new file mode 100644 index 000000000000..15ce2d01573f --- /dev/null +++ b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/QueryOpTest.java @@ -0,0 +1,195 @@ +/** + * 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.compat; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.BufferedReader; +import java.io.IOException; +import java.io.StringReader; +import org.apache.pinot.spi.utils.JsonUtils; +import org.testng.annotations.Test; +import org.yaml.snakeyaml.DumperOptions; +import org.yaml.snakeyaml.LoaderOptions; +import org.yaml.snakeyaml.Yaml; +import org.yaml.snakeyaml.representer.Representer; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +/// Tests expected query-error matching used by mixed-version compatibility suites. +public class QueryOpTest { + + @Test + public void testMatchesExpectedErrorInExceptionArray() + throws Exception { + JsonNode response = JsonUtils.stringToJsonNode( + "{\"exceptions\":[" + + "{\"errorCode\":200,\"message\":\"java.lang.IllegalStateException: " + + "Unsupported proto ColumnDataType: UNRECOGNIZED\"}," + + "{\"errorCode\":200,\"message\":\"Unsupported proto ColumnDataType: UNRECOGNIZED\"}]}"); + + assertTrue(QueryOp.hasExpectedError(response, "Unsupported proto ColumnDataType: UNRECOGNIZED")); + assertFalse(QueryOp.hasExpectedError(response, "Unsupported proto ColumnDataType: VARIANT")); + + JsonNode mixedResponse = JsonUtils.stringToJsonNode( + "{\"exceptions\":[" + + "{\"errorCode\":200,\"message\":\"Unsupported proto ColumnDataType: UNRECOGNIZED\"}," + + "{\"errorCode\":200,\"message\":\"Server request timed out\"}]}"); + assertFalse(QueryOp.hasExpectedError(mixedResponse, "Unsupported proto ColumnDataType: UNRECOGNIZED")); + + JsonNode partialRowsResponse = JsonUtils.stringToJsonNode( + "{\"resultTable\":{\"rows\":[[10]]},\"numRowsResultSet\":1," + + "\"exceptions\":[{\"message\":\"Unsupported proto ColumnDataType: UNRECOGNIZED\"}]}"); + assertFalse(QueryOp.hasExpectedError(partialRowsResponse, "Unsupported proto ColumnDataType: UNRECOGNIZED")); + } + + @Test + public void testMatchesExpectedErrorInExceptionObject() + throws Exception { + JsonNode response = + JsonUtils.stringToJsonNode("{\"exceptions\":{\"message\":\"Unsupported proto ColumnDataType: UNRECOGNIZED\"}}"); + + assertTrue(QueryOp.hasExpectedError(response, "ColumnDataType: UNRECOGNIZED")); + } + + @Test + public void testDoesNotMatchMissingOrMalformedExceptions() + throws Exception { + String expectedMessage = "Unsupported proto ColumnDataType: UNRECOGNIZED"; + assertFalse(QueryOp.hasExpectedError(null, expectedMessage)); + assertFalse(QueryOp.hasExpectedError(JsonUtils.stringToJsonNode("{}"), expectedMessage)); + assertFalse(QueryOp.hasExpectedError(JsonUtils.stringToJsonNode("{\"exceptions\":[]}"), + expectedMessage)); + assertFalse(QueryOp.hasExpectedError( + JsonUtils.stringToJsonNode("{\"exceptions\":[{\"errorCode\":200}]}"), expectedMessage)); + assertFalse(QueryOp.hasExpectedError( + JsonUtils.stringToJsonNode("{\"exceptions\":[{\"message\":42}]}"), expectedMessage)); + assertFalse(QueryOp.hasExpectedError( + JsonUtils.stringToJsonNode("{\"exceptions\":[{\"message\":\"" + expectedMessage + "\"}]}"), "")); + } + + @Test + public void testExpectedOutcomeConfigurationRequiresExactlyOneMode() { + QueryOp queryOp = new QueryOp(); + assertFalse(queryOp.hasValidExpectedOutcomeConfiguration()); + + queryOp.setExpectedResultsFileName("query-results/results.json"); + assertTrue(queryOp.hasValidExpectedOutcomeConfiguration()); + + queryOp.setExpectedErrorMessageContains(" "); + assertFalse(queryOp.hasValidExpectedOutcomeConfiguration()); + + queryOp.setExpectedErrorMessageContains("Unsupported proto ColumnDataType: UNRECOGNIZED"); + assertFalse(queryOp.hasValidExpectedOutcomeConfiguration()); + + queryOp.setExpectedResultsFileName(null); + assertTrue(queryOp.hasValidExpectedOutcomeConfiguration()); + + queryOp.setExpectedErrorMessageContains(" "); + assertFalse(queryOp.hasValidExpectedOutcomeConfiguration()); + } + + @Test + public void testRetriesAreExplicitAndBounded() { + int[] attempts = {0}; + QueryOp queryOp = new QueryOp() { + @Override + boolean verifyQueries(int generationNumber) { + return ++attempts[0] == 3; + } + }; + queryOp.setQueryFileName("queries/routing-ready.queries"); + queryOp.setExpectedResultsFileName("query-results/variant-wire.results"); + queryOp.setMaxAttempts(3); + queryOp.setRetryDelayMs(0); + + assertTrue(queryOp.runOp(1)); + assertEquals(attempts[0], 3); + + queryOp.setMaxAttempts(0); + assertFalse(queryOp.runOp(1)); + queryOp.setMaxAttempts(1); + queryOp.setRetryDelayMs(-1); + assertFalse(queryOp.runOp(1)); + } + + @Test + public void testExpectedResultReaderReportsTruncatedFiles() + throws IOException { + BufferedReader reader = new BufferedReader(new StringReader("# comment\n\n{\"resultTable\":{}}\n")); + assertEquals(QueryOp.readNextExpectedResult(reader, 21), "{\"resultTable\":{}}"); + assertThrows(IOException.class, () -> QueryOp.readNextExpectedResult(reader, 22)); + } + + @Test + public void testNormalResultComparisonUsesSelectedEngine() + throws Exception { + JsonNode actualSubset = successfulResponse(1, 1); + JsonNode expectedSuperset = JsonUtils.stringToJsonNode( + "{\"resultTable\":{\"dataSchema\":{\"columnNames\":[\"EXPR$0\"],\"columnDataTypes\":[\"LONG\"]}," + + "\"rows\":[[1],[2]]},\"exceptions\":[],\"numDocsScanned\":1,\"isSuperset\":true}"); + String query = "SELECT COUNT(*) FROM testTable"; + + assertTrue(QueryOp.matchesExpectedResponse(actualSubset, expectedSuperset, null, false, query)); + assertFalse(QueryOp.matchesExpectedResponse(actualSubset, expectedSuperset, null, true, query)); + + JsonNode moreExpensiveActual = successfulResponse(1, 2); + JsonNode expected = successfulResponse(1, 1); + assertFalse(QueryOp.matchesExpectedResponse(moreExpensiveActual, expected, null, false, query)); + assertTrue(QueryOp.matchesExpectedResponse(moreExpensiveActual, expected, null, true, query)); + } + + @Test + public void testNullHandlingQueryOptionsAreOptIn() { + assertEquals(Utils.getSingleStageQueryOptions(false), "groupByMode=sql;responseFormat=sql"); + assertEquals(Utils.getSingleStageQueryOptions(true), + "groupByMode=sql;responseFormat=sql;enableNullHandling=true"); + assertEquals(Utils.getMultiStageQueryOptions(false), "useMultistageEngine=true"); + assertEquals(Utils.getMultiStageQueryOptions(true), "useMultistageEngine=true;enableNullHandling=true"); + } + + @Test + public void testNullHandlingYamlPropertyUsesBooleanNaming() { + Representer representer = new Representer(new DumperOptions()); + representer.getPropertyUtils().setSkipMissingProperties(true); + Yaml yaml = new Yaml(new CompatibilityOpsRunner.CustomConstructor(new LoaderOptions()), representer); + + CompatTestOperation operation = yaml.loadAs(""" + description: Test null handling property + operations: + - type: queryOp + queryFileName: queries/test.queries + expectedResultsFileName: query-results/test.results + nullHandlingEnabled: true + """, CompatTestOperation.class); + QueryOp queryOp = (QueryOp) operation.getOperations().get(0); + + assertTrue(queryOp.isNullHandlingEnabled()); + } + + private static JsonNode successfulResponse(int value, int numDocsScanned) + throws Exception { + return JsonUtils.stringToJsonNode( + "{\"resultTable\":{\"dataSchema\":{\"columnNames\":[\"EXPR$0\"],\"columnDataTypes\":[\"LONG\"]}," + + "\"rows\":[[" + value + "]]},\"exceptions\":[],\"numDocsScanned\":" + numDocsScanned + "}"); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java index fcdc49f3cf9a..ba8d7781d99a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java @@ -36,6 +36,7 @@ import org.apache.pinot.core.operator.transform.TransformResultMetadata; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.exception.BadQueryRequestException; +import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder; import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.VariantEnvelope; import org.roaringbitmap.RoaringBitmap; @@ -197,6 +198,40 @@ public void testTryVariantGetNullBitmap() { assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(1, 2, 3)); } + @Test + public void testMissingPathUsesTypedNullPlaceholders() { + byte[] variant = VariantUtils.parseJsonToVariant("{}"); + ValueBlock block = valueBlock(1); + + VariantGetTransformFunction function = typedFunction(variant, "$.missing", "INT"); + assertEquals(function.transformToIntValuesSV(block)[0], NullValuePlaceHolder.INT); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + function = typedFunction(variant, "$.missing", "LONG"); + assertEquals(function.transformToLongValuesSV(block)[0], NullValuePlaceHolder.LONG); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + function = typedFunction(variant, "$.missing", "FLOAT"); + assertEquals(function.transformToFloatValuesSV(block)[0], NullValuePlaceHolder.FLOAT); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + function = typedFunction(variant, "$.missing", "DOUBLE"); + assertEquals(function.transformToDoubleValuesSV(block)[0], NullValuePlaceHolder.DOUBLE); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + function = typedFunction(variant, "$.missing", "BIG_DECIMAL"); + assertEquals(function.transformToBigDecimalValuesSV(block)[0], NullValuePlaceHolder.BIG_DECIMAL); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + function = typedFunction(variant, "$.missing", "STRING"); + assertEquals(function.transformToStringValuesSV(block)[0], NullValuePlaceHolder.STRING); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + + function = typedFunction(variant, "$.missing", "BYTES"); + assertEquals(function.transformToBytesValuesSV(block)[0], NullValuePlaceHolder.BYTES); + assertEquals(function.getNullBitmap(block), RoaringBitmap.bitmapOf(0)); + } + @Test public void testVariantExtractionIsCachedPerBlock() { byte[] variant = VariantUtils.parseJsonToVariant("{\"eventType\":\"click\"}"); @@ -353,8 +388,12 @@ private static List arguments(TransformFunction input, String } private static VariantGetTransformFunction typedFunction(byte[] variant, String targetType) { + return typedFunction(variant, "$", targetType); + } + + private static VariantGetTransformFunction typedFunction(byte[] variant, String path, String targetType) { VariantGetTransformFunction function = new VariantGetTransformFunction(); - function.init(arguments(new BytesTransformFunction(new byte[][]{variant}, null), "$", targetType), Map.of(), true); + function.init(arguments(new BytesTransformFunction(new byte[][]{variant}, null), path, targetType), Map.of(), true); return function; } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java index 10eaceeecb56..e050286c8ded 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java @@ -348,10 +348,30 @@ public void testRawVariantWindowKeysAreRejectedButTypedPathWorks() } JsonNode response = postVariantQuery( - "SELECT " + EVENT_ID + ", COUNT(*) OVER (PARTITION BY variant_get(" + PAYLOAD + "SELECT " + EVENT_ID + ", variant_get(" + PAYLOAD + ", '$.eventType', 'STRING'), " + + "COUNT(*) OVER (PARTITION BY variant_get(" + PAYLOAD + ", '$.eventType', 'STRING')) FROM " + TABLE_NAME + " ORDER BY " + EVENT_ID); assertNoExceptions(response); - Assert.assertEquals(response.get("resultTable").get("rows").size(), NUM_DOCS); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"STRING\",\"STRING\",\"LONG\"]"); + JsonNode rows = response.get("resultTable").get("rows"); + Assert.assertEquals(rows.size(), NUM_DOCS); + Assert.assertEquals(rows.get(0).get(0).asText(), "evt-001"); + Assert.assertEquals(rows.get(0).get(1).asText(), "checkout"); + Assert.assertEquals(rows.get(0).get(2).asLong(), 2L); + Assert.assertEquals(rows.get(1).get(0).asText(), "evt-002"); + Assert.assertEquals(rows.get(1).get(1).asText(), "view"); + Assert.assertEquals(rows.get(1).get(2).asLong(), 1L); + Assert.assertEquals(rows.get(2).get(0).asText(), "evt-003"); + Assert.assertEquals(rows.get(2).get(1).asText(), "checkout"); + Assert.assertEquals(rows.get(2).get(2).asLong(), 2L); + Assert.assertEquals(rows.get(3).get(0).asText(), "evt-004"); + Assert.assertTrue(rows.get(3).get(1).isNull(), "Variant null must extract to the SQL-null partition"); + Assert.assertEquals(rows.get(3).get(2).asLong(), 2L); + Assert.assertEquals(rows.get(4).get(0).asText(), "evt-005"); + Assert.assertTrue(rows.get(4).get(1).isNull(), "SQL null must remain in the SQL-null partition"); + Assert.assertEquals(rows.get(4).get(2).asLong(), 2L); } @Test diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java new file mode 100644 index 000000000000..ee6261b511e8 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.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.query.planner.validation; + +import com.google.common.base.Preconditions; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.query.planner.plannode.JoinNode; + + +/** + * Validates the logical types used as join keys. + * + *

This stateless validator is shared by planning and execution so that mixed-version plans are rejected with the + * same semantics and error messages regardless of where validation first occurs. + */ +public final class JoinKeyTypeValidator { + private static final String JOIN_KEY_ERROR = + "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"; + private static final String ASOF_MATCH_KEY_ERROR = + "Raw VARIANT values do not support ASOF JOIN match keys; extract a typed path with variantGet first"; + + private JoinKeyTypeValidator() { + } + + /** + * Validates equality/hash keys and, for ASOF joins, ordering keys. + * + *

The right schema can be absent for legacy SEMI and ANTI join plans whose output schema contains only left + * columns. In that case, the caller can validate only the left keys. + */ + public static void validate(JoinNode joinNode, DataSchema leftSchema, @Nullable DataSchema rightSchema) { + validateEqualityAndHashing(joinNode.getLeftKeys(), leftSchema); + if (rightSchema == null) { + return; + } + validateEqualityAndHashing(joinNode.getRightKeys(), rightSchema); + if (joinNode.getJoinStrategy() == JoinNode.JoinStrategy.ASOF) { + validateAsofMatchKeys(joinNode, leftSchema, rightSchema); + } + } + + private static void validateEqualityAndHashing(List keys, DataSchema schema) { + for (int key : keys) { + DataSchema.ColumnDataType dataType = schema.getColumnDataType(key); + Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), JOIN_KEY_ERROR); + } + } + + private static void validateAsofMatchKeys(JoinNode joinNode, DataSchema leftSchema, DataSchema rightSchema) { + RexExpression.FunctionCall matchCondition = (RexExpression.FunctionCall) joinNode.getMatchCondition(); + List matchKeys = matchCondition.getFunctionOperands(); + int leftMatchKey = ((RexExpression.InputRef) matchKeys.get(0)).getIndex(); + int rightMatchKey = ((RexExpression.InputRef) matchKeys.get(1)).getIndex() - leftSchema.size(); + Preconditions.checkArgument(leftSchema.getColumnDataType(leftMatchKey).supportsOrdering() + && rightSchema.getColumnDataType(rightMatchKey).supportsOrdering(), ASOF_MATCH_KEY_ERROR); + } +} diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java index ff6405ce5cd4..955c78ede633 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java @@ -35,9 +35,9 @@ /** * Rejects operations that would otherwise assign physical byte ordering, equality, or hashing semantics to a raw - * VARIANT value. + * VARIANT value. The visitor has no mutable state and is thread-safe, so callers may share {@link #INSTANCE}. */ -public class VariantTypeValidationVisitor extends PlanNodeVisitor.DepthFirstVisitor { +public final class VariantTypeValidationVisitor extends PlanNodeVisitor.DepthFirstVisitor { public static final VariantTypeValidationVisitor INSTANCE = new VariantTypeValidationVisitor(); private VariantTypeValidationVisitor() { @@ -123,27 +123,10 @@ public Void visitJoin(JoinNode node, Void context) { * {@code HashJoinOperator}. */ public static void validateJoinInputs(JoinNode node, DataSchema leftSchema, DataSchema rightSchema) { - for (int leftKey : node.getLeftKeys()) { - DataSchema.ColumnDataType dataType = leftSchema.getColumnDataType(leftKey); - if (!dataType.supportsEquality() || !dataType.supportsHashing()) { - throw unsupported("JOIN keys"); - } - } - for (int rightKey : node.getRightKeys()) { - DataSchema.ColumnDataType dataType = rightSchema.getColumnDataType(rightKey); - if (!dataType.supportsEquality() || !dataType.supportsHashing()) { - throw unsupported("JOIN keys"); - } - } - if (node.getJoinStrategy() == JoinNode.JoinStrategy.ASOF) { - RexExpression.FunctionCall matchCondition = (RexExpression.FunctionCall) node.getMatchCondition(); - List matchKeys = matchCondition.getFunctionOperands(); - int leftMatchKey = ((RexExpression.InputRef) matchKeys.get(0)).getIndex(); - int rightMatchKey = ((RexExpression.InputRef) matchKeys.get(1)).getIndex() - leftSchema.size(); - if (!leftSchema.getColumnDataType(leftMatchKey).supportsOrdering() - || !rightSchema.getColumnDataType(rightMatchKey).supportsOrdering()) { - throw unsupported("ASOF JOIN MATCH_CONDITION"); - } + try { + JoinKeyTypeValidator.validate(node, leftSchema, rightSchema); + } catch (IllegalArgumentException e) { + throw new QueryException(QueryErrorCode.QUERY_PLANNING, e.getMessage(), e); } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java index 321694d718e5..7c390843f00f 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java @@ -95,7 +95,7 @@ public void testRejectsRawVariantAsofMatchKeys() { new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.LONG}); QueryException exception = Assert.expectThrows(QueryException.class, () -> asofJoin(leftVariant, rightTyped).visit(VariantTypeValidationVisitor.INSTANCE, null)); - Assert.assertTrue(exception.getMessage().contains("ASOF JOIN MATCH_CONDITION")); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); DataSchema leftTyped = new DataSchema(new String[]{"key", "match"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.LONG}); @@ -103,7 +103,7 @@ public void testRejectsRawVariantAsofMatchKeys() { new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.VARIANT}); exception = Assert.expectThrows(QueryException.class, () -> asofJoin(leftTyped, rightVariant).visit(VariantTypeValidationVisitor.INSTANCE, null)); - Assert.assertTrue(exception.getMessage().contains("ASOF JOIN MATCH_CONDITION")); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); } @Test diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java index 55cf3a6ab575..a7ec713a7e47 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java @@ -33,6 +33,7 @@ import org.apache.pinot.query.planner.partitioning.KeySelectorFactory; import org.apache.pinot.query.planner.plannode.JoinNode; import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.validation.JoinKeyTypeValidator; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.operator.join.DoubleLookupTable; import org.apache.pinot.query.runtime.operator.join.FloatLookupTable; @@ -91,7 +92,7 @@ private HashJoinOperator(OpChainExecutionContext context, MultiStageOperator lef List leftKeys = node.getLeftKeys(); Preconditions.checkState(!leftKeys.isEmpty(), "Hash join operator requires join keys"); Preconditions.checkArgument(!rightSchemaRequired || rightSchema != null, "Right input schema must not be null"); - validateVariantJoinKeys(leftKeys, node.getRightKeys(), leftSchema, rightSchema); + JoinKeyTypeValidator.validate(node, leftSchema, rightSchema); _leftKeySelector = KeySelectorFactory.getKeySelector(leftKeys); _rightKeySelector = KeySelectorFactory.getKeySelector(node.getRightKeys()); _rightTable = createLookupTable(leftKeys, leftSchema); @@ -126,7 +127,7 @@ private HashJoinOperator(OpChainExecutionContext context, MultiStageOperator lef List leftKeys = node.getLeftKeys(); Preconditions.checkState(!leftKeys.isEmpty(), "Hash join operator requires join keys"); Preconditions.checkArgument(!rightSchemaRequired || rightSchema != null, "Right input schema must not be null"); - validateVariantJoinKeys(leftKeys, node.getRightKeys(), leftSchema, rightSchema); + JoinKeyTypeValidator.validate(node, leftSchema, rightSchema); _leftKeySelector = KeySelectorFactory.getKeySelector(leftKeys); _rightKeySelector = KeySelectorFactory.getKeySelector(node.getRightKeys()); _rightTable = createLookupTable(leftKeys, leftSchema); @@ -169,22 +170,6 @@ private static LookupTable createLookupTable(List joinKeys, DataSchema } } - private static void validateVariantJoinKeys(List leftKeys, List rightKeys, DataSchema leftSchema, - @Nullable DataSchema rightSchema) { - for (int leftKey : leftKeys) { - DataSchema.ColumnDataType dataType = leftSchema.getColumnDataType(leftKey); - Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), - "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); - } - if (rightSchema != null) { - for (int rightKey : rightKeys) { - DataSchema.ColumnDataType dataType = rightSchema.getColumnDataType(rightKey); - Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), - "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); - } - } - } - @Override public String toExplainString() { return EXPLAIN_NAME; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java index 329a2176795f..7634d2500d6f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/factory/DefaultJoinOperatorFactory.java @@ -18,12 +18,11 @@ */ package org.apache.pinot.query.runtime.operator.factory; -import com.google.common.base.Preconditions; import org.apache.pinot.common.utils.DataSchema; -import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.EnrichedJoinNode; import org.apache.pinot.query.planner.plannode.JoinNode; import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.validation.JoinKeyTypeValidator; import org.apache.pinot.query.runtime.operator.AsofJoinOperator; import org.apache.pinot.query.runtime.operator.HashJoinOperator; import org.apache.pinot.query.runtime.operator.LookupJoinOperator; @@ -40,7 +39,7 @@ public MultiStageOperator createJoinOperator(OpChainExecutionContext context, Mu JoinNode.JoinStrategy joinStrategy = joinNode.getJoinStrategy(); DataSchema leftSchema = leftPlanNode.getDataSchema(); DataSchema rightSchema = rightPlanNode.getDataSchema(); - validateJoinKeys(joinNode, leftSchema, rightSchema); + JoinKeyTypeValidator.validate(joinNode, leftSchema, rightSchema); switch (joinStrategy) { case HASH: if (joinNode.getLeftKeys().isEmpty()) { @@ -58,28 +57,6 @@ public MultiStageOperator createJoinOperator(OpChainExecutionContext context, Mu } } - private static void validateJoinKeys(JoinNode joinNode, DataSchema leftSchema, DataSchema rightSchema) { - for (int leftKey : joinNode.getLeftKeys()) { - DataSchema.ColumnDataType dataType = leftSchema.getColumnDataType(leftKey); - Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), - "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); - } - for (int rightKey : joinNode.getRightKeys()) { - DataSchema.ColumnDataType dataType = rightSchema.getColumnDataType(rightKey); - Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), - "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"); - } - if (joinNode.getJoinStrategy() == JoinNode.JoinStrategy.ASOF) { - RexExpression.FunctionCall matchCondition = (RexExpression.FunctionCall) joinNode.getMatchCondition(); - int leftMatchKey = ((RexExpression.InputRef) matchCondition.getFunctionOperands().get(0)).getIndex(); - int rightMatchKey = - ((RexExpression.InputRef) matchCondition.getFunctionOperands().get(1)).getIndex() - leftSchema.size(); - Preconditions.checkArgument(leftSchema.getColumnDataType(leftMatchKey).supportsOrdering() - && rightSchema.getColumnDataType(rightMatchKey).supportsOrdering(), - "Raw VARIANT values do not support ASOF JOIN match keys; extract a typed path with variantGet first"); - } - } - /// Enriched joins have been removed. This method is retained only for backward compatibility of the /// [JoinOperatorFactory] interface and always throws. A current broker never produces an /// [EnrichedJoinNode], so this is only reachable if a plan from an older-version broker is executed. diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 431accc81659..b3ad5286006a 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java @@ -137,17 +137,35 @@ public void testRejectsRawVariantWindowKeys() { @Test public void testAllowsWindowKeysOverTypedVariantExtraction() { - DataSchema inputSchema = new DataSchema(new String[]{"typedPayload"}, new ColumnDataType[]{STRING}); + DataSchema inputSchema = + new DataSchema(new String[]{"eventType", "eventId"}, new ColumnDataType[]{STRING, STRING}); DataSchema resultSchema = - new DataSchema(new String[]{"typedPayload", "count"}, new ColumnDataType[]{STRING, LONG}); - MultiStageOperator input = new BlockListMultiStageOperator.Builder(inputSchema).buildWithEos(); - List aggCalls = List.of(getCount(new RexExpression.InputRef(0))); + new DataSchema(new String[]{"eventType", "eventId", "count"}, new ColumnDataType[]{STRING, STRING, LONG}); + MultiStageOperator input = new BlockListMultiStageOperator.Builder(inputSchema) + .addBlock(new Object[]{"checkout", "evt-001"}) + .addBlock(new Object[]{"view", "evt-002"}) + .addBlock(new Object[]{"checkout", "evt-003"}) + // Both an encoded Variant null and a SQL null extract to SQL null before reaching the window operator. + .addBlock(new Object[]{null, "evt-004"}) + .addBlock(new Object[]{null, "evt-005"}) + .buildWithEos(); + List aggCalls = List.of(getCountStar()); WindowAggregateOperator operator = - getOperator(inputSchema, resultSchema, List.of(0), List.of(new RelFieldCollation(0)), aggCalls, ROWS, + getOperator(inputSchema, resultSchema, List.of(0), List.of(), aggCalls, RANGE, Integer.MIN_VALUE, Integer.MAX_VALUE, input); - assertTrue(operator.nextBlock().isSuccess()); + List resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + Map> expectedRows = new HashMap<>(); + expectedRows.put("checkout", List.of( + new Object[]{"checkout", "evt-001", 2L}, + new Object[]{"checkout", "evt-003", 2L})); + expectedRows.put("view", List.of(new Object[]{"view", "evt-002", 1L})); + expectedRows.put(null, List.of( + new Object[]{null, "evt-004", 2L}, + new Object[]{null, "evt-005", 2L})); + verifyResultRows(resultRows, List.of(0), expectedRows); + assertTrue(operator.nextBlock().isSuccess(), "Second block is EOS (done processing)"); } @Test @@ -3563,6 +3581,10 @@ private static RexExpression.FunctionCall getCount(RexExpression arg) { return new RexExpression.FunctionCall(ColumnDataType.LONG, SqlKind.COUNT.name(), List.of(arg)); } + private static RexExpression.FunctionCall getCountStar() { + return new RexExpression.FunctionCall(ColumnDataType.LONG, SqlKind.COUNT.name(), List.of()); + } + private static RexExpression.FunctionCall getMin(RexExpression arg) { return new RexExpression.FunctionCall(ColumnDataType.INT, SqlKind.MIN.name(), List.of(arg)); } diff --git a/pinot-spi/VARIANT_DESIGN.md b/pinot-spi/VARIANT_DESIGN.md index 22eab9a70d45..512934cd3819 100644 --- a/pinot-spi/VARIANT_DESIGN.md +++ b/pinot-spi/VARIANT_DESIGN.md @@ -281,6 +281,10 @@ preserve the distinction among SQL null, Variant null, and the Variant string `" There is no type negotiation or safe downgrade to `BYTES`. An old node can read existing non-Variant traffic, but it cannot plan or execute an active Variant query. +The rolling-upgrade compatibility suite exercises both directions: a new broker proves +legacy plans still work against old servers and observes the deterministic +`UNRECOGNIZED` error for a leaf plan containing `VARIANT`, while an old broker +continues sending legacy plans to upgraded servers. ## 9. Deployment and rollback @@ -308,12 +312,14 @@ to activate. The review dependency order is: 2. schema/table/index validation and Parquet ingestion; 3. function semantics plus single-stage and multi-stage guards; 4. DDL, response, Java, JDBC, gRPC, JSON, and Arrow propagation; and -5. quickstart, integration tests, compatibility tests, and benchmarks. +5. quickstart, integration tests, and compatibility tests. The initial implementation is presented as one end-to-end draft so reviewers can -evaluate one activation contract and run one acceptance test. It is not merge-ready -until SPI/wire, Parquet, both query engines, and client/response owners approve their -areas. If maintainers prefer independent rollback units, the five groups above form a +evaluate one activation contract and run one acceptance test. Performance benchmarks +are intentionally kept in a follow-up change so they do not couple the activation +contract to optional tooling. This change is not merge-ready until SPI/wire, Parquet, +both query engines, and client/response owners approve their areas. If maintainers +prefer additional independent rollback units, the five groups above form a dependency-ordered PR stack; every intermediate PR must compile, keep Variant unactivatable until the safety guards land, and preserve the permanent wire allocation. From 0068ded7a21e536396cf7d0878fdce8aa708dc58 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 31 Jul 2026 17:04:08 -0700 Subject: [PATCH 4/8] Harden VARIANT review edge cases Raw VARIANT projection is only lossless with query null handling: otherwise the reserved empty-byte SQL-null placeholder cannot preserve null semantics. Enforce this contract at the multi-stage root plan, single-stage result block, and broker normal, empty, and materialized-view split response paths. Preserve legacy ResultSet.getString() behavior for established types while mapping only VARIANT SQL null to Java null across JSON, gRPC, and JDBC clients. Make Parquet VARIANT validation honor selected fields so an unselected nested VARIANT does not reject an unrelated projection, while selected and extract-all paths remain fail-closed. Reject raw VARIANT operands for IS DISTINCT FROM and IS NOT DISTINCT FROM through shared operand type resolution, add focused and end-to-end coverage, and normalize the touched documentation to the repository style. --- .../BaseSingleStageBrokerRequestHandler.java | 21 ++ ...seSingleStageBrokerRequestHandlerTest.java | 101 ++++++++ .../pinot/client/AbstractResultSet.java | 10 +- .../org/apache/pinot/client/ResultSet.java | 6 +- .../pinot/client/ResultTableResultSet.java | 2 + .../pinot/client/grpc/GrpcResultSet.java | 13 +- .../pinot/client/grpc/GrpcResultSetTest.java | 85 +++++++ .../client/grpc/PinotGrpcResultSetTest.java | 6 +- .../evaluator/InbuiltFunctionEvaluator.java | 10 +- .../function/scalar/VariantFunctions.java | 4 +- .../pinot/common/utils/VariantUtils.java | 220 +++++++----------- .../function/scalar/VariantFunctionsTest.java | 4 +- .../pinot/common/utils/VariantUtilsTest.java | 12 + .../java/org/apache/pinot/compat/BaseOp.java | 10 +- .../apache/pinot/compat/FileContainsOp.java | 6 +- .../java/org/apache/pinot/compat/QueryOp.java | 20 +- .../org/apache/pinot/compat/BaseOpTest.java | 4 +- .../blocks/results/SelectionResultsBlock.java | 5 + .../BaseVariantTransformFunction.java | 42 ++-- .../IsVariantNullTransformFunction.java | 10 +- .../ParseJsonToVariantTransformFunction.java | 12 +- .../VariantExistsTransformFunction.java | 12 +- .../function/VariantGetTransformFunction.java | 16 +- .../VariantTypeOfTransformFunction.java | 10 +- .../results/SelectionResultsBlockTest.java | 54 +++++ .../tests/custom/VariantTypeTest.java | 30 ++- .../parquet/ParquetNativeRecordExtractor.java | 5 +- .../ParquetNativeRecordExtractorConfig.java | 12 +- .../inputformat/parquet/ParquetUtils.java | 8 +- .../parquet/ParquetVariantConverter.java | 19 ++ .../ParquetVariantRecordReaderTest.java | 35 ++- .../physical/PinotDispatchPlanner.java | 4 + .../validation/JoinKeyTypeValidator.java | 20 +- .../VariantTypeValidationVisitor.java | 51 ++-- .../VariantTypeValidationVisitorTest.java | 10 + .../runtime/operator/HashJoinOperator.java | 24 +- .../operator/operands/FunctionOperand.java | 12 +- .../operands/LiteralParseJsonOperand.java | 12 +- .../operands/TransformOperandFactory.java | 28 +++ .../operator/operands/VariantOperand.java | 12 +- .../operator/operands/FilterOperandTest.java | 11 + pinot-spi/VARIANT_DESIGN.md | 2 + .../apache/pinot/spi/utils/PinotDataType.java | 4 +- .../pinot/spi/utils/VariantEnvelope.java | 122 ++++------ .../apache/pinot/tools/VariantQuickStart.java | 4 +- 45 files changed, 695 insertions(+), 425 deletions(-) create mode 100644 pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/grpc/GrpcResultSetTest.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlockTest.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index b61abf35f083..7acad35391c4 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -80,6 +80,7 @@ import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.DatabaseUtils; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.request.QueryFingerprintUtils; import org.apache.pinot.common.utils.request.RequestUtils; @@ -1049,6 +1050,7 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn // server returns STRING as default dataType for all columns in (some) scenarios where no rows are returned // this is an attempt to return more faithful information based on other sources fillEmptyResponseSchema(pinotQuery, brokerResponse, schema, database, query); + validateRawVariantResult(brokerResponse, serverPinotQuery, requestContext, tableName); // Set total query processing time long totalTimeMs = System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(); @@ -1524,6 +1526,7 @@ private BrokerResponseNative getEmptyBrokerOnlyResponse(PinotQuery pinotQuery, P LOGGER.warn("Caught exception while building empty response for request {}: {}, {}", requestContext.getRequestId(), query, e.getMessage()); } + validateRawVariantResult(brokerResponse, serverPinotQuery, requestContext, tableName); brokerResponse.setTablesQueried(Set.of(TableNameBuilder.extractRawTableName(tableName))); brokerResponse.setTimeUsedMs(System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis()); _queryLogger.logQueryCompleted(new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, @@ -1540,6 +1543,23 @@ private void fillEmptyResponseSchema(PinotQuery pinotQuery, BrokerResponse broke } } + private void validateRawVariantResult(BrokerResponseNative brokerResponse, PinotQuery serverPinotQuery, + RequestContext requestContext, String tableName) { + ResultTable resultTable = brokerResponse.getResultTable(); + Map queryOptions = serverPinotQuery.getQueryOptions(); + boolean nullHandlingEnabled = queryOptions != null && QueryOptionsUtils.isNullHandlingEnabled(queryOptions); + if (resultTable == null || !VariantUtils.requiresNullHandlingForRawVariantResult(resultTable.getDataSchema(), + nullHandlingEnabled)) { + return; + } + requestContext.setErrorCode(QueryErrorCode.QUERY_VALIDATION); + _brokerMetrics.addMeteredTableValue(TableNameBuilder.extractRawTableName(tableName), + BrokerMeter.QUERY_VALIDATION_EXCEPTIONS, 1); + brokerResponse.setResultTable(null); + brokerResponse.addException(new QueryProcessingException(QueryErrorCode.QUERY_VALIDATION, + VariantUtils.RAW_VARIANT_REQUIRES_NULL_HANDLING_ERROR)); + } + private void handleTimestampIndexOverride(PinotQuery pinotQuery, @Nullable TableConfig tableConfig) { if (tableConfig == null || tableConfig.getFieldConfigList() == null) { return; @@ -2478,6 +2498,7 @@ private BrokerResponseNative tryExecuteMaterializedViewSplit(long requestId, Str } viewSplitResponse.setNumSegmentsPrunedByBroker(numPrunedSegmentsTotal); fillEmptyResponseSchema(brokerRequest.getPinotQuery(), viewSplitResponse, schema, database, query); + validateRawVariantResult(viewSplitResponse, serverPinotQuery, requestContext, tableName); long totalTimeMs = System.currentTimeMillis() - requestContext.getRequestArrivalTimeMillis(); viewSplitResponse.setTimeUsedMs(totalTimeMs); augmentStatistics(requestContext, viewSplitResponse); diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java index ae661c0107a4..fae6b69b46a9 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java @@ -46,6 +46,8 @@ import org.apache.pinot.common.request.PinotQuery; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.core.routing.RoutingTable; import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TableRouteInfo; @@ -73,6 +75,7 @@ import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.exception.BadQueryRequestException; +import org.apache.pinot.spi.exception.QueryErrorCode; import org.apache.pinot.spi.trace.LoggerConstants; import org.apache.pinot.spi.trace.RequestContext; import org.apache.pinot.spi.utils.CommonConstants; @@ -1603,6 +1606,104 @@ protected BrokerResponseNative processMaterializedViewSplitBrokerRequest(long re "Response must not report a materializedViewQueried when the cascade guard fired"); } + /// A split MV response returns before the ordinary single-stage response finalization. Verify + /// that this early-return path applies the same raw-VARIANT/null-handling contract, including + /// when both split branches are empty and their merged schema is repaired at the broker. + @Test + public void testMaterializedViewSplitEmptyRawVariantRequiresNullHandling() + throws Exception { + String baseOfflineTable = "baseTable_OFFLINE"; + String materializedViewOfflineTable = "mv_baseTable_OFFLINE"; + String baseRawTable = "baseTable"; + String materializedViewRawTable = "mv_baseTable"; + + String userSql = "SELECT payload FROM baseTable LIMIT 10"; + PinotQuery materializedViewServerQuery = CalciteSqlParser.compileToPinotQuery( + "SELECT payload FROM mv_baseTable_OFFLINE LIMIT 10"); + MaterializedViewRewritePlan plan = new MaterializedViewRewritePlan( + materializedViewOfflineTable, MatchType.EXACT, ExecutionMode.SPLIT_REWRITE, + materializedViewServerQuery, 1.0); + + Schema baseSchema = new Schema.SchemaBuilder() + .setSchemaName(baseRawTable) + .addSingleValueDimension("payload", DataType.VARIANT) + .build(); + Schema materializedViewSchema = new Schema.SchemaBuilder() + .setSchemaName(materializedViewRawTable) + .addSingleValueDimension("payload", DataType.VARIANT) + .build(); + + // Empty server responses can report STRING regardless of the selected column's actual type. The broker must + // repair this to VARIANT from the query/schema before enforcing the null-handling contract. + DataSchema emptySplitDataSchema = new DataSchema(new String[]{"payload"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + BrokerResponseNative splitResponse = new BrokerResponseNative(); + splitResponse.setResultTable(new ResultTable(emptySplitDataSchema, List.of())); + + MaterializedViewHandler materializedViewHandler = mock(MaterializedViewHandler.class); + when(materializedViewHandler.compile(any(MaterializedViewCompileContext.class))) + .thenReturn(MaterializedViewContext.forSplitRewrite( + plan, materializedViewServerQuery, materializedViewOfflineTable, materializedViewSchema)); + when(materializedViewHandler.executeSplit(any(MaterializedViewSplitExecutionContext.class))) + .thenReturn(splitResponse); + + TableCache tableCache = mock(TableCache.class); + when(tableCache.getActualTableName(baseRawTable)).thenReturn(baseRawTable); + when(tableCache.getSchema(baseRawTable)).thenReturn(baseSchema); + when(tableCache.getSchema(materializedViewRawTable)).thenReturn(materializedViewSchema); + when(tableCache.getColumnNameMap(anyString())).thenReturn(Map.of("payload", "payload")); + TableConfig tableConfig = mock(TableConfig.class); + when(tableConfig.getTenantConfig()).thenReturn(new TenantConfig("t_BROKER", "t_SERVER", null)); + when(tableCache.getTableConfig(baseOfflineTable)).thenReturn(tableConfig); + + BrokerRoutingManager routingManager = mock(BrokerRoutingManager.class); + when(routingManager.routingExists(baseOfflineTable)).thenReturn(true); + when(routingManager.getQueryTimeoutMs(anyString())).thenReturn(10000L); + RoutingTable routingTable = mock(RoutingTable.class); + when(routingTable.getServerInstanceToSegmentsMap()).thenReturn( + Map.of(new ServerInstance(new InstanceConfig("server01_9000")), + new SegmentsToQuery(List.of("seg01"), List.of()))); + when(routingManager.getRoutingTable(any(), Mockito.anyLong())).thenReturn(routingTable); + + QueryQuotaManager quotaManager = mock(QueryQuotaManager.class); + when(quotaManager.acquireDatabase(anyString())).thenReturn(true); + when(quotaManager.acquireApplication(anyString())).thenReturn(true); + when(quotaManager.acquire(anyString())).thenReturn(true); + + BrokerMetrics.register(mock(BrokerMetrics.class)); + PinotConfiguration config = new PinotConfiguration(); + BrokerQueryEventListenerFactory.init(config); + + BaseSingleStageBrokerRequestHandler handler = + new BaseSingleStageBrokerRequestHandler(config, "broker1", new BrokerRequestIdGenerator(), + routingManager, new AllowAllAccessControlFactory(), quotaManager, tableCache, + ThreadAccountantUtils.getNoOpAccountant(), null, materializedViewHandler) { + @Override + public void start() { + } + + @Override + public void shutDown() { + } + + @Override + protected BrokerResponseNative processBrokerRequest(long requestId, + BrokerRequest originalBrokerRequest, BrokerRequest serverBrokerRequest, + TableRouteInfo route, long timeoutMs, ServerStats serverStats, + RequestContext requestContext) { + Assert.fail("Split response must return before the ordinary broker request path"); + return null; + } + }; + + BrokerResponseNative response = (BrokerResponseNative) handler.handleRequest(userSql); + Assert.assertNull(response.getResultTable()); + Assert.assertEquals(response.getExceptionsSize(), 1); + Assert.assertEquals(response.getExceptions().get(0).getErrorCode(), QueryErrorCode.QUERY_VALIDATION.getId()); + Assert.assertTrue(response.getExceptions().get(0).getMessage().contains("null handling")); + verify(materializedViewHandler).executeSplit(any(MaterializedViewSplitExecutionContext.class)); + } + /// Pins the fix for the reviewer-flagged regression: `EXPLAIN PLAN FOR ` against a /// SPLIT-eligible MV must NOT enter `tryExecuteMaterializedViewSplit` (which would dispatch /// dual scatter-gather to base+MV and merge live `DataTable`s). The broker must fall diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/AbstractResultSet.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/AbstractResultSet.java index a67199cbe565..e7825de2d3c8 100644 --- a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/AbstractResultSet.java +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/AbstractResultSet.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.client; +import java.util.Objects; import javax.annotation.Nullable; /// Shared implementation between the different ResultSets. @@ -49,6 +50,7 @@ public double getDouble(int rowIndex) { return getDouble(rowIndex, 0); } + @Nullable @Override public String getString(int rowIndex) { return getString(rowIndex, 0); @@ -56,22 +58,22 @@ public String getString(int rowIndex) { @Override public int getInt(int rowIndex, int columnIndex) { - return Integer.parseInt(getString(rowIndex, columnIndex)); + return Integer.parseInt(Objects.requireNonNull(getString(rowIndex, columnIndex))); } @Override public long getLong(int rowIndex, int columnIndex) { - return Long.parseLong(getString(rowIndex, columnIndex)); + return Long.parseLong(Objects.requireNonNull(getString(rowIndex, columnIndex))); } @Override public float getFloat(int rowIndex, int columnIndex) { - return Float.parseFloat(getString(rowIndex, columnIndex)); + return Float.parseFloat(Objects.requireNonNull(getString(rowIndex, columnIndex))); } @Override public double getDouble(int rowIndex, int columnIndex) { - return Double.parseDouble(getString(rowIndex, columnIndex)); + return Double.parseDouble(Objects.requireNonNull(getString(rowIndex, columnIndex))); } @Override diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultSet.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultSet.java index 35af1130be55..dd4269f120cb 100644 --- a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultSet.java +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultSet.java @@ -72,7 +72,8 @@ public interface ResultSet { /// Obtains the String value for the given row. /// /// @param rowIndex The index of the row - /// @return The String value for the given row + /// @return The String value for the given row, or Java `null` when a VARIANT cell is SQL null + @Nullable String getString(int rowIndex); /// Obtains the integer value for the given row and column. @@ -107,7 +108,8 @@ public interface ResultSet { /// /// @param rowIndex The index of the row /// @param columnIndex The index of the column for which to fetch the value - /// @return The String value for the given row and column + /// @return The String value for the given row and column, or Java `null` when a VARIANT cell is SQL null + @Nullable String getString(int rowIndex, int columnIndex); /// Obtains the length of the group key, or 0 if there is no grouping key. diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java index 0e15e411f0df..f696925a0734 100644 --- a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/ResultTableResultSet.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.JsonNode; import java.util.ArrayList; import java.util.List; +import javax.annotation.Nullable; /// ResultSet which contains the ResultTable from the broker response of a sql query. @@ -56,6 +57,7 @@ public String getColumnDataType(int columnIndex) { return _columnDataTypesArray.get(columnIndex).asText(); } + @Nullable @Override public String getString(int rowIndex, int columnIndex) { JsonNode jsonValue = _rowsArray.get(rowIndex).get(columnIndex); diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/grpc/GrpcResultSet.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/grpc/GrpcResultSet.java index 915cfe785209..0d5f73300039 100644 --- a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/grpc/GrpcResultSet.java +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/grpc/GrpcResultSet.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.List; +import javax.annotation.Nullable; import org.apache.pinot.client.AbstractResultSet; import org.apache.pinot.client.TextTable; import org.apache.pinot.common.proto.Broker; @@ -68,9 +69,16 @@ public String getColumnDataType(int columnIndex) { return _columnDataTypesArray.get(columnIndex); } + @Nullable @Override public String getString(int rowIndex, int columnIndex) { - return _currentBatchRows.getRows().get(rowIndex)[columnIndex].toString(); + Object value = _currentBatchRows.getRows().get(rowIndex)[columnIndex]; + if (value != null) { + return value.toString(); + } + // Keep the established getString() contract for existing types. VARIANT alone needs Java null so callers can + // distinguish SQL null from the encoded Variant null, whose canonical JSON representation is the string "null". + return "VARIANT".equals(getColumnDataType(columnIndex)) ? null : "null"; } public List getAllColumns() { @@ -114,7 +122,8 @@ public String toString() { String[] columnValues = new String[numColumns]; for (int c = 0; c < numColumns; c++) { try { - columnValues[c] = getString(r, c); + String value = getString(r, c); + columnValues[c] = value != null ? value : "null"; } catch (Exception e) { columnNames[c] = "ERROR"; } diff --git a/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/grpc/GrpcResultSetTest.java b/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/grpc/GrpcResultSetTest.java new file mode 100644 index 000000000000..61de024bf4db --- /dev/null +++ b/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/grpc/GrpcResultSetTest.java @@ -0,0 +1,85 @@ +/** + * 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.client.grpc; + +import com.google.protobuf.ByteString; +import java.io.IOException; +import java.util.List; +import org.apache.pinot.common.proto.Broker; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.response.encoder.JsonResponseEncoder; +import org.apache.pinot.common.utils.DataSchema; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/// Verifies the Java client result contract for data blocks emitted by the gRPC broker endpoint. +public class GrpcResultSetTest { + + @Test + public void testEstablishedTypeSqlNullRetainsStringContract() + throws IOException { + DataSchema schema = new DataSchema(new String[]{"message"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING}); + List rows = List.of(new Object[]{null}); + byte[] payload = new JsonResponseEncoder().encodeResultTable( + new ResultTable(schema, rows), 0, rows.size()); + Broker.BrokerResponse response = Broker.BrokerResponse.newBuilder() + .setPayload(ByteString.copyFrom(payload)) + .putMetadata("rowSize", Integer.toString(rows.size())) + .putMetadata("compression", "NONE") + .putMetadata("encoding", "JSON") + .build(); + + GrpcResultSet resultSet = new GrpcResultSet(schema, response); + assertEquals(resultSet.getString(0, 0), "null", + "Existing column types must retain the legacy getString() representation for SQL null"); + } + + @Test + public void testVariantGetStringPreservesVariantNullAndSqlNull() + throws IOException { + DataSchema schema = new DataSchema(new String[]{"payload"}, + new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.VARIANT}); + List rows = List.of( + new Object[]{"{\"answer\":42}"}, + new Object[]{"null"}, + new Object[]{"\"null\""}, + new Object[]{null} + ); + byte[] payload = new JsonResponseEncoder().encodeResultTable( + new ResultTable(schema, rows), 0, rows.size()); + Broker.BrokerResponse response = Broker.BrokerResponse.newBuilder() + .setPayload(ByteString.copyFrom(payload)) + .putMetadata("rowSize", Integer.toString(rows.size())) + .putMetadata("compression", "NONE") + .putMetadata("encoding", "JSON") + .build(); + + GrpcResultSet resultSet = new GrpcResultSet(schema, response); + assertEquals(resultSet.getString(0, 0), "{\"answer\":42}"); + assertEquals(resultSet.getString(1, 0), "null", "A Variant null is not SQL null"); + assertEquals(resultSet.getString(2, 0), "\"null\"", "A Variant string containing null is not SQL null"); + assertNull(resultSet.getString(3, 0), "SQL null must map to Java null"); + assertTrue(resultSet.toString().contains("null"), "Text rendering must tolerate SQL null values"); + } +} diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java index 92123145ab6a..630265903502 100644 --- a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java +++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java @@ -35,10 +35,8 @@ import static org.testng.Assert.assertTrue; -/** - * Verifies the JDBC result contract over the same metadata, schema, and data block sequence emitted by the gRPC - * broker endpoint. - */ +/// Verifies the JDBC result contract over the same metadata, schema, and data block sequence emitted by the gRPC +/// broker endpoint. public class PinotGrpcResultSetTest { @Test diff --git a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java index 62d0ebf04798..a9374c1cf147 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java @@ -142,12 +142,10 @@ public String toString() { return _functionExpression; } - /** - * Planned ingestion evaluator for Variant scalar functions with literal path and target-type operands. - * - *

Each node compiles its literals once and owns a reusable cursor result. Like the enclosing evaluator, instances - * are intended to be confined to the record-transformer thread and are not thread-safe. - */ + /// Planned ingestion evaluator for Variant scalar functions with literal path and target-type operands. + /// + ///

Each node compiles its literals once and owns a reusable cursor result. Like the enclosing evaluator, instances + /// are intended to be confined to the record-transformer thread and are not thread-safe. private static class VariantExecutionNode implements ExecutableNode { private static final String VARIANT_GET = "variantget"; private static final String TRY_VARIANT_GET = "tryvariantget"; diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java index a5d968319270..668fb0cc402d 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java @@ -23,9 +23,7 @@ import org.apache.pinot.spi.annotations.ScalarFunction; -/** - * Scalar functions for the Pinot {@code VARIANT} logical type. - */ +/// Scalar functions for the Pinot {@code VARIANT} logical type. public final class VariantFunctions { private VariantFunctions() { } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java index 0239a9f13bff..9b7d66e47b69 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java @@ -52,15 +52,16 @@ import org.apache.pinot.spi.utils.VariantEnvelope; -/** - * Query-side operations for Pinot {@code VARIANT} values. - * - *

The utility navigates the Parquet Variant binary representation directly. It never materializes a JSON tree. - * Instances are not required, and stateless convenience methods are thread-safe. Overloads that accept a - * caller-provided {@link ReusableResult} require that result to be thread-confined and not shared by concurrent calls. - * An empty byte array is Pinot's SQL-null placeholder and is never decoded as an envelope. - */ +/// Query-side operations for Pinot {@code VARIANT} values. +/// +///

The utility navigates the Parquet Variant binary representation directly. It never materializes a JSON tree. +/// Instances are not required, and stateless convenience methods are thread-safe. Overloads that accept a +/// caller-provided {@link ReusableResult} require that result to be thread-confined and not shared by concurrent calls. +/// An empty byte array is Pinot's SQL-null placeholder and is never decoded as an envelope. public final class VariantUtils { + public static final String RAW_VARIANT_REQUIRES_NULL_HANDLING_ERROR = + "Raw VARIANT projection requires query null handling to be enabled; set enableNullHandling=true"; + private static final JsonFactory JSON_FACTORY = new JsonFactory(); private static final BigDecimal MIN_INT_DECIMAL = BigDecimal.valueOf(Integer.MIN_VALUE); private static final BigDecimal MAX_INT_DECIMAL = BigDecimal.valueOf(Integer.MAX_VALUE); @@ -107,9 +108,22 @@ public final class VariantUtils { private VariantUtils() { } - /** - * Statically supported result types for {@code variantGet} and {@code tryVariantGet}. - */ + /// Returns whether a final result containing raw VARIANT values requires query null handling. Without a null bitmap, + /// Pinot's reserved empty-byte SQL-null placeholder cannot be distinguished from a logical Variant value. + public static boolean requiresNullHandlingForRawVariantResult(DataSchema resultSchema, + boolean nullHandlingEnabled) { + if (nullHandlingEnabled) { + return false; + } + for (DataSchema.ColumnDataType dataType : resultSchema.getColumnDataTypes()) { + if (dataType == DataSchema.ColumnDataType.VARIANT) { + return true; + } + } + return false; + } + + /// Statically supported result types for {@code variantGet} and {@code tryVariantGet}. public enum ResultType { BOOLEAN(DataType.BOOLEAN, SqlTypeName.BOOLEAN), INT(DataType.INT, SqlTypeName.INTEGER), @@ -141,10 +155,8 @@ public SqlTypeName getSqlTypeName() { } } - /** - * An immutable, pre-parsed Variant path. The v1 grammar supports {@code $}, dot-separated object fields, and - * non-negative array subscripts. - */ + /// An immutable, pre-parsed Variant path. The v1 grammar supports {@code $}, dot-separated object fields, and + /// non-negative array subscripts. public static final class VariantPath { private final PathElement[] _elements; @@ -153,15 +165,13 @@ private VariantPath(PathElement[] elements) { } } - /** - * Reusable, unboxed destination for vectorized Variant extraction. - * - *

Only the getter corresponding to the requested {@link ResultType} is defined after a successful extraction. - * The instance is mutable and not thread-safe; callers should retain one per transform-function instance. Every - * extraction may replace its state. Each successful byte-valued extraction installs a newly materialized array. - * Values returned as {@code byte[]} or as a {@link ByteArray} may be retained after this result is reused, but they - * are read-only by contract and must be copied before mutation. - */ + /// Reusable, unboxed destination for vectorized Variant extraction. + /// + ///

Only the getter corresponding to the requested {@link ResultType} is defined after a successful extraction. + /// The instance is mutable and not thread-safe; callers should retain one per transform-function instance. Every + /// extraction may replace its state. Each successful byte-valued extraction installs a newly materialized array. + /// Values returned as {@code byte[]} or as a {@link ByteArray} may be retained after this result is reused, but they + /// are read-only by contract and must be copied before mutation. public static final class ReusableResult { private final Cursor _cursor = new Cursor(); private int _intValue; @@ -196,12 +206,10 @@ public String getStringValue() { return _stringValue; } - /** - * Returns the extracted BYTES, VARIANT, or direct 16-byte UUID representation. - * - *

The returned array is replaced, but not mutated, by the next byte-valued extraction. It may be retained after - * this result is reused, but must be treated as immutable and copied before mutation. - */ + /// Returns the extracted BYTES, VARIANT, or direct 16-byte UUID representation. + /// + ///

The returned array is replaced, but not mutated, by the next byte-valued extraction. It may be retained after + /// this result is reused, but must be treated as immutable and copied before mutation. public byte[] getBytesValue() { return _bytesValue; } @@ -210,12 +218,10 @@ public UUID getUuidValue() { return UuidUtils.toUUID(_bytesValue); } - /** - * Materializes the extracted value in the external representation used by scalar functions and ingestion. - * - *

For BYTES and VARIANT, the returned {@code byte[]} may be retained after this result is reused. It must be - * treated as immutable and copied before mutation. - */ + /// Materializes the extracted value in the external representation used by scalar functions and ingestion. + /// + ///

For BYTES and VARIANT, the returned {@code byte[]} may be retained after this result is reused. It must be + /// treated as immutable and copied before mutation. public Object getExternalValue(ResultType resultType) { switch (resultType) { case BOOLEAN: @@ -245,14 +251,12 @@ public Object getExternalValue(ResultType resultType) { } } - /** - * Materializes the extracted value in {@link DataSchema}'s internal representation. - * - *

TIMESTAMP remains epoch milliseconds and UUID wraps the directly copied 16-byte value, avoiding an - * external-object round trip in the multi-stage engine. For BYTES, UUID, and VARIANT, the returned - * {@link ByteArray} wraps a newly materialized array that may be retained after this result is reused. Neither the - * wrapper nor its array may be mutated; callers must copy the array before mutation. - */ + /// Materializes the extracted value in {@link DataSchema}'s internal representation. + /// + ///

TIMESTAMP remains epoch milliseconds and UUID wraps the directly copied 16-byte value, avoiding an + /// external-object round trip in the multi-stage engine. For BYTES, UUID, and VARIANT, the returned + /// {@link ByteArray} wraps a newly materialized array that may be retained after this result is reused. Neither the + /// wrapper nor its array may be mutated; callers must copy the array before mutation. public Object getInternalValue(ResultType resultType) { switch (resultType) { case BOOLEAN: @@ -281,9 +285,7 @@ public Object getInternalValue(ResultType resultType) { } } - /** - * Parses a target type literal once for reuse by a transform function. - */ + /// Parses a target type literal once for reuse by a transform function. public static ResultType parseResultType(String targetType) { if (targetType == null) { throw new IllegalArgumentException("Variant target type must not be null"); @@ -295,9 +297,7 @@ public static ResultType parseResultType(String targetType) { } } - /** - * Compiles a v1 Variant path. - */ + /// Compiles a v1 Variant path. public static VariantPath compilePath(String path) { if (path == null || path.isEmpty() || path.charAt(0) != '$') { throw new IllegalArgumentException("Variant path must start with '$': " + path); @@ -340,39 +340,31 @@ public static VariantPath compilePath(String path) { return new VariantPath(elements.toArray(new PathElement[0])); } - /** - * Extracts a Variant value. A missing path or SQL null returns Java null; a Variant null remains an encoded Variant - * value. - */ + /// Extracts a Variant value. A missing path or SQL null returns Java null; a Variant null remains an encoded Variant + /// value. @Nullable public static byte[] variantGet(@Nullable byte[] envelope, String path) { return (byte[]) variantGet(envelope, compilePath(path), ResultType.VARIANT); } - /** - * Strictly extracts and converts a value. A missing path or SQL null returns Java null. A Variant null remains - * encoded when the target type is {@link ResultType#VARIANT}, and returns Java null for other target types. An - * incompatible non-null value throws. - */ + /// Strictly extracts and converts a value. A missing path or SQL null returns Java null. A Variant null remains + /// encoded when the target type is {@link ResultType#VARIANT}, and returns Java null for other target types. An + /// incompatible non-null value throws. @Nullable public static Object variantGet(@Nullable byte[] envelope, String path, String targetType) { return variantGet(envelope, compilePath(path), parseResultType(targetType)); } - /** - * Strictly extracts using pre-parsed path and type values. - */ + /// Strictly extracts using pre-parsed path and type values. @Nullable public static Object variantGet(@Nullable byte[] envelope, VariantPath path, ResultType targetType) { ReusableResult result = new ReusableResult(); return extractInto(envelope, path, targetType, result) ? result.getExternalValue(targetType) : null; } - /** - * Strictly extracts into a reusable, unboxed result. - * - * @return {@code false} for SQL null, a missing path, or Variant null converted to a non-Variant target - */ + /// Strictly extracts into a reusable, unboxed result. + /// + /// @return {@code false} for SQL null, a missing path, or Variant null converted to a non-Variant target public static boolean extractInto(@Nullable byte[] envelope, VariantPath path, ResultType targetType, ReusableResult result) { Objects.requireNonNull(result, "result must not be null"); @@ -392,17 +384,13 @@ public static boolean extractInto(@Nullable byte[] envelope, VariantPath path, R return true; } - /** - * Tolerant Variant extraction. Malformed input returns Java null. - */ + /// Tolerant Variant extraction. Malformed input returns Java null. @Nullable public static byte[] tryVariantGet(@Nullable byte[] envelope, String path) { return (byte[]) tryVariantGet(envelope, compilePath(path), ResultType.VARIANT); } - /** - * Tolerant typed extraction. Malformed input and incompatible types return Java null. - */ + /// Tolerant typed extraction. Malformed input and incompatible types return Java null. @Nullable public static Object tryVariantGet(@Nullable byte[] envelope, String path, String targetType) { try { @@ -412,9 +400,7 @@ public static Object tryVariantGet(@Nullable byte[] envelope, String path, Strin } } - /** - * Tolerant extraction using pre-parsed path and type values. - */ + /// Tolerant extraction using pre-parsed path and type values. @Nullable public static Object tryVariantGet(@Nullable byte[] envelope, VariantPath path, ResultType targetType) { try { @@ -425,12 +411,10 @@ public static Object tryVariantGet(@Nullable byte[] envelope, VariantPath path, } } - /** - * Tolerantly extracts into a reusable, unboxed result. - * - * @return {@code false} for SQL null, missing paths, Variant null converted to a non-Variant target, malformed input, - * or an incompatible conversion - */ + /// Tolerantly extracts into a reusable, unboxed result. + /// + /// @return {@code false} for SQL null, missing paths, Variant null converted to a non-Variant target, + /// malformed input, or an incompatible conversion public static boolean tryExtractInto(@Nullable byte[] envelope, VariantPath path, ResultType targetType, ReusableResult result) { Objects.requireNonNull(result, "result must not be null"); @@ -455,26 +439,20 @@ public static boolean tryExtractInto(@Nullable byte[] envelope, VariantPath path } } - /** - * Returns whether the path is present. A present Variant null counts as present. - */ + /// Returns whether the path is present. A present Variant null counts as present. @Nullable public static Boolean variantExists(@Nullable byte[] envelope, String path) { return variantExists(envelope, compilePath(path)); } - /** - * Returns whether a compiled path is present. A present Variant null counts as present. - */ + /// Returns whether a compiled path is present. A present Variant null counts as present. @Nullable public static Boolean variantExists(@Nullable byte[] envelope, VariantPath path) { return variantExists(envelope, path, new ReusableResult()); } - /** - * Allocation-free compiled-path form of {@link #variantExists(byte[], VariantPath)} when the caller retains the - * supplied result between rows. - */ + /// Allocation-free compiled-path form of {@link #variantExists(byte[], VariantPath)} when the caller retains the + /// supplied result between rows. @Nullable public static Boolean variantExists(@Nullable byte[] envelope, VariantPath path, ReusableResult result) { Objects.requireNonNull(result, "result must not be null"); @@ -484,31 +462,23 @@ public static Boolean variantExists(@Nullable byte[] envelope, VariantPath path, return result._cursor.navigate(envelope, Objects.requireNonNull(path, "path must not be null")); } - /** - * Returns whether the root value is a Variant null. SQL null is not a Variant null. - */ + /// Returns whether the root value is a Variant null. SQL null is not a Variant null. public static boolean isVariantNull(@Nullable byte[] envelope) { return isVariantNull(envelope, ROOT_PATH, new ReusableResult()); } - /** - * Returns whether a present value at the path is a Variant null. SQL null and missing paths return false. - */ + /// Returns whether a present value at the path is a Variant null. SQL null and missing paths return false. public static boolean isVariantNull(@Nullable byte[] envelope, String path) { return isVariantNull(envelope, compilePath(path)); } - /** - * Returns whether a present value at a compiled path is a Variant null. SQL null and missing paths return false. - */ + /// Returns whether a present value at a compiled path is a Variant null. SQL null and missing paths return false. public static boolean isVariantNull(@Nullable byte[] envelope, VariantPath path) { return isVariantNull(envelope, path, new ReusableResult()); } - /** - * Allocation-free compiled-path form of {@link #isVariantNull(byte[], VariantPath)} when the caller retains the - * supplied result between rows. - */ + /// Allocation-free compiled-path form of {@link #isVariantNull(byte[], VariantPath)} when the caller retains the + /// supplied result between rows. public static boolean isVariantNull(@Nullable byte[] envelope, VariantPath path, ReusableResult result) { Objects.requireNonNull(result, "result must not be null"); if (isSqlNull(envelope)) { @@ -519,34 +489,26 @@ public static boolean isVariantNull(@Nullable byte[] envelope, VariantPath path, && cursor.getType() == Variant.Type.NULL; } - /** - * Returns the Variant type name at the root, or Java null for SQL null. - */ + /// Returns the Variant type name at the root, or Java null for SQL null. @Nullable public static String variantTypeOf(@Nullable byte[] envelope) { return variantTypeOf(envelope, ROOT_PATH, new ReusableResult()); } - /** - * Returns the Variant type name at a path, or Java null for SQL null or a missing path. - */ + /// Returns the Variant type name at a path, or Java null for SQL null or a missing path. @Nullable public static String variantTypeOf(@Nullable byte[] envelope, String path) { return variantTypeOf(envelope, compilePath(path)); } - /** - * Returns the Variant type name at a compiled path, or Java null for SQL null or a missing path. - */ + /// Returns the Variant type name at a compiled path, or Java null for SQL null or a missing path. @Nullable public static String variantTypeOf(@Nullable byte[] envelope, VariantPath path) { return variantTypeOf(envelope, path, new ReusableResult()); } - /** - * Allocation-free compiled-path form of {@link #variantTypeOf(byte[], VariantPath)} when the caller retains the - * supplied result between rows. - */ + /// Allocation-free compiled-path form of {@link #variantTypeOf(byte[], VariantPath)} when the caller retains the + /// supplied result between rows. @Nullable public static String variantTypeOf(@Nullable byte[] envelope, VariantPath path, ReusableResult result) { Objects.requireNonNull(result, "result must not be null"); @@ -558,9 +520,7 @@ public static String variantTypeOf(@Nullable byte[] envelope, VariantPath path, ? typeName(cursor.getType()) : null; } - /** - * Renders the Variant value as canonical JSON text without constructing a JSON tree. - */ + /// Renders the Variant value as canonical JSON text without constructing a JSON tree. @Nullable public static String variantToJson(@Nullable byte[] envelope) { if (isSqlNull(envelope)) { @@ -572,9 +532,7 @@ public static String variantToJson(@Nullable byte[] envelope) { return variantToJson(cursor.asVariant()); } - /** - * Parses JSON text into a Pinot Variant envelope without constructing a JSON tree. - */ + /// Parses JSON text into a Pinot Variant envelope without constructing a JSON tree. @Nullable public static byte[] parseJsonToVariant(@Nullable String json) { if (json == null) { @@ -597,9 +555,7 @@ public static byte[] parseJsonToVariant(@Nullable String json) { } } - /** - * Tolerant JSON parser. Malformed or unsupported input returns Java null. - */ + /// Tolerant JSON parser. Malformed or unsupported input returns Java null. @Nullable public static byte[] tryParseJsonToVariant(@Nullable String json) { try { @@ -1205,12 +1161,10 @@ private static IllegalArgumentException unsupportedVariantDecimal(BigDecimal val + ", scale=" + value.scale()); } - /** - * Mutable zero-copy view over one selected value in a Pinot envelope. - * - *

The constants and layouts used here mirror Parquet Variant encoding version 1. Keeping this cursor on - * {@link ReusableResult} avoids allocating envelope views, Variant wrappers, and navigation wrappers for every row. - */ + /// Mutable zero-copy view over one selected value in a Pinot envelope. + /// + ///

The constants and layouts used here mirror Parquet Variant encoding version 1. Keeping this cursor on + /// {@link ReusableResult} avoids allocating envelope views, Variant wrappers, and navigation wrappers for every row. private static final class Cursor { private byte[] _envelope; private int _metadataOffset; diff --git a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java index 7ad55045abe1..e3d99235a0c9 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java @@ -26,9 +26,7 @@ import static org.testng.Assert.assertTrue; -/** - * Tests strict and tolerant behavior exposed by the public scalar VARIANT function facade. - */ +/// Tests strict and tolerant behavior exposed by the public scalar VARIANT function facade. public class VariantFunctionsTest { @Test public void testScalarFunctionFacade() { diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java index 9a7cb3317112..4d1fe9fdc9ea 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java @@ -47,6 +47,18 @@ public class VariantUtilsTest { + @Test + public void testRawVariantResultRequiresNullHandling() { + DataSchema variantSchema = + new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + DataSchema typedSchema = + new DataSchema(new String[]{"eventType"}, new ColumnDataType[]{ColumnDataType.STRING}); + + assertTrue(VariantUtils.requiresNullHandlingForRawVariantResult(variantSchema, false)); + assertFalse(VariantUtils.requiresNullHandlingForRawVariantResult(variantSchema, true)); + assertFalse(VariantUtils.requiresNullHandlingForRawVariantResult(typedSchema, false)); + } + @Test public void testResultTypeContract() { assertEquals(ResultType.BOOLEAN.getDataType(), DataType.BOOLEAN); diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java index 193dd7df1c23..061f5d58809a 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/BaseOp.java @@ -87,9 +87,7 @@ public String getAbsoluteFileName(String fileName) { return _parentDir + CONFIG_PLACEHOLDER + fileName; } - /** - * Returns the system property that gates this operation, or {@code null} when the operation is unconditional. - */ + /// Returns the system property that gates this operation, or {@code null} when the operation is unconditional. @Nullable public String getRunIfSystemProperty() { return _runIfSystemProperty; @@ -99,10 +97,8 @@ public void setRunIfSystemProperty(@Nullable String runIfSystemProperty) { _runIfSystemProperty = runIfSystemProperty; } - /** - * Returns the exact property value required to run this operation, or {@code null} when the operation is - * unconditional. - */ + /// Returns the exact property value required to run this operation, or {@code null} when the operation is + /// unconditional. @Nullable public String getRunIfSystemPropertyValue() { return _runIfSystemPropertyValue; diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java index 0448da66f5b8..f3d156f7d0d5 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java @@ -28,10 +28,8 @@ import org.slf4j.LoggerFactory; -/** - * Verifies that at least one file matching a glob in a system-property directory contains an expected text fragment. - * Instances are configured and invoked serially by the compatibility runner and are not thread-safe. - */ +/// Verifies that at least one file matching a glob in a system-property directory contains an expected text fragment. +/// Instances are configured and invoked serially by the compatibility runner and are not thread-safe. public class FileContainsOp extends BaseOp { private static final Logger LOGGER = LoggerFactory.getLogger(FileContainsOp.class); private static final int MAX_ATTEMPTS = 10; diff --git a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java index 94c5bb7fcb42..101bc2dedef9 100644 --- a/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/QueryOp.java @@ -70,32 +70,24 @@ public void setQueryFileName(String queryFileName) { _queryFileName = queryFileName; } - /** - * Returns the expected-results file, or {@code null} when expected-error mode is configured. - */ + /// Returns the expected-results file, or {@code null} when expected-error mode is configured. @Nullable public String getExpectedResultsFileName() { return _expectedResultsFileName; } - /** - * Configures result-comparison mode. Pass {@code null} to clear it before configuring expected-error mode. - */ + /// Configures result-comparison mode. Pass {@code null} to clear it before configuring expected-error mode. public void setExpectedResultsFileName(@Nullable String expectedResultsFileName) { _expectedResultsFileName = expectedResultsFileName; } - /** - * Returns the required error-message substring, or {@code null} when result-comparison mode is configured. - */ + /// Returns the required error-message substring, or {@code null} when result-comparison mode is configured. @Nullable public String getExpectedErrorMessageContains() { return _expectedErrorMessageContains; } - /** - * Configures expected-error mode. Pass {@code null} to clear it before configuring result-comparison mode. - */ + /// Configures expected-error mode. Pass {@code null} to clear it before configuring result-comparison mode. public void setExpectedErrorMessageContains(@Nullable String expectedErrorMessageContains) { _expectedErrorMessageContains = expectedErrorMessageContains; } @@ -301,9 +293,7 @@ static boolean matchesExpectedResponse(JsonNode actual, @Nullable JsonNode expec : SqlResultComparator.areEqual(actual, expected, query); } - /** - * Returns whether a response has no result rows and contains only errors matching the required substring. - */ + /// Returns whether a response has no result rows and contains only errors matching the required substring. static boolean hasExpectedError(@Nullable JsonNode response, @Nullable String expectedErrorMessageContains) { if (response == null || expectedErrorMessageContains == null || expectedErrorMessageContains.isEmpty()) { return false; diff --git a/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java index 6a6dc6c55723..b3e41a81d24e 100644 --- a/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java +++ b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java @@ -25,9 +25,7 @@ import static org.testng.Assert.assertTrue; -/** - * Tests capability-gated compatibility operations. - */ +/// Tests capability-gated compatibility operations. public class BaseOpTest { private static final String TEST_PROPERTY = BaseOpTest.class.getName() + ".enabled"; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlock.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlock.java index 61f6f0d815d8..5e9d97b38c50 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlock.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlock.java @@ -25,8 +25,10 @@ import javax.annotation.Nullable; import org.apache.pinot.common.datatable.DataTable; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.selection.SelectionOperatorUtils; +import org.apache.pinot.spi.exception.QueryErrorCode; /// Results block for selection queries. @@ -80,6 +82,9 @@ public List getRows() { @Override public DataTable getDataTable() throws IOException { + if (VariantUtils.requiresNullHandlingForRawVariantResult(_dataSchema, _queryContext.isNullHandlingEnabled())) { + throw QueryErrorCode.QUERY_VALIDATION.asException(VariantUtils.RAW_VARIANT_REQUIRES_NULL_HANDLING_ERROR); + } return SelectionOperatorUtils.getDataTableFromRows(_rows, _dataSchema, _queryContext.isNullHandlingEnabled()); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java index 33756eaa0f99..f62971f03d25 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java @@ -29,13 +29,11 @@ import org.roaringbitmap.RoaringBitmap; -/** - * Shared single-stage lifecycle for functions that inspect or extract a VARIANT value. - * - *

The first argument is consistently validated as a single-value VARIANT-compatible operand, the optional path is - * compiled once, and each input block is evaluated at most once even when the engine asks for both values and a null - * bitmap. Instances are query-local and not thread-safe. - */ +/// Shared single-stage lifecycle for functions that inspect or extract a VARIANT value. +/// +///

The first argument is consistently validated as a single-value VARIANT-compatible operand, the optional path is +/// compiled once, and each input block is evaluated at most once even when the engine asks for both values and a null +/// bitmap. Instances are query-local and not thread-safe. abstract class BaseVariantTransformFunction extends BaseTransformFunction { private static final VariantPath ROOT_PATH = VariantUtils.compilePath("$"); @@ -47,14 +45,12 @@ abstract class BaseVariantTransformFunction extends BaseTransformFunction { @Nullable private RoaringBitmap _cachedNullBitmap; - /** - * Initializes the common VARIANT operand and path contract. - * - * @param arguments function arguments - * @param minArguments minimum accepted argument count - * @param maxArguments maximum accepted argument count - * @param pathRequired whether argument 1 is mandatory - */ + /// Initializes the common VARIANT operand and path contract. + /// + /// @param arguments function arguments + /// @param minArguments minimum accepted argument count + /// @param maxArguments maximum accepted argument count + /// @param pathRequired whether argument 1 is mandatory protected final void initVariantArguments(List arguments, int minArguments, int maxArguments, boolean pathRequired) { int numArguments = arguments.size(); @@ -137,25 +133,17 @@ protected final void ensureEvaluated(ValueBlock valueBlock) { _cachedValueBlock = valueBlock; } - /** - * Returns whether an input SQL null should remain SQL null in the result. - */ + /// Returns whether an input SQL null should remain SQL null in the result. protected boolean inputNullIsResultNull() { return true; } - /** - * Initializes the output array for a block. - */ + /// Initializes the output array for a block. protected abstract void initResultValues(int numDocs); - /** - * Evaluates one non-SQL-null VARIANT. Returns {@code false} when the result should be SQL null. - */ + /// Evaluates one non-SQL-null VARIANT. Returns {@code false} when the result should be SQL null. protected abstract boolean evaluateVariant(@Nullable byte[] variant, int index); - /** - * Stores the type-specific placeholder for a SQL-null result. - */ + /// Stores the type-specific placeholder for a SQL-null result. protected abstract void setNullValue(int index); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java index b99be20e70d5..fc54703ca191 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.java @@ -28,12 +28,10 @@ import org.apache.pinot.core.operator.transform.TransformResultMetadata; -/** - * Returns whether a Variant root value or a value selected by a literal path is an encoded Variant null. - * - *

SQL null and missing paths return a non-null {@code false}; only a present encoded Variant null returns - * {@code true}. Instances are query-local and not thread-safe. - */ +/// Returns whether a Variant root value or a value selected by a literal path is an encoded Variant null. +/// +///

SQL null and missing paths return a non-null {@code false}; only a present encoded Variant null returns +/// {@code true}. Instances are query-local and not thread-safe. public class IsVariantNullTransformFunction extends BaseVariantTransformFunction { @Override diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java index eb8486e5ba53..03d9892a66a8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.java @@ -31,11 +31,9 @@ import org.roaringbitmap.RoaringBitmap; -/** - * Parses JSON text into a logical Variant value for single-stage queries and ingestion transforms. - * - *

Instances are query-local and not thread-safe. - */ +/// Parses JSON text into a logical Variant value for single-stage queries and ingestion transforms. +/// +///

Instances are query-local and not thread-safe. public class ParseJsonToVariantTransformFunction extends BaseTransformFunction { public static final String FUNCTION_NAME = "parseJson"; private static final TransformResultMetadata RESULT_METADATA = @@ -155,9 +153,7 @@ private byte[] parse(@Nullable String json) { return _tolerant ? VariantUtils.tryParseJsonToVariant(json) : VariantUtils.parseJsonToVariant(json); } - /** - * Tolerant JSON-to-Variant parsing. - */ + /// Tolerant JSON-to-Variant parsing. public static final class Try extends ParseJsonToVariantTransformFunction { public static final String FUNCTION_NAME = "tryParseJson"; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java index 0dec8fe85882..af99b272ac7a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.java @@ -27,13 +27,11 @@ import org.apache.pinot.core.operator.transform.TransformResultMetadata; -/** - * Vectorized single-stage implementation of {@code variantExists}. - * - *

The path must be a string literal and is compiled once during initialization. A present encoded Variant null - * counts as present, a missing path returns {@code false}, and SQL null remains SQL null. Instances are query-local - * and not thread-safe. - */ +/// Vectorized single-stage implementation of {@code variantExists}. +/// +///

The path must be a string literal and is compiled once during initialization. A present encoded Variant null +/// counts as present, a missing path returns {@code false}, and SQL null remains SQL null. Instances are query-local +/// and not thread-safe. public class VariantExistsTransformFunction extends BaseVariantTransformFunction { public static final String FUNCTION_NAME = "variantExists"; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java index 43d8e688505d..592bba90a77c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java @@ -31,13 +31,11 @@ import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder; -/** - * Single-stage typed extraction from a Variant envelope. - * - *

The path and optional target type must be literals, so they are parsed once during initialization. Omitting the - * target type returns a Variant. A missing path returns SQL null; the strict form throws for an incompatible non-null - * value, while {@link Try} maps that failure to SQL null. Instances are query-local and not thread-safe. - */ +/// Single-stage typed extraction from a Variant envelope. +/// +///

The path and optional target type must be literals, so they are parsed once during initialization. Omitting the +/// target type returns a Variant. A missing path returns SQL null; the strict form throws for an incompatible non-null +/// value, while {@link Try} maps that failure to SQL null. Instances are query-local and not thread-safe. public class VariantGetTransformFunction extends BaseVariantTransformFunction { public static final String FUNCTION_NAME = "variantGet"; private final boolean _tolerant; @@ -240,9 +238,7 @@ private void setExtractedValue(int index) { } } - /** - * Tolerant Variant extraction. - */ + /// Tolerant Variant extraction. public static final class Try extends VariantGetTransformFunction { public static final String FUNCTION_NAME = "tryVariantGet"; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java index 75bd1f9af027..c6843f06134f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.java @@ -28,12 +28,10 @@ import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder; -/** - * Returns the type name of a Variant root value or a value selected by a literal path. - * - *

SQL null and missing paths produce SQL null. An encoded Variant null is a present value whose type name is - * {@code NULL}. Instances are query-local and not thread-safe. - */ +/// Returns the type name of a Variant root value or a value selected by a literal path. +/// +///

SQL null and missing paths produce SQL null. An encoded Variant null is a present value whose type name is +/// {@code NULL}. Instances are query-local and not thread-safe. public class VariantTypeOfTransformFunction extends BaseVariantTransformFunction { public static final String FUNCTION_NAME = "variantTypeOf"; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlockTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlockTest.java new file mode 100644 index 000000000000..12277b18e00c --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/blocks/results/SelectionResultsBlockTest.java @@ -0,0 +1,54 @@ +/** + * 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.operator.blocks.results; + +import java.util.List; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; +import org.apache.pinot.spi.utils.ByteArray; +import org.testng.Assert; +import org.testng.annotations.Test; + + +/// Tests the final selection-result contracts for logical VARIANT values. +public class SelectionResultsBlockTest { + private static final DataSchema VARIANT_SCHEMA = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.VARIANT}); + + @Test + public void testRawVariantProjectionRequiresNullHandling() + throws Exception { + List rows = List.of( + new Object[]{new ByteArray(VariantUtils.parseJsonToVariant("{\"answer\":42}"))}); + SelectionResultsBlock disabledBlock = + new SelectionResultsBlock(VARIANT_SCHEMA, rows, new QueryContext.Builder().build()); + + QueryException exception = Assert.expectThrows(QueryException.class, disabledBlock::getDataTable); + Assert.assertEquals(exception.getErrorCode(), QueryErrorCode.QUERY_VALIDATION); + Assert.assertTrue(exception.getMessage().contains("requires query null handling")); + + QueryContext nullAwareContext = new QueryContext.Builder().build(); + nullAwareContext.setNullHandlingEnabled(true); + SelectionResultsBlock enabledBlock = new SelectionResultsBlock(VARIANT_SCHEMA, rows, nullAwareContext); + Assert.assertEquals(enabledBlock.getDataTable().getNumberOfRows(), 1); + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java index e050286c8ded..0edf5a46abdb 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java @@ -35,10 +35,8 @@ import org.testng.annotations.Test; -/** - * End-to-end coverage for creating a VARIANT table, ingesting Parquet VARIANT(1), materializing a hot path, and - * querying nested values with both Pinot query engines. - */ +/// End-to-end coverage for creating a VARIANT table, ingesting Parquet VARIANT(1), materializing a hot path, and +/// querying nested values with both Pinot query engines. @Test(suiteName = "CustomClusterIntegrationTest") public class VariantTypeTest extends CustomDataQueryClusterIntegrationTest { private static final String RESOURCE_DIRECTORY = "examples/batch/variantEvents/"; @@ -273,6 +271,12 @@ public void testRawVariantComparisonGroupingAndDistinctAreRejected(boolean useMu "SELECT " + EVENT_ID + " FROM " + TABLE_NAME + " WHERE " + PAYLOAD + " = parse_json('{\"eventType\":\"checkout\"}')"); assertExceptionContains(response, "raw variant", "comparison"); + for (String operator : List.of("IS DISTINCT FROM", "IS NOT DISTINCT FROM")) { + response = postVariantQuery( + "SELECT " + EVENT_ID + " FROM " + TABLE_NAME + " WHERE " + PAYLOAD + " " + operator + + " parse_json(variant_to_json(" + PAYLOAD + "))"); + assertExceptionContains(response, "raw variant", "comparison"); + } response = postVariantQuery( "SELECT " + PAYLOAD + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + PAYLOAD); assertExceptionContains(response, "raw variant", "group by"); @@ -315,6 +319,24 @@ public void testVariantFunctionsRequireQueryNullHandling(boolean useMultiStageQu assertExceptionContains(response, "requires query null handling"); } + @Test(dataProvider = "useBothQueryEngines") + public void testRawVariantProjectionRequiresQueryNullHandling(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postQuery( + "SELECT " + PAYLOAD + " FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " IN ('evt-001', 'evt-005') ORDER BY " + EVENT_ID); + assertExceptionContains(response, "raw variant", "requires query null handling"); + + if (!useMultiStageQueryEngine) { + for (String predicate : List.of("1 = 0", "eventTime < 0")) { + response = postQuery("SELECT " + PAYLOAD + " FROM " + TABLE_NAME + " WHERE " + predicate); + assertExceptionContains(response, "raw variant", "requires query null handling"); + } + } + } + @Test public void testRawVariantJoinIsRejectedButTypedPathJoinWorks() throws Exception { diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java index 481dd2b0df17..e667c2f7fe24 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractor.java @@ -108,8 +108,9 @@ public GenericRow extract(Group from, GenericRow to) { } private void initializeFieldPlans(GroupType schema) { - ParquetVariantConverter[] variantConverters = - ParquetVariantConverter.createTopLevelVariantConverters(schema); + ParquetVariantConverter[] variantConverters = _extractAll + ? ParquetVariantConverter.createTopLevelVariantConverters(schema) + : ParquetVariantConverter.createTopLevelVariantConverters(schema, _fields::contains); int fieldCount = schema.getFieldCount(); int selectedFieldCount = 0; for (int fieldIndex = 0; fieldIndex < fieldCount; fieldIndex++) { diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java index b9e348305b33..4d51d22d4b33 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetNativeRecordExtractorConfig.java @@ -46,13 +46,11 @@ public void setExtractRawTimeValues(boolean extractRawTimeValues) { _extractRawTimeValues = extractRawTimeValues; } - /** - * Supplies the immutable Parquet record schema used to initialize schema-bound logical-type converters. - * - *

The native record reader sets this before initializing the extractor. Direct extractor users should do the - * same when the schema contains VARIANT columns; otherwise the extractor initializes those converters from the - * first record as a compatibility fallback. - */ + /// Supplies the immutable Parquet record schema used to initialize schema-bound logical-type converters. + /// + ///

The native record reader sets this before initializing the extractor. Direct extractor users should do the + /// same when the schema contains VARIANT columns; otherwise the extractor initializes those converters from the + /// first record as a compatibility fallback. public void setParquetSchema(GroupType parquetSchema) { _parquetSchema = parquetSchema; } diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java index d928485ffec3..64b687cad16d 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetUtils.java @@ -83,17 +83,13 @@ public static Schema getParquetAvroSchema(Path path) } } - /** - * Returns the physical Parquet schema for the given file path. - */ + /// Returns the physical Parquet schema for the given file path. public static MessageType getParquetSchema(Path path) throws IOException { return getParquetFileMetadata(path).getSchema(); } - /** - * Returns the immutable footer metadata for the given Parquet file path. - */ + /// Returns the immutable footer metadata for the given Parquet file path. public static FileMetaData getParquetFileMetadata(Path path) throws IOException { InputFile inputFile = HadoopInputFile.fromPath(path, getParquetHadoopConfiguration()); diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java index 6862e49331af..0b624cf60c1e 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java @@ -22,6 +22,7 @@ import java.util.HashSet; import java.util.Set; import java.util.function.Consumer; +import java.util.function.Predicate; import org.apache.parquet.example.data.Group; import org.apache.parquet.io.api.Binary; import org.apache.parquet.io.api.Converter; @@ -69,8 +70,15 @@ private ParquetVariantConverter(GroupType variantType) { /// Pinot currently supports only non-repeated, top-level `VARIANT(1)` values. Failing during reader /// initialization avoids silently surfacing an unsupported nested/repeated Variant as an ordinary struct. static Set validateAndGetTopLevelVariantFields(GroupType schema) { + return validateAndGetTopLevelVariantFields(schema, fieldName -> true); + } + + private static Set validateAndGetTopLevelVariantFields(GroupType schema, Predicate fieldSelector) { Set variantFields = new HashSet<>(); for (Type field : schema.getFields()) { + if (!fieldSelector.test(field.getName())) { + continue; + } if (isVariant(field)) { validateTopLevelVariant(field); variantFields.add(field.getName()); @@ -84,6 +92,17 @@ static Set validateAndGetTopLevelVariantFields(GroupType schema) { /// Validates the file schema and builds an index-aligned reusable converter tree for each top-level VARIANT column. static ParquetVariantConverter[] createTopLevelVariantConverters(GroupType schema) { Set variantFields = validateAndGetTopLevelVariantFields(schema); + return buildTopLevelVariantConverters(schema, variantFields); + } + + /// Validates selected top-level fields and builds index-aligned converters for selected VARIANT columns. + static ParquetVariantConverter[] createTopLevelVariantConverters(GroupType schema, Predicate fieldSelector) { + Set variantFields = validateAndGetTopLevelVariantFields(schema, fieldSelector); + return buildTopLevelVariantConverters(schema, variantFields); + } + + private static ParquetVariantConverter[] buildTopLevelVariantConverters(GroupType schema, + Set variantFields) { ParquetVariantConverter[] variantConverters = new ParquetVariantConverter[schema.getFieldCount()]; for (int fieldIndex = 0; fieldIndex < schema.getFieldCount(); fieldIndex++) { Type field = schema.getType(fieldIndex); diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java index e6a2058f238d..450f31cf5a1d 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java @@ -66,9 +66,7 @@ import static org.testng.Assert.expectThrows; -/** - * End-to-end coverage for Parquet `VARIANT(1)` reconstruction and reader selection. - */ +/// End-to-end coverage for Parquet `VARIANT(1)` reconstruction and reader selection. public class ParquetVariantRecordReaderTest { private static final String VARIANT_FIELD = "variant_col"; @@ -428,6 +426,37 @@ public void testVariantDetectionAndUnsupportedShapes() { assertTrue(int96Exception.getMessage().matches("(?i).*int96.*")); } + @Test + public void testNestedVariantValidationHonorsSelectedFields() { + MessageType schema = MessageTypeParser.parseMessageType( + "message nested {" + + " required int32 id;" + + " optional group wrapper {" + + " optional group variant_col (VARIANT(1)) {" + + " required binary metadata;" + + " optional binary value;" + + " }" + + " }" + + "}"); + ParquetNativeRecordExtractorConfig config = new ParquetNativeRecordExtractorConfig(); + config.setParquetSchema(schema); + + ParquetNativeRecordExtractor idExtractor = new ParquetNativeRecordExtractor(); + idExtractor.init(Set.of("id"), config); + GenericRow row = idExtractor.extract(new SimpleGroupFactory(schema).newGroup().append("id", 7), new GenericRow()); + assertEquals(row.getFieldToValueMap(), Map.of("id", 7)); + + ParquetNativeRecordExtractor wrapperExtractor = new ParquetNativeRecordExtractor(); + UnsupportedOperationException selectedException = expectThrows(UnsupportedOperationException.class, + () -> wrapperExtractor.init(Set.of("wrapper"), config)); + assertTrue(selectedException.getMessage().contains("Nested")); + + ParquetNativeRecordExtractor extractAll = new ParquetNativeRecordExtractor(); + UnsupportedOperationException extractAllException = expectThrows(UnsupportedOperationException.class, + () -> extractAll.init(null, config)); + assertTrue(extractAllException.getMessage().contains("Nested")); + } + private void assertScalarRows(RecordReader reader, File dataFile) throws IOException { try (reader) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java index 09e0a735467a..227089558869 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java @@ -156,6 +156,10 @@ private static void trackEmptyLeafStages(DispatchablePlanContext context) { /// Runs validations on the plan. private void runValidations(PlanFragment planFragment, DispatchablePlanContext context) { PlanNode rootPlanNode = planFragment.getFragmentRoot(); + if (rootPlanNode.getStageId() == 0) { + VariantTypeValidationVisitor.validateResultSchema(rootPlanNode.getDataSchema(), + context.getPlannerContext().getEnvConfig().isNullHandlingEnabled()); + } boolean isIntermediateStage = context.getDispatchablePlanMetadataMap().get(rootPlanNode.getStageId()).getScannedTables().isEmpty(); rootPlanNode.visit(ArrayToMvValidationVisitor.INSTANCE, isIntermediateStage); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java index ee6261b511e8..9939aa8f8388 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java @@ -26,12 +26,10 @@ import org.apache.pinot.query.planner.plannode.JoinNode; -/** - * Validates the logical types used as join keys. - * - *

This stateless validator is shared by planning and execution so that mixed-version plans are rejected with the - * same semantics and error messages regardless of where validation first occurs. - */ +/// Validates the logical types used as join keys. +/// +///

This stateless validator is shared by planning and execution so that mixed-version plans are rejected with the +/// same semantics and error messages regardless of where validation first occurs. public final class JoinKeyTypeValidator { private static final String JOIN_KEY_ERROR = "Raw VARIANT values do not support JOIN keys; extract a typed path with variantGet first"; @@ -41,12 +39,10 @@ public final class JoinKeyTypeValidator { private JoinKeyTypeValidator() { } - /** - * Validates equality/hash keys and, for ASOF joins, ordering keys. - * - *

The right schema can be absent for legacy SEMI and ANTI join plans whose output schema contains only left - * columns. In that case, the caller can validate only the left keys. - */ + /// Validates equality/hash keys and, for ASOF joins, ordering keys. + /// + ///

The right schema can be absent for legacy SEMI and ANTI join plans whose output schema contains only left + /// columns. In that case, the caller can validate only the left keys. public static void validate(JoinNode joinNode, DataSchema leftSchema, @Nullable DataSchema rightSchema) { validateEqualityAndHashing(joinNode.getLeftKeys(), leftSchema); if (rightSchema == null) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java index 955c78ede633..2e017f215bf2 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java @@ -21,6 +21,7 @@ import java.util.List; import org.apache.calcite.rel.RelFieldCollation; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.JoinNode; @@ -33,10 +34,8 @@ import org.apache.pinot.spi.exception.QueryException; -/** - * Rejects operations that would otherwise assign physical byte ordering, equality, or hashing semantics to a raw - * VARIANT value. The visitor has no mutable state and is thread-safe, so callers may share {@link #INSTANCE}. - */ +/// Rejects operations that would otherwise assign physical byte ordering, equality, or hashing semantics to a raw +/// VARIANT value. The visitor has no mutable state and is thread-safe, so callers may share {@link #INSTANCE}. public final class VariantTypeValidationVisitor extends PlanNodeVisitor.DepthFirstVisitor { public static final VariantTypeValidationVisitor INSTANCE = new VariantTypeValidationVisitor(); @@ -57,19 +56,15 @@ public Void visitAggregate(AggregateNode node, Void context) { return super.visitAggregate(node, context); } - /** - * Validates aggregate operands against their logical input schema. - * - *

This method is also invoked by the runtime as a defensive check for plans that did not pass through the - * current broker planner. - */ + /// Validates aggregate operands against their logical input schema. + /// + ///

This method is also invoked by the runtime as a defensive check for plans that did not pass through the + /// current broker planner. public static void validateAggregateInputs(AggregateNode node, DataSchema inputSchema) { validateAggregateInputs(node.getAggCalls(), inputSchema); } - /** - * Validates aggregate or window-function operands against their logical input schema. - */ + /// Validates aggregate or window-function operands against their logical input schema. public static void validateAggregateInputs(List aggCalls, DataSchema inputSchema) { for (RexExpression.FunctionCall aggCall : aggCalls) { if (isRawVariantIndependent(aggCall)) { @@ -83,6 +78,16 @@ public static void validateAggregateInputs(List aggC } } + /// Rejects a raw VARIANT result when query null handling is disabled. Without the null bitmap, the reserved empty + /// byte placeholder cannot participate in normal disabled-null semantics while also remaining distinguishable from + /// an encoded Variant null. + public static void validateResultSchema(DataSchema resultSchema, boolean nullHandlingEnabled) { + if (VariantUtils.requiresNullHandlingForRawVariantResult(resultSchema, nullHandlingEnabled)) { + throw new QueryException(QueryErrorCode.QUERY_PLANNING, + VariantUtils.RAW_VARIANT_REQUIRES_NULL_HANDLING_ERROR); + } + } + @Override public Void visitSort(SortNode node, Void context) { DataSchema dataSchema = node.getDataSchema(); @@ -116,12 +121,10 @@ public Void visitJoin(JoinNode node, Void context) { return super.visitJoin(node, context); } - /** - * Validates equality/hash join keys against both logical input schemas. - * - *

The runtime also invokes this for mixed-version plans, including LOOKUP joins that do not construct a - * {@code HashJoinOperator}. - */ + /// Validates equality/hash join keys against both logical input schemas. + /// + ///

The runtime also invokes this for mixed-version plans, including LOOKUP joins that do not construct a + /// {@code HashJoinOperator}. public static void validateJoinInputs(JoinNode node, DataSchema leftSchema, DataSchema rightSchema) { try { JoinKeyTypeValidator.validate(node, leftSchema, rightSchema); @@ -139,12 +142,10 @@ public Void visitWindow(WindowNode node, Void context) { return super.visitWindow(node, context); } - /** - * Validates window partition keys, ordering keys, and function operands against their logical input schema. - * - *

This method is also invoked by the runtime as a defensive check for plans that did not pass through the - * current broker planner. - */ + /// Validates window partition keys, ordering keys, and function operands against their logical input schema. + /// + ///

This method is also invoked by the runtime as a defensive check for plans that did not pass through the + /// current broker planner. public static void validateWindowInputs(WindowNode node, DataSchema inputSchema) { for (int key : node.getKeys()) { DataSchema.ColumnDataType dataType = inputSchema.getColumnDataType(key); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java index 7c390843f00f..c558e15464db 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java @@ -51,6 +51,16 @@ public void testRejectsVariantOrderBy() { Assert.assertTrue(exception.getMessage().contains("ORDER BY")); } + @Test + public void testRawVariantProjectionRequiresNullHandling() { + QueryException exception = Assert.expectThrows(QueryException.class, + () -> VariantTypeValidationVisitor.validateResultSchema(VARIANT_SCHEMA, false)); + Assert.assertTrue(exception.getMessage().contains("requires query null handling")); + + VariantTypeValidationVisitor.validateResultSchema(VARIANT_SCHEMA, true); + VariantTypeValidationVisitor.validateResultSchema(TYPED_EXTRACTION_SCHEMA, false); + } + @Test public void testRejectsEqualityDependentSetOperations() { List unsupportedNodes = List.of( diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java index a7ec713a7e47..33921e82ccff 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java @@ -69,13 +69,11 @@ public class HashJoinOperator extends BaseJoinOperator { @Nullable private List _nullKeyRightRows; - /** - * Creates a hash join using schemas available on the join node. - * - *

For SEMI and ANTI joins whose node does not carry its inputs, the result schema contains only left columns, so - * this legacy constructor cannot validate the right key's logical type. New callers that need right-side VARIANT - * validation must use the overload that accepts {@code rightSchema}. - */ + /// Creates a hash join using schemas available on the join node. + /// + ///

For SEMI and ANTI joins whose node does not carry its inputs, the result schema contains only left columns, so + /// this legacy constructor cannot validate the right key's logical type. New callers that need right-side VARIANT + /// validation must use the overload that accepts {@code rightSchema}. public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, MultiStageOperator rightInput, JoinNode node) { this(context, leftInput, leftSchema, rightInput, tryInferRightSchema(leftSchema, node), node, false); @@ -101,13 +99,11 @@ private HashJoinOperator(OpChainExecutionContext context, MultiStageOperator lef _nullKeyRightRows = needUnmatchedRightRows() ? new ArrayList<>() : null; } - /** - * Constructor that takes the schema for NonEquiEvaluator as an argument. - * - *

For SEMI and ANTI joins whose node does not carry its inputs, the result schema contains only left columns, so - * this legacy constructor cannot validate the right key's logical type. New callers that need right-side VARIANT - * validation must use the overload that accepts {@code rightSchema}. - */ + /// Constructor that takes the schema for NonEquiEvaluator as an argument. + /// + ///

For SEMI and ANTI joins whose node does not carry its inputs, the result schema contains only left columns, so + /// this legacy constructor cannot validate the right key's logical type. New callers that need right-side VARIANT + /// validation must use the overload that accepts {@code rightSchema}. public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema, MultiStageOperator rightInput, JoinNode node, DataSchema nonEquiEvaluationSchema) { this(context, leftInput, leftSchema, rightInput, tryInferRightSchema(leftSchema, node), node, diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FunctionOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FunctionOperand.java index c9d01d0f8cda..e8e1c4afb64a 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FunctionOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FunctionOperand.java @@ -51,17 +51,7 @@ public FunctionOperand(RexExpression.FunctionCall functionCall, DataSchema dataS int numOperands = operands.size(); ColumnDataType[] argumentTypes = new ColumnDataType[numOperands]; for (int i = 0; i < numOperands; i++) { - RexExpression operand = operands.get(i); - ColumnDataType argumentType; - if (operand instanceof RexExpression.InputRef) { - argumentType = dataSchema.getColumnDataType(((RexExpression.InputRef) operand).getIndex()); - } else if (operand instanceof RexExpression.Literal) { - argumentType = ((RexExpression.Literal) operand).getDataType(); - } else { - assert operand instanceof RexExpression.FunctionCall; - argumentType = ((RexExpression.FunctionCall) operand).getDataType(); - } - argumentTypes[i] = argumentType; + argumentTypes[i] = TransformOperandFactory.getResultType(operands.get(i), dataSchema); } String functionName = functionCall.getFunctionName(); String canonicalName = FunctionRegistry.canonicalize(functionName); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java index 709ce6112fbf..2075d0afe29e 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java @@ -26,13 +26,11 @@ import org.apache.pinot.query.planner.logical.RexExpression; -/** - * Query-local constant operand for parsing a JSON literal into Variant. - * - *

The literal is parsed exactly once while the expression tree is constructed. The cached internal value is - * immutable by convention and can therefore be reused for every input row. Instances are thread-safe after - * construction. - */ +/// Query-local constant operand for parsing a JSON literal into Variant. +/// +///

The literal is parsed exactly once while the expression tree is constructed. The cached internal value is +/// immutable by convention and can therefore be reused for every input row. Instances are thread-safe after +/// construction. final class LiteralParseJsonOperand implements TransformOperand { private static final String PARSE_JSON = "parsejson"; private static final String PARSE_JSON_TO_VARIANT = "parsejsontovariant"; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java index b7e0470f4795..8581f7b7d8af 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java @@ -21,11 +21,17 @@ import com.google.common.base.Preconditions; import java.util.List; import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.TransformFunctionType; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.query.planner.logical.RexExpression; public class TransformOperandFactory { + private static final String IS_DISTINCT_FROM = + FunctionRegistry.canonicalize(TransformFunctionType.IS_DISTINCT_FROM.getName()); + private static final String IS_NOT_DISTINCT_FROM = + FunctionRegistry.canonicalize(TransformFunctionType.IS_NOT_DISTINCT_FROM.getName()); + private TransformOperandFactory() { } @@ -52,6 +58,15 @@ private static TransformOperand getTransformOperand(RexExpression.FunctionCall f && operands.get(0) instanceof RexExpression.Literal) { return new LiteralParseJsonOperand(functionCall, canonicalName); } + if (canonicalName.equals(IS_DISTINCT_FROM) || canonicalName.equals(IS_NOT_DISTINCT_FROM)) { + Preconditions.checkState(numOperands == 2, "%s takes 2 arguments, got: %s", functionCall.getFunctionName(), + numOperands); + for (RexExpression operand : operands) { + Preconditions.checkArgument(getResultType(operand, dataSchema).supportsEquality(), + "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); + } + return new FunctionOperand(functionCall, dataSchema); + } switch (functionCall.getFunctionName()) { case "AND": Preconditions.checkState(numOperands >= 2, "AND takes >=2 arguments, got: %s", numOperands); @@ -90,4 +105,17 @@ private static TransformOperand getTransformOperand(RexExpression.FunctionCall f return new FunctionOperand(functionCall, dataSchema); } } + + static DataSchema.ColumnDataType getResultType(RexExpression rexExpression, DataSchema dataSchema) { + if (rexExpression instanceof RexExpression.InputRef) { + return dataSchema.getColumnDataType(((RexExpression.InputRef) rexExpression).getIndex()); + } + if (rexExpression instanceof RexExpression.Literal) { + return ((RexExpression.Literal) rexExpression).getDataType(); + } + if (rexExpression instanceof RexExpression.FunctionCall) { + return ((RexExpression.FunctionCall) rexExpression).getDataType(); + } + throw new UnsupportedOperationException("Unsupported RexExpression: " + rexExpression); + } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java index b76bfef52a81..71f7c8e0425e 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java @@ -30,13 +30,11 @@ import org.apache.pinot.query.planner.logical.RexExpression; -/** - * Query-local multi-stage operand for Variant scalar operations. - * - *

Literal paths and target types are compiled once at construction. Values enter and leave the operand in - * {@link DataSchema}'s internal representation, which {@link VariantUtils.ReusableResult} materializes directly after - * extraction. The reusable extraction result makes instances not thread-safe. - */ +/// Query-local multi-stage operand for Variant scalar operations. +/// +///

Literal paths and target types are compiled once at construction. Values enter and leave the operand in +/// {@link DataSchema}'s internal representation, which {@link VariantUtils.ReusableResult} materializes directly after +/// extraction. The reusable extraction result makes instances not thread-safe. final class VariantOperand implements TransformOperand { private static final VariantPath ROOT_PATH = VariantUtils.compilePath("$"); private static final String VARIANT_GET = "variantget"; diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java index 62ed1d5f50f4..e8768419c4dd 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java @@ -69,4 +69,15 @@ public void testRawVariantNotInIsRejected() { () -> new FilterOperand.In(VARIANT_OPERANDS, VARIANT_SCHEMA, true)); Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support IN")); } + + @Test + public void testRawVariantDistinctFromIsRejected() { + for (String functionName : List.of("IS_DISTINCT_FROM", "isNotDistinctFrom")) { + RexExpression.FunctionCall functionCall = + new RexExpression.FunctionCall(ColumnDataType.BOOLEAN, functionName, VARIANT_OPERANDS); + IllegalArgumentException exception = Assert.expectThrows(IllegalArgumentException.class, + () -> TransformOperandFactory.getTransformOperand(functionCall, VARIANT_SCHEMA)); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support comparison")); + } + } } diff --git a/pinot-spi/VARIANT_DESIGN.md b/pinot-spi/VARIANT_DESIGN.md index 512934cd3819..eedadc2e5b27 100644 --- a/pinot-spi/VARIANT_DESIGN.md +++ b/pinot-spi/VARIANT_DESIGN.md @@ -267,6 +267,8 @@ Storage and query null handling are required because four states must remain dis Extracting a Variant null as `VARIANT` retains its non-empty envelope. `variant_to_json` returns SQL null for SQL null, JSON text `null` for Variant null, and JSON text `"null"` (including quotes) for the Variant string. +Raw Variant projection is rejected when query null handling is disabled because the +reserved empty-byte placeholder cannot preserve these states under disabled-null semantics. ## 8. Query wire and client compatibility diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java index 3cb3ee746b42..6401294522f4 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/PinotDataType.java @@ -1019,9 +1019,7 @@ public Object convert(Object value, PinotDataType sourceType) { } }, - /** - * Pinot's external representation of a VARIANT value: a validated PVAR envelope in a {@code byte[]}. - */ + /// Pinot's external representation of a VARIANT value: a validated PVAR envelope in a `byte[]`. VARIANT { @Override public int toInt(Object value) { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java index ab9579f91b12..f908cf99fabc 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java @@ -23,27 +23,25 @@ import javax.annotation.Nullable; -/** - * Pinot-owned framing for the two buffers that make up a Parquet Variant value. - * - *

The version-1 wire format is: - * 

- *   0        4 bytes  ASCII magic "PVAR"
- *   4        1 byte   envelope version (1)
- *   5        1 byte   flags (0)
- *   6        2 bytes  reserved (0)
- *   8        4 bytes  metadata length, unsigned range restricted to Java array sizes
- *   12       4 bytes  value length, unsigned range restricted to Java array sizes
- *   16       M bytes  Parquet Variant metadata
- *   16 + M   V bytes  Parquet Variant value
- *
- * - *

An empty byte array is deliberately not an envelope. Pinot reserves it as the default null value for a - * {@code VARIANT} field, allowing the null-value vector to distinguish SQL null from an encoded Variant null. - * - *

This class validates only Pinot's stable outer framing. Producers and consumers remain responsible for validating - * the Parquet Variant metadata and value payloads. - */ +/// Pinot-owned framing for the two buffers that make up a Parquet Variant value. +/// +/// The version-1 wire format is: +/// ``` +/// 0 4 bytes ASCII magic "PVAR" +/// 4 1 byte envelope version (1) +/// 5 1 byte flags (0) +/// 6 2 bytes reserved (0) +/// 8 4 bytes metadata length, unsigned range restricted to Java array sizes +/// 12 4 bytes value length, unsigned range restricted to Java array sizes +/// 16 M bytes Parquet Variant metadata +/// 16 + M V bytes Parquet Variant value +/// ``` +/// +/// An empty byte array is deliberately not an envelope. Pinot reserves it as the default null value for a +/// `VARIANT` field, allowing the null-value vector to distinguish SQL null from an encoded Variant null. +/// +/// This class validates only Pinot's stable outer framing. Producers and consumers remain responsible for validating +/// the Parquet Variant metadata and value payloads. public final class VariantEnvelope { public static final int HEADER_SIZE = 16; public static final byte VERSION = 1; @@ -54,13 +52,11 @@ public final class VariantEnvelope { private VariantEnvelope() { } - /** - * Encodes the remaining bytes of the supplied metadata and value buffers without changing their positions or - * limits. - * - *

Array-backed buffers are copied directly from their backing arrays. Other buffers, including direct and - * read-only buffers, are read through independent duplicate views. - */ + /// Encodes the remaining bytes of the supplied metadata and value buffers without changing their positions or + /// limits. + /// + /// Array-backed buffers are copied directly from their backing arrays. Other buffers, including direct and + /// read-only buffers, are read through independent duplicate views. public static byte[] encode(ByteBuffer metadata, ByteBuffer value) { Objects.requireNonNull(metadata, "metadata must not be null"); Objects.requireNonNull(value, "value must not be null"); @@ -73,9 +69,7 @@ public static byte[] encode(ByteBuffer metadata, ByteBuffer value) { return envelope; } - /** - * Encodes slices of the supplied arrays without allocating intermediate buffer views. - */ + /// Encodes slices of the supplied arrays without allocating intermediate buffer views. public static byte[] encode(byte[] metadata, int metadataOffset, int metadataLength, byte[] value, int valueOffset, int valueLength) { Objects.requireNonNull(metadata, "metadata must not be null"); @@ -89,19 +83,17 @@ public static byte[] encode(byte[] metadata, int metadataOffset, int metadataLen return envelope; } - /** - * Decodes and validates an envelope, returning zero-copy, read-only views over its metadata and value buffers. - * - *

The returned views alias {@code envelope}; this method does not copy either payload. The decoded object and - * any views obtained from it keep the backing array alive, so the caller does not need to retain a separate - * reference to {@code envelope}. Mutations made to the input array after this method returns are visible through - * the views and can corrupt the decoded payload. Callers must therefore treat the input array as immutable for as - * long as the decoded object or any returned view may be used. - * - *

The decoded holder is safe for concurrent reads when the aliased input array is not mutated. Each accessor - * returns a read-only view with independent position and limit, so cursor movement by one reader does not affect - * another reader. - */ + /// Decodes and validates an envelope, returning zero-copy, read-only views over its metadata and value buffers. + /// + /// The returned views alias `envelope`; this method does not copy either payload. The decoded object and any views + /// obtained from it keep the backing array alive, so the caller does not need to retain a separate reference to + /// `envelope`. Mutations made to the input array after this method returns are visible through the views and can + /// corrupt the decoded payload. Callers must therefore treat the input array as immutable for as long as the decoded + /// object or any returned view may be used. + /// + /// The decoded holder is safe for concurrent reads when the aliased input array is not mutated. Each accessor returns + /// a read-only view with independent position and limit, so cursor movement by one reader does not affect another + /// reader. public static Decoded decode(byte[] envelope) { int metadataLength = validateAndGetMetadataLength(envelope); int valueLength = envelope.length - HEADER_SIZE - metadataLength; @@ -112,11 +104,9 @@ public static Decoded decode(byte[] envelope) { return new Decoded(metadata, value); } - /** - * Validates the stable outer framing and returns the metadata length without allocating buffer views. - * - *

The value begins at {@code HEADER_SIZE + metadataLength}; its length is the remaining envelope length. - */ + /// Validates the stable outer framing and returns the metadata length without allocating buffer views. + /// + /// The value begins at `HEADER_SIZE + metadataLength`; its length is the remaining envelope length. public static int validateAndGetMetadataLength(byte[] envelope) { Objects.requireNonNull(envelope, "envelope must not be null"); if (envelope.length < HEADER_SIZE) { @@ -156,9 +146,7 @@ public static int validateAndGetMetadataLength(byte[] envelope) { return metadataLength; } - /** - * Returns whether the bytes form a complete, supported Variant envelope. - */ + /// Returns whether the bytes form a complete, supported Variant envelope. public static boolean isEnvelope(@Nullable byte[] envelope) { if (envelope == null) { return false; @@ -171,13 +159,11 @@ public static boolean isEnvelope(@Nullable byte[] envelope) { } } - /** - * Allocates an initialized envelope with writable, zero-filled metadata and value regions. - * - *

This is the zero-intermediate-copy producer API for sources that can write directly into a destination array. - * The metadata region begins at {@link #HEADER_SIZE}; the value region begins at - * {@code HEADER_SIZE + metadataLength}. Callers must completely fill both regions before publishing the envelope. - */ + /// Allocates an initialized envelope with writable, zero-filled metadata and value regions. + /// + /// This is the zero-intermediate-copy producer API for sources that can write directly into a destination array. + /// The metadata region begins at [#HEADER_SIZE]; the value region begins at `HEADER_SIZE + metadataLength`. Callers + /// must completely fill both regions before publishing the envelope. public static byte[] allocate(int metadataLength, int valueLength) { if (metadataLength < 0 || valueLength < 0) { throw new IllegalArgumentException( @@ -231,12 +217,10 @@ private static void writeInt(byte[] bytes, int offset, int value) { bytes[offset + 3] = (byte) value; } - /** - * Read-only views of the two Parquet Variant buffers stored in an envelope. - * - *

Instances retain and alias the envelope array supplied to {@link VariantEnvelope#decode(byte[])}. They are - * safe for concurrent reads only while that array remains unmodified. - */ + /// Read-only views of the two Parquet Variant buffers stored in an envelope. + /// + /// Instances retain and alias the envelope array supplied to [VariantEnvelope#decode(byte[])]. They are safe for + /// concurrent reads only while that array remains unmodified. public static final class Decoded { private final ByteBuffer _metadata; private final ByteBuffer _value; @@ -246,16 +230,12 @@ private Decoded(ByteBuffer metadata, ByteBuffer value) { _value = value; } - /** - * Returns a read-only metadata view with independent position and limit. - */ + /// Returns a read-only metadata view with independent position and limit. public ByteBuffer getMetadata() { return _metadata.asReadOnlyBuffer(); } - /** - * Returns a read-only value view with independent position and limit. - */ + /// Returns a read-only value view with independent position and limit. public ByteBuffer getValue() { return _value.asReadOnlyBuffer(); } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java b/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java index 5b60836fdd76..976c71912cfd 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java @@ -27,9 +27,7 @@ import org.apache.pinot.tools.admin.command.QuickstartRunner; -/** - * Batch quickstart for ingesting and querying an Apache Parquet VARIANT column. - */ +/// Batch quickstart for ingesting and querying an Apache Parquet VARIANT column. public class VariantQuickStart extends Quickstart { private static final String[] VARIANT_TABLE_DIRECTORIES = {"examples/batch/variantEvents"}; private static final int EXPECTED_NUM_ROWS = 5; From a1869eacd0fa06df8b93ec2cf1620e5b54ec048f Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sat, 8 Aug 2026 23:44:50 -0700 Subject: [PATCH 5/8] Fix VARIANT opacity guards (OBJECT regression), effective partial-upsert validation, and type-accurate error messages Addresses the four review findings on #19101 (2 Medium: OBJECT-comparison regression, partial-upsert default-strategy gap; 2 Low: type-accurate opacity messages, planner GROUP BY key validation). Tests added across pinot-core, pinot-query-runtime, pinot-query-planner, pinot-segment-local. --- .../BinaryOperatorTransformFunction.java | 5 +- .../function/InTransformFunction.java | 3 +- .../query/utils/OrderByComparatorFactory.java | 9 ++- .../VariantTypeValidationVisitor.java | 35 ++++++++---- .../VariantTypeValidationVisitorTest.java | 19 +++++++ .../query/runtime/operator/SortOperator.java | 8 ++- .../SortedMailboxReceiveOperator.java | 7 ++- .../operator/operands/FilterOperand.java | 9 ++- .../operands/TransformOperandFactory.java | 3 +- .../operator/operands/FilterOperandTest.java | 15 +++++ .../segment/local/utils/TableConfigUtils.java | 24 ++++++-- .../VariantTableConfigValidationTest.java | 57 +++++++++++++++++++ 12 files changed, 163 insertions(+), 31 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java index ad2809ed8728..2d81fa6f8531 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BinaryOperatorTransformFunction.java @@ -108,8 +108,9 @@ public void init(List arguments, Map c _rightTransformFunction = arguments.get(1); DataType leftDataType = _leftTransformFunction.getResultMetadata().getDataType(); DataType rightDataType = _rightTransformFunction.getResultMetadata().getDataType(); - Preconditions.checkArgument((leftDataType == DataType.UNKNOWN || leftDataType.supportsOrdering()) - && (rightDataType == DataType.UNKNOWN || rightDataType.supportsOrdering()), + // Reject raw VARIANT operands only; other types keep their existing comparison behavior. VARIANT is opaque + // because its PVAR byte encoding is not a canonical semantic ordering. + Preconditions.checkArgument(leftDataType != DataType.VARIANT && rightDataType != DataType.VARIANT, "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); if (leftDataType == DataType.UNKNOWN || rightDataType == DataType.UNKNOWN) { _alwaysNull = true; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java index c75b52d9313a..6a511d58e9a4 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InTransformFunction.java @@ -65,8 +65,9 @@ public void init(List arguments, Map c + "transform function: (expression, values)", getName()); _mainFunction = arguments.get(0); for (TransformFunction argument : arguments) { + // Reject raw VARIANT operands only; preserve existing IN behavior for all other types. DataType dataType = argument.getResultMetadata().getDataType(); - Preconditions.checkArgument(dataType.supportsEquality() && dataType.supportsHashing(), + Preconditions.checkArgument(dataType != DataType.VARIANT, "Raw VARIANT values do not support IN; extract a typed path with variantGet first"); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java index a11c206751ea..9a7ed1efffbc 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java @@ -27,6 +27,7 @@ import org.apache.pinot.common.request.context.OrderByExpressionContext; import org.apache.pinot.core.data.table.Record; import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.exception.BadQueryRequestException; @@ -51,9 +52,11 @@ public static Comparator getComparator(List throw new BadQueryRequestException("MV expression: " + orderByExpressions.get(i) + " should not be included in the ORDER-BY clause"); } - if (!orderByColumnContexts[i].getDataType().supportsOrdering()) { - throw new BadQueryRequestException( - "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first"); + FieldSpec.DataType dataType = orderByColumnContexts[i].getDataType(); + if (!dataType.supportsOrdering()) { + throw new BadQueryRequestException(dataType == FieldSpec.DataType.VARIANT + ? "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first" + : "ORDER BY does not support " + dataType + " values"); } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java index 2e017f215bf2..ad1178af1333 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java @@ -61,6 +61,12 @@ public Void visitAggregate(AggregateNode node, Void context) { ///

This method is also invoked by the runtime as a defensive check for plans that did not pass through the /// current broker planner. public static void validateAggregateInputs(AggregateNode node, DataSchema inputSchema) { + for (int key : node.getGroupKeys()) { + DataSchema.ColumnDataType dataType = inputSchema.getColumnDataType(key); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + throw unsupported("GROUP BY", dataType); + } + } validateAggregateInputs(node.getAggCalls(), inputSchema); } @@ -71,8 +77,9 @@ public static void validateAggregateInputs(List aggC continue; } for (RexExpression operand : aggCall.getFunctionOperands()) { - if (!getLogicalType(operand, inputSchema).supportsDirectAggregation()) { - throw unsupported("Aggregate function " + aggCall.getFunctionName()); + DataSchema.ColumnDataType dataType = getLogicalType(operand, inputSchema); + if (!dataType.supportsDirectAggregation()) { + throw unsupported("Aggregate function " + aggCall.getFunctionName(), dataType); } } } @@ -93,8 +100,9 @@ public Void visitSort(SortNode node, Void context) { DataSchema dataSchema = node.getDataSchema(); for (RelFieldCollation collation : node.getCollations()) { int fieldIndex = collation.getFieldIndex(); - if (!dataSchema.getColumnDataType(fieldIndex).supportsOrdering()) { - throw unsupported("ORDER BY"); + DataSchema.ColumnDataType dataType = dataSchema.getColumnDataType(fieldIndex); + if (!dataType.supportsOrdering()) { + throw unsupported("ORDER BY", dataType); } } return super.visitSort(node, context); @@ -105,7 +113,7 @@ public Void visitSetOp(SetOpNode node, Void context) { if (!(node.getSetOpType() == SetOpNode.SetOpType.UNION && node.isAll())) { for (DataSchema.ColumnDataType dataType : node.getDataSchema().getColumnDataTypes()) { if (!dataType.supportsEquality() || !dataType.supportsHashing()) { - throw unsupported(node.explain().replace('_', ' ')); + throw unsupported(node.explain().replace('_', ' '), dataType); } } } @@ -150,12 +158,13 @@ public static void validateWindowInputs(WindowNode node, DataSchema inputSchema) for (int key : node.getKeys()) { DataSchema.ColumnDataType dataType = inputSchema.getColumnDataType(key); if (!dataType.supportsEquality() || !dataType.supportsHashing()) { - throw unsupported("Window PARTITION BY"); + throw unsupported("Window PARTITION BY", dataType); } } for (RelFieldCollation collation : node.getCollations()) { - if (!inputSchema.getColumnDataType(collation.getFieldIndex()).supportsOrdering()) { - throw unsupported("Window ORDER BY"); + DataSchema.ColumnDataType dataType = inputSchema.getColumnDataType(collation.getFieldIndex()); + if (!dataType.supportsOrdering()) { + throw unsupported("Window ORDER BY", dataType); } } validateAggregateInputs(node.getAggCalls(), inputSchema); @@ -178,8 +187,12 @@ private static DataSchema.ColumnDataType getLogicalType(RexExpression expression throw new IllegalStateException("Unsupported aggregate operand: " + expression.getClass().getName()); } - private static QueryException unsupported(String operation) { - return new QueryException(QueryErrorCode.QUERY_PLANNING, - operation + " does not support raw VARIANT values; extract a typed path with variantGet first"); + private static QueryException unsupported(String operation, DataSchema.ColumnDataType dataType) { + // Name the actual unsupported type so the error is accurate for non-VARIANT opaque types (OBJECT, arrays, MAP), + // while preserving the raw-VARIANT wording and remediation guidance for the VARIANT case. + String message = dataType == DataSchema.ColumnDataType.VARIANT + ? operation + " does not support raw VARIANT values; extract a typed path with variantGet first" + : operation + " does not support " + dataType + " values"; + return new QueryException(QueryErrorCode.QUERY_PLANNING, message); } } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java index c558e15464db..ae5be0c5b003 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java @@ -141,6 +141,16 @@ public void testAllowsRawVariantCount() { aggregate("COUNT", false, VARIANT_SCHEMA).visit(VariantTypeValidationVisitor.INSTANCE, null); } + @Test + public void testRejectsRawVariantGroupByKey() { + // COUNT is raw-variant-independent, so the rejection must come from the VARIANT GROUP BY key itself. + AggregateNode node = aggregateGroupBy("COUNT", VARIANT_SCHEMA, List.of(0)); + QueryException exception = + Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("GROUP BY")); + Assert.assertTrue(exception.getMessage().contains("raw VARIANT")); + } + @Test public void testUsesLogicalTypeInsteadOfVariantStorageType() { DataSchema bytesSchema = @@ -225,6 +235,15 @@ private static AggregateNode aggregate(String functionName, boolean distinct, Da List.of(), AggregateNode.AggType.DIRECT, false, List.of(), 0); } + private static AggregateNode aggregateGroupBy(String functionName, DataSchema inputSchema, List groupKeys) { + ValueNode input = new ValueNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + RexExpression.FunctionCall aggCall = functionCall(functionName, false, new RexExpression.InputRef(0)); + DataSchema resultSchema = + new DataSchema(new String[]{"result"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.LONG}); + return new AggregateNode(0, resultSchema, PlanNode.NodeHint.EMPTY, List.of(input), List.of(aggCall), List.of(-1), + groupKeys, AggregateNode.AggType.DIRECT, false, List.of(), 0); + } + private static WindowNode window(String functionName) { return window(functionName, VARIANT_SCHEMA, List.of(), List.of()); } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java index be69c8b60a6b..d78a6404727a 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java @@ -28,6 +28,7 @@ import org.apache.calcite.rel.RelFieldCollation; import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.query.selection.SelectionOperatorUtils; import org.apache.pinot.query.planner.plannode.SortNode; import org.apache.pinot.query.runtime.blocks.MseBlock; @@ -75,9 +76,10 @@ public SortOperator(OpChainExecutionContext context, MultiStageOperator input, S // - Input is already sorted List collations = node.getCollations(); for (RelFieldCollation collation : collations) { - Preconditions.checkArgument( - _dataSchema.getColumnDataType(collation.getFieldIndex()).supportsOrdering(), - "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first"); + ColumnDataType dataType = _dataSchema.getColumnDataType(collation.getFieldIndex()); + Preconditions.checkArgument(dataType.supportsOrdering(), dataType == ColumnDataType.VARIANT + ? "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first" + : "ORDER BY does not support " + dataType + " values"); } if (collations.isEmpty() || input instanceof SortedMailboxReceiveOperator) { _priorityQueue = null; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java index a4b22e8c80eb..da51fabee35d 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperator.java @@ -57,9 +57,10 @@ public SortedMailboxReceiveOperator(OpChainExecutionContext context, MailboxRece _dataSchema = node.getDataSchema(); _collations = node.getCollations(); for (RelFieldCollation collation : _collations) { - Preconditions.checkArgument( - _dataSchema.getColumnDataType(collation.getFieldIndex()).supportsOrdering(), - "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first"); + DataSchema.ColumnDataType dataType = _dataSchema.getColumnDataType(collation.getFieldIndex()); + Preconditions.checkArgument(dataType.supportsOrdering(), dataType == DataSchema.ColumnDataType.VARIANT + ? "ORDER BY does not support raw VARIANT values; extract a typed path with variantGet first" + : "ORDER BY does not support " + dataType + " values"); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java index 808f117b83f4..3796fcbecd2a 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/FilterOperand.java @@ -116,7 +116,8 @@ public In(List children, DataSchema dataSchema, boolean isNotIn) _childOperands = new ArrayList<>(children.size()); for (RexExpression child : children) { TransformOperand operand = TransformOperandFactory.getTransformOperand(child, dataSchema); - Preconditions.checkArgument(operand.getResultType().supportsEquality(), + // Reject raw VARIANT operands only; preserve existing IN behavior for all other types. + Preconditions.checkArgument(operand.getResultType() != ColumnDataType.VARIANT, "Raw VARIANT values do not support IN; extract a typed path with variantGet first"); _childOperands.add(operand); } @@ -196,8 +197,10 @@ public Predicate(List operands, DataSchema dataSchema, IntPredica ColumnDataType lhsType = _lhs.getResultType(); ColumnDataType rhsType = _rhs.getResultType(); - Preconditions.checkArgument((lhsType == ColumnDataType.UNKNOWN || lhsType.supportsOrdering()) - && (rhsType == ColumnDataType.UNKNOWN || rhsType.supportsOrdering()), + // Reject raw VARIANT operands only; other non-orderable types (OBJECT, arrays, MAP) keep their existing + // best-effort comparison behavior. VARIANT is opaque because its PVAR byte encoding is not a canonical + // semantic ordering, so a comparison must extract a typed scalar first. + Preconditions.checkArgument(lhsType != ColumnDataType.VARIANT && rhsType != ColumnDataType.VARIANT, "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); if (lhsType == ColumnDataType.UNKNOWN || rhsType == ColumnDataType.UNKNOWN || lhsType == rhsType) { _requireCasting = false; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java index 8581f7b7d8af..dbee885f59fb 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/TransformOperandFactory.java @@ -62,7 +62,8 @@ private static TransformOperand getTransformOperand(RexExpression.FunctionCall f Preconditions.checkState(numOperands == 2, "%s takes 2 arguments, got: %s", functionCall.getFunctionName(), numOperands); for (RexExpression operand : operands) { - Preconditions.checkArgument(getResultType(operand, dataSchema).supportsEquality(), + // Reject raw VARIANT operands only; preserve existing distinct-from behavior for all other types. + Preconditions.checkArgument(getResultType(operand, dataSchema) != DataSchema.ColumnDataType.VARIANT, "Raw VARIANT values do not support comparison; extract a typed path with variantGet first"); } return new FunctionOperand(functionCall, dataSchema); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java index e8768419c4dd..5b1e146f7cd3 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java @@ -31,6 +31,8 @@ public class FilterOperandTest { new DataSchema(new String[]{"value"}, new ColumnDataType[]{ColumnDataType.INT}); private static final DataSchema VARIANT_SCHEMA = new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.VARIANT}); + private static final DataSchema OBJECT_SCHEMA = + new DataSchema(new String[]{"payload"}, new ColumnDataType[]{ColumnDataType.OBJECT}); private static final RexExpression NULL_LITERAL = new RexExpression.Literal(ColumnDataType.UNKNOWN, null); private static final List VARIANT_OPERANDS = List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(0)); @@ -70,6 +72,19 @@ public void testRawVariantNotInIsRejected() { Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support IN")); } + @Test + public void testNonVariantOpaqueTypesAreNotRejectedAsVariant() { + // OBJECT (and other non-orderable types) must keep their existing best-effort comparison/IN behavior rather + // than being rejected with the VARIANT-specific guard. Constructing the operands must not throw. + new FilterOperand.Predicate(VARIANT_OPERANDS, OBJECT_SCHEMA, value -> value == 0); + new FilterOperand.In(VARIANT_OPERANDS, OBJECT_SCHEMA, false); + for (String functionName : List.of("IS_DISTINCT_FROM", "isNotDistinctFrom")) { + RexExpression.FunctionCall functionCall = + new RexExpression.FunctionCall(ColumnDataType.BOOLEAN, functionName, VARIANT_OPERANDS); + TransformOperandFactory.getTransformOperand(functionCall, OBJECT_SCHEMA); + } + } + @Test public void testRawVariantDistinctFromIsRejected() { for (String functionName : List.of("IS_DISTINCT_FROM", "isNotDistinctFrom")) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java index b10be0ac202d..0334267bd699 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java @@ -1437,6 +1437,25 @@ static void validatePartialUpsertStrategies(TableConfig tableConfig, Schema sche Map partialUpsertStrategies = upsertConfig.getPartialUpsertStrategies(); String partialUpsertMergerClass = upsertConfig.getPartialUpsertMergerClass(); + // VARIANT columns support only the OVERWRITE partial-upsert strategy. A column that is not listed in + // partialUpsertStrategies is merged at runtime with defaultPartialUpsertStrategy, so validating only the listed + // entries would let a non-OVERWRITE default (or a custom merger class) silently apply INCREMENT/APPEND/UNION/IGNORE + // to a VARIANT byte envelope. Validate the effective strategy for every VARIANT column here. + for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) { + if (fieldSpec.getDataType() != DataType.VARIANT) { + continue; + } + String column = fieldSpec.getName(); + Preconditions.checkState(StringUtils.isBlank(partialUpsertMergerClass), + "VARIANT column supports only OVERWRITE partial-upsert strategy and cannot be merged by a custom " + + "partialUpsertMergerClass: %s", column); + UpsertConfig.Strategy effectiveStrategy = + partialUpsertStrategies != null && partialUpsertStrategies.containsKey(column) + ? partialUpsertStrategies.get(column) : upsertConfig.getDefaultPartialUpsertStrategy(); + Preconditions.checkState(effectiveStrategy == UpsertConfig.Strategy.OVERWRITE, + "VARIANT column supports only OVERWRITE partial-upsert strategy: %s", column); + } + // check if partialUpsertMergerClass is provided then partialUpsertStrategies should be empty if (StringUtils.isNotBlank(partialUpsertMergerClass)) { Preconditions.checkState(MapUtils.isEmpty(partialUpsertStrategies), @@ -1460,10 +1479,7 @@ static void validatePartialUpsertStrategies(TableConfig tableConfig, Schema sche FieldSpec fieldSpec = schema.getFieldSpecFor(column); Preconditions.checkState(fieldSpec != null, "Merger cannot be applied to non-existing column: %s", column); - if (fieldSpec.getDataType() == DataType.VARIANT) { - Preconditions.checkState(columnStrategy == UpsertConfig.Strategy.OVERWRITE, - "VARIANT column supports only OVERWRITE partial-upsert strategy: %s", column); - } + // VARIANT columns are validated up front against their effective strategy (default or explicit). if (columnStrategy == UpsertConfig.Strategy.INCREMENT) { Preconditions.checkState(fieldSpec.getDataType().getStoredType().isNumeric(), diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java index cab555d446be..240ff896b8fd 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java @@ -226,6 +226,63 @@ public void testOverwritePartialUpsertStrategyIsValid() { partialUpsertTable(UpsertConfig.Strategy.OVERWRITE), UPSERT_SCHEMA); } + @Test + public void testNonOverwriteDefaultPartialUpsertStrategyIsRejectedForUnlistedVariant() { + // The VARIANT column is not listed in partialUpsertStrategies, so at merge time it uses + // defaultPartialUpsertStrategy. A non-OVERWRITE default must still be rejected. + UpsertConfig upsertConfig = new UpsertConfig(UpsertConfig.Mode.PARTIAL); + upsertConfig.setComparisonColumn(ID_COLUMN); + upsertConfig.setPartialUpsertStrategies(Map.of(ID_COLUMN, UpsertConfig.Strategy.OVERWRITE)); + upsertConfig.setDefaultPartialUpsertStrategy(UpsertConfig.Strategy.UNION); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setUpsertConfig(upsertConfig) + .build(); + try { + TableConfigUtils.validatePartialUpsertStrategies(tableConfig, UPSERT_SCHEMA); + fail("Expected non-OVERWRITE default VARIANT partial-upsert strategy validation to fail"); + } catch (IllegalStateException e) { + assertTrue(e.getMessage() != null + && e.getMessage().contains("VARIANT column supports only OVERWRITE partial-upsert strategy"), + "Unexpected validation error: " + e.getMessage()); + } + } + + @Test + public void testDefaultOverwritePartialUpsertStrategyIsValidForUnlistedVariant() { + // The VARIANT column is unlisted and defaultPartialUpsertStrategy defaults to OVERWRITE, which is valid. + UpsertConfig upsertConfig = new UpsertConfig(UpsertConfig.Mode.PARTIAL); + upsertConfig.setComparisonColumn(ID_COLUMN); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setUpsertConfig(upsertConfig) + .build(); + TableConfigUtils.validatePartialUpsertStrategies(tableConfig, UPSERT_SCHEMA); + } + + @Test + public void testCustomPartialUpsertMergerIsRejectedForVariant() { + // A custom merger class cannot be validated against the OVERWRITE-only VARIANT contract, so it must be rejected. + UpsertConfig upsertConfig = new UpsertConfig(UpsertConfig.Mode.PARTIAL); + upsertConfig.setComparisonColumn(ID_COLUMN); + upsertConfig.setPartialUpsertMergerClass("com.example.CustomMerger"); + TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setUpsertConfig(upsertConfig) + .build(); + try { + TableConfigUtils.validatePartialUpsertStrategies(tableConfig, UPSERT_SCHEMA); + fail("Expected custom partialUpsertMergerClass with a VARIANT column to fail"); + } catch (IllegalStateException e) { + assertTrue(e.getMessage() != null && e.getMessage().contains("custom") + && e.getMessage().contains("VARIANT column supports only OVERWRITE partial-upsert strategy"), + "Unexpected validation error: " + e.getMessage()); + } + } + @Test public void testMetricsAggregationAndDefaultStarTreeAreRejected() { TableConfig aggregateMetricsTable = new TableConfigBuilder(TableType.OFFLINE) From a488814d553ef66d9f371323ffb7b2cf90abcc2e Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 11 Aug 2026 00:23:05 -0700 Subject: [PATCH 6/8] Fix VARIANT validation for partial metadata aggregation --- .../function/AggregationFunctionUtils.java | 16 +++++++++++++++- ...ataAndDictionaryAggregationPlanMakerTest.java | 15 +++++++++------ 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index ee61ce12dc6b..52925f880cd5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -400,6 +400,18 @@ public AggregationInfo(AggregationFunction[] functions, BaseProjectOperator p } } + private AggregationInfo(AggregationFunction[] functions, BaseProjectOperator projectOperator) { + _functions = functions; + _projectOperator = projectOperator; + _useStarTree = false; + } + + private static AggregationInfo forPartialProjection(AggregationFunction[] allFunctions, + BaseProjectOperator projectOperator, AggregationFunction[] projectionFunctions) { + validateRawVariantAggregationInputs(projectionFunctions, projectOperator); + return new AggregationInfo(allFunctions, projectOperator); + } + public AggregationFunction[] getFunctions() { return _functions; } @@ -470,12 +482,14 @@ public static AggregationInfo buildAggregationInfoWithoutStarTree(SegmentContext public static AggregationInfo buildAggregationInfoWithoutStarTree(SegmentContext segmentContext, QueryContext queryContext, AggregationFunction[] allFunctions, AggregationFunction[] projectionFunctions, BaseFilterOperator filterOperator) { + // This builder is public, so do not rely on its current caller having performed the star-tree validation first. + validateRawVariantIdentifierInputs(allFunctions, segmentContext, queryContext); Set expressionsToTransform = collectExpressionsToTransform(projectionFunctions, queryContext.getGroupByExpressions()); BaseProjectOperator projectOperator = new ProjectPlanNode(segmentContext, queryContext, expressionsToTransform, DocIdSetPlanNode.MAX_DOC_PER_CALL, filterOperator).run(); - return new AggregationInfo(allFunctions, projectOperator, false); + return AggregationInfo.forPartialProjection(allFunctions, projectOperator, projectionFunctions); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java index 8cd0278f7340..ecfb5d3d075d 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/MetadataAndDictionaryAggregationPlanMakerTest.java @@ -35,6 +35,7 @@ import org.apache.pinot.core.operator.query.GroupByOperator; import org.apache.pinot.core.operator.query.NonScanBasedAggregationOperator; import org.apache.pinot.core.operator.query.SelectionOnlyOperator; +import org.apache.pinot.core.operator.transform.TransformOperator; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.segment.local.data.manager.TableDataManager; @@ -359,13 +360,15 @@ public void testMinMaxOnNonNumericColumnFallsToScan() { @Test public void testResolvedOnlyColumnExcludedFromProjection() { QueryContext queryContext = QueryContextConverterUtils.getQueryContext( - "select max(metricCol), sum(intCol) from " + PREDICTABLE_TABLE_NAME); + "select max(metricCol), sum(add(intCol, intCol)) from " + PREDICTABLE_TABLE_NAME); Operator operator = PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(_predictableSegment), queryContext).run(); assertTrue(operator instanceof AggregationOperator); BaseProjectOperator projectOperator = ((AggregationOperator) operator).getChildOperators().get(0); - // Only intCol (used by the scanned sum) is projected; metricCol (used solely by the resolved max) is excluded. + // The expression forces a TransformOperator whose input contains only intCol; metricCol is resolved from metadata. + assertTrue(projectOperator instanceof TransformOperator, + projectOperator.getClass() + ": " + queryContext.getAggregationFunctions()[1].getInputExpressions()); assertEquals(projectOperator.getNumColumnsProjected(), 1); assertTrue(projectOperator.getSourceColumnContextMap().containsKey("intCol")); assertFalse(projectOperator.getSourceColumnContextMap().containsKey("metricCol")); @@ -373,12 +376,12 @@ public void testResolvedOnlyColumnExcludedFromProjection() { AggregationResultsBlock resultsBlock = (AggregationResultsBlock) operator.nextBlock(); List results = resultsBlock.getResults(); assertNotNull(results); - // max(metricCol) is resolved from the dictionary; sum(intCol) is scanned. - // intCol = i + 10 for i = 1..10 over the two row copies, so sum(intCol) = 2 * (11 + 12 + ... + 20) = 310. + // max(metricCol) is resolved from the dictionary; sum(add(intCol, intCol)) is scanned. + // intCol = i + 10 for i = 1..10 over the two row copies, so the sum is 4 * (11 + 12 + ... + 20) = 620. assertEquals(((Number) results.get(0)).doubleValue(), 10.0); - assertEquals(((Number) results.get(1)).doubleValue(), 310.0); + assertEquals(((Number) results.get(1)).doubleValue(), 620.0); - // All 20 docs are still scanned for the unresolved sum(intCol). + // All 20 docs are still scanned for the unresolved sum(add(intCol, intCol)). assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), 20); // With metricCol excluded, only 1 column is projected, so entries scanned post-filter is 20 docs * 1 column = 20 // (it would be 40 if the resolved-only metricCol were still projected). From d4bb30698e2a2fc5a1489e6320be21f07ffbc83b Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Thu, 13 Aug 2026 22:45:07 -0700 Subject: [PATCH 7/8] Harden VARIANT after upstream rebase Restore the parquet-common runtime dependency required by parquet-variant and validate grouping-set keys against the expanded repeat schema. Reuse the query cursor for JSON rendering, cover every Parquet Variant type, prevent timestamp-derived-name validation bypasses, and clarify the type-capability validation APIs. --- pinot-common/pom.xml | 8 +- .../pinot/common/utils/VariantUtils.java | 8 +- .../pinot/common/utils/VariantUtilsTest.java | 81 +++++++++++++++++++ .../utils/OrderByComparatorFactoryTest.java | 14 ++++ .../physical/PinotDispatchPlanner.java | 6 +- ...a => TypeCapabilityValidationVisitor.java} | 11 +-- ... TypeCapabilityValidationVisitorTest.java} | 56 ++++++++----- .../runtime/operator/AggregateOperator.java | 10 ++- .../operator/MultistageGroupByExecutor.java | 12 +-- .../runtime/operator/RepeatOperator.java | 5 ++ .../operator/WindowAggregateOperator.java | 4 +- .../operator/operands/VariantOperand.java | 2 +- .../operator/set/BinarySetOperator.java | 2 +- .../runtime/operator/set/SetOperator.java | 10 ++- .../runtime/operator/set/UnionOperator.java | 2 +- .../operator/AggregateOperatorTest.java | 22 +++++ .../runtime/operator/SortOperatorTest.java | 11 +++ .../SortedMailboxReceiveOperatorTest.java | 12 +++ .../operator/set/UnionOperatorTest.java | 6 ++ .../segment/local/utils/TableConfigUtils.java | 10 ++- .../VariantTableConfigValidationTest.java | 30 +++++++ 21 files changed, 263 insertions(+), 59 deletions(-) rename pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/{VariantTypeValidationVisitor.java => TypeCapabilityValidationVisitor.java} (93%) rename pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/{VariantTypeValidationVisitorTest.java => TypeCapabilityValidationVisitorTest.java} (84%) diff --git a/pinot-common/pom.xml b/pinot-common/pom.xml index f9789fc77bfb..88510efd4d88 100644 --- a/pinot-common/pom.xml +++ b/pinot-common/pom.xml @@ -203,17 +203,13 @@ org.apache.parquet parquet-variant - + org.apache.parquet parquet-column - - org.apache.parquet - parquet-common - diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java index 9b7d66e47b69..eb393401e6be 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java @@ -523,10 +523,16 @@ public static String variantTypeOf(@Nullable byte[] envelope, VariantPath path, /// Renders the Variant value as canonical JSON text without constructing a JSON tree. @Nullable public static String variantToJson(@Nullable byte[] envelope) { + return variantToJson(envelope, new ReusableResult()); + } + + /// Allocation-reduced form of [#variantToJson(byte[])] when the caller retains the supplied result between rows. + @Nullable + public static String variantToJson(@Nullable byte[] envelope, ReusableResult result) { + Objects.requireNonNull(result, "result must not be null"); if (isSqlNull(envelope)) { return null; } - ReusableResult result = new ReusableResult(); Cursor cursor = result._cursor; cursor.navigate(envelope, ROOT_PATH); return variantToJson(cursor.asVariant()); diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java index 4d1fe9fdc9ea..02774a9616f3 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java @@ -22,9 +22,12 @@ import java.nio.ByteBuffer; import java.sql.Timestamp; import java.util.Arrays; +import java.util.EnumSet; import java.util.List; +import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.parquet.variant.Variant; import org.apache.parquet.variant.VariantBuilder; @@ -36,6 +39,7 @@ import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.VariantEnvelope; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; @@ -379,6 +383,10 @@ public void testSqlNullAndVariantNullRemainDistinct() { public void testParseRenderAndLogicalDataSchema() { byte[] variant = VariantUtils.parseJsonToVariant("{\"a\":[1,true,null],\"b\":\"text\"}"); assertEquals(VariantUtils.variantToJson(variant), "{\"a\":[1,true,null],\"b\":\"text\"}"); + ReusableResult reusableResult = new ReusableResult(); + assertEquals(VariantUtils.variantToJson(variant, reusableResult), "{\"a\":[1,true,null],\"b\":\"text\"}"); + assertEquals(VariantUtils.variantToJson(VariantUtils.parseJsonToVariant("[2,3]"), reusableResult), "[2,3]"); + assertNull(VariantUtils.variantToJson(new byte[0], reusableResult)); assertNull(VariantUtils.tryParseJsonToVariant("{not-json")); assertThrows(IllegalArgumentException.class, () -> VariantUtils.parseJsonToVariant("{not-json")); @@ -389,6 +397,68 @@ public void testParseRenderAndLogicalDataSchema() { "{\"a\":[1,true,null],\"b\":\"text\"}"); } + @DataProvider(name = "variantJsonRenderingCases") + public Object[][] variantJsonRenderingCases() { + return new Object[][]{ + jsonRenderingCase(Variant.Type.OBJECT, VariantUtils.parseJsonToVariant("{\"a\":1}"), "{\"a\":1}"), + jsonRenderingCase(Variant.Type.ARRAY, VariantUtils.parseJsonToVariant("[1,true]"), "[1,true]"), + jsonRenderingCase(Variant.Type.NULL, "null", VariantBuilder::appendNull), + jsonRenderingCase(Variant.Type.BOOLEAN, "true", builder -> builder.appendBoolean(true)), + jsonRenderingCase(Variant.Type.BYTE, "-8", builder -> builder.appendByte((byte) -8)), + jsonRenderingCase(Variant.Type.SHORT, "32000", builder -> builder.appendShort((short) 32_000)), + jsonRenderingCase(Variant.Type.INT, "-123456", builder -> builder.appendInt(-123_456)), + jsonRenderingCase(Variant.Type.LONG, "9876543210", builder -> builder.appendLong(9_876_543_210L)), + jsonRenderingCase(Variant.Type.STRING, "\"text\"", builder -> builder.appendString("text")), + jsonRenderingCase(Variant.Type.DOUBLE, "-123.5", builder -> builder.appendDouble(-123.5D)), + jsonRenderingCase(Variant.Type.DECIMAL4, "12.34", builder -> builder.appendDecimal(new BigDecimal("12.34"))), + jsonRenderingCase(Variant.Type.DECIMAL8, "1234567890.12", + builder -> builder.appendDecimal(new BigDecimal("1234567890.12"))), + jsonRenderingCase(Variant.Type.DECIMAL16, "12345678901234567890.1234", + builder -> builder.appendDecimal(new BigDecimal("12345678901234567890.1234"))), + jsonRenderingCase(Variant.Type.DATE, "\"1970-01-02\"", builder -> builder.appendDate(1)), + jsonRenderingCase(Variant.Type.TIMESTAMP_TZ, "\"1970-01-01T00:00:01.234567Z\"", + builder -> builder.appendTimestampTz(1_234_567L)), + jsonRenderingCase(Variant.Type.TIMESTAMP_NTZ, "\"1970-01-01T00:00:01.234567\"", + builder -> builder.appendTimestampNtz(1_234_567L)), + jsonRenderingCase(Variant.Type.FLOAT, "1.25", builder -> builder.appendFloat(1.25F)), + jsonRenderingCase(Variant.Type.BINARY, "\"AAH/Kg==\"", + builder -> builder.appendBinary(ByteBuffer.wrap(new byte[]{0, 1, -1, 42}))), + jsonRenderingCase(Variant.Type.TIME, "\"01:02:03.004005\"", + builder -> builder.appendTime(3_723_004_005L)), + jsonRenderingCase(Variant.Type.TIMESTAMP_NANOS_TZ, "\"1970-01-01T00:00:01.234567891Z\"", + builder -> builder.appendTimestampNanosTz(1_234_567_891L)), + jsonRenderingCase(Variant.Type.TIMESTAMP_NANOS_NTZ, "\"1970-01-01T00:00:01.234567891\"", + builder -> builder.appendTimestampNanosNtz(1_234_567_891L)), + jsonRenderingCase(Variant.Type.UUID, "\"00112233-4455-6677-8899-aabbccddeeff\"", + builder -> builder.appendUUID(UUID.fromString("00112233-4455-6677-8899-aabbccddeeff"))) + }; + } + + @Test(dataProvider = "variantJsonRenderingCases") + public void testVariantToJsonForEveryParquetType(Variant.Type expectedType, byte[] envelope, String expectedJson) { + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + assertEquals(new Variant(decoded.getValue(), decoded.getMetadata()).getType(), expectedType); + assertEquals(VariantUtils.variantToJson(envelope), expectedJson); + assertEquals(VariantUtils.variantToJson(envelope, new ReusableResult()), expectedJson); + } + + @Test + public void testVariantToJsonCoversEveryParquetTypeAndRejectsMalformedEnvelope() { + Set coveredTypes = EnumSet.noneOf(Variant.Type.class); + for (Object[] testCase : variantJsonRenderingCases()) { + coveredTypes.add((Variant.Type) testCase[0]); + } + assertEquals(coveredTypes, EnumSet.allOf(Variant.Type.class)); + + byte[] valid = VariantUtils.parseJsonToVariant("{\"a\":1}"); + byte[] badMagic = Arrays.copyOf(valid, valid.length); + badMagic[0] = 0; + ReusableResult reusableResult = new ReusableResult(); + assertThrows(IllegalArgumentException.class, () -> VariantUtils.variantToJson(badMagic)); + assertThrows(IllegalArgumentException.class, () -> VariantUtils.variantToJson(badMagic, reusableResult)); + assertEquals(VariantUtils.variantToJson(valid, reusableResult), "{\"a\":1}"); + } + @Test public void testJsonIntegerBoundsAndBigIntegerFallback() { byte[] variant = VariantUtils.parseJsonToVariant( @@ -544,6 +614,17 @@ private static byte[] encode(VariantBuilder builder) { return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); } + private static Object[] jsonRenderingCase(Variant.Type type, String expectedJson, + Consumer appender) { + VariantBuilder builder = new VariantBuilder(); + appender.accept(builder); + return jsonRenderingCase(type, encode(builder), expectedJson); + } + + private static Object[] jsonRenderingCase(Variant.Type type, byte[] envelope, String expectedJson) { + return new Object[]{type, envelope, expectedJson}; + } + private static int readBigEndianInt(byte[] bytes, int offset) { return Byte.toUnsignedInt(bytes[offset]) << 24 | Byte.toUnsignedInt(bytes[offset + 1]) << 16 diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java index 1ee8e7629962..862fe6804cc8 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java @@ -123,4 +123,18 @@ public void testRejectsRawVariant() { ENABLE_NULL_HANDLING)); Assert.assertTrue(exception.getMessage().contains("ORDER BY does not support raw VARIANT")); } + + @Test + public void testNamesUnsupportedNonVariantType() { + List orderBys = + List.of(new OrderByExpressionContext(COLUMN1, ASC, NULLS_LAST)); + ColumnContext columnContext = Mockito.mock(ColumnContext.class); + Mockito.when(columnContext.isSingleValue()).thenReturn(true); + Mockito.when(columnContext.getDataType()).thenReturn(DataType.MAP); + + BadQueryRequestException exception = Assert.expectThrows(BadQueryRequestException.class, + () -> OrderByComparatorFactory.getComparator(orderBys, new ColumnContext[]{columnContext}, + ENABLE_NULL_HANDLING)); + Assert.assertTrue(exception.getMessage().contains("ORDER BY does not support MAP values")); + } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java index 227089558869..58e0b0fd7726 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/PinotDispatchPlanner.java @@ -37,7 +37,7 @@ import org.apache.pinot.query.planner.plannode.TableScanNode; import org.apache.pinot.query.planner.plannode.ValueNode; import org.apache.pinot.query.planner.validation.ArrayToMvValidationVisitor; -import org.apache.pinot.query.planner.validation.VariantTypeValidationVisitor; +import org.apache.pinot.query.planner.validation.TypeCapabilityValidationVisitor; import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.query.routing.WorkerMetadata; @@ -157,13 +157,13 @@ private static void trackEmptyLeafStages(DispatchablePlanContext context) { private void runValidations(PlanFragment planFragment, DispatchablePlanContext context) { PlanNode rootPlanNode = planFragment.getFragmentRoot(); if (rootPlanNode.getStageId() == 0) { - VariantTypeValidationVisitor.validateResultSchema(rootPlanNode.getDataSchema(), + TypeCapabilityValidationVisitor.validateResultSchema(rootPlanNode.getDataSchema(), context.getPlannerContext().getEnvConfig().isNullHandlingEnabled()); } boolean isIntermediateStage = context.getDispatchablePlanMetadataMap().get(rootPlanNode.getStageId()).getScannedTables().isEmpty(); rootPlanNode.visit(ArrayToMvValidationVisitor.INSTANCE, isIntermediateStage); - rootPlanNode.visit(VariantTypeValidationVisitor.INSTANCE, null); + rootPlanNode.visit(TypeCapabilityValidationVisitor.INSTANCE, null); for (PlanFragment child : planFragment.getChildren()) { runValidations(child, context); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitor.java similarity index 93% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java rename to pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitor.java index ad1178af1333..b6b216265ab9 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitor.java @@ -34,12 +34,13 @@ import org.apache.pinot.spi.exception.QueryException; -/// Rejects operations that would otherwise assign physical byte ordering, equality, or hashing semantics to a raw -/// VARIANT value. The visitor has no mutable state and is thread-safe, so callers may share {@link #INSTANCE}. -public final class VariantTypeValidationVisitor extends PlanNodeVisitor.DepthFirstVisitor { - public static final VariantTypeValidationVisitor INSTANCE = new VariantTypeValidationVisitor(); +/// Validates that each logical input type supports the capabilities required by its operation, including equality, +/// hashing, ordering, aggregation, and lossless result projection. The visitor has no mutable state and is thread-safe, +/// so callers may share {@link #INSTANCE}. +public final class TypeCapabilityValidationVisitor extends PlanNodeVisitor.DepthFirstVisitor { + public static final TypeCapabilityValidationVisitor INSTANCE = new TypeCapabilityValidationVisitor(); - private VariantTypeValidationVisitor() { + private TypeCapabilityValidationVisitor() { } @Override diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitorTest.java similarity index 84% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java rename to pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitorTest.java index ae5be0c5b003..9a891f20727f 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/VariantTypeValidationVisitorTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitorTest.java @@ -35,7 +35,7 @@ import org.testng.annotations.Test; -public class VariantTypeValidationVisitorTest { +public class TypeCapabilityValidationVisitorTest { private static final DataSchema VARIANT_SCHEMA = new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.VARIANT}); private static final DataSchema TYPED_EXTRACTION_SCHEMA = @@ -47,18 +47,30 @@ public void testRejectsVariantOrderBy() { List.of(new RelFieldCollation(0)), 10, 0); QueryException exception = - Assert.expectThrows(QueryException.class, () -> sortNode.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.expectThrows(QueryException.class, () -> sortNode.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("ORDER BY")); } + @Test + public void testNamesUnsupportedNonVariantOrderByType() { + DataSchema objectSchema = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.OBJECT}); + SortNode sortNode = new SortNode(0, objectSchema, PlanNode.NodeHint.EMPTY, List.of(), + List.of(new RelFieldCollation(0)), 10, 0); + + QueryException exception = + Assert.expectThrows(QueryException.class, () -> sortNode.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("ORDER BY does not support OBJECT values")); + } + @Test public void testRawVariantProjectionRequiresNullHandling() { QueryException exception = Assert.expectThrows(QueryException.class, - () -> VariantTypeValidationVisitor.validateResultSchema(VARIANT_SCHEMA, false)); + () -> TypeCapabilityValidationVisitor.validateResultSchema(VARIANT_SCHEMA, false)); Assert.assertTrue(exception.getMessage().contains("requires query null handling")); - VariantTypeValidationVisitor.validateResultSchema(VARIANT_SCHEMA, true); - VariantTypeValidationVisitor.validateResultSchema(TYPED_EXTRACTION_SCHEMA, false); + TypeCapabilityValidationVisitor.validateResultSchema(VARIANT_SCHEMA, true); + TypeCapabilityValidationVisitor.validateResultSchema(TYPED_EXTRACTION_SCHEMA, false); } @Test @@ -72,14 +84,14 @@ public void testRejectsEqualityDependentSetOperations() { for (SetOpNode node : unsupportedNodes) { QueryException exception = - Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.expectThrows(QueryException.class, () -> node.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("raw VARIANT")); } } @Test public void testAllowsVariantUnionAll() { - setOp(SetOpNode.SetOpType.UNION, true).visit(VariantTypeValidationVisitor.INSTANCE, null); + setOp(SetOpNode.SetOpType.UNION, true).visit(TypeCapabilityValidationVisitor.INSTANCE, null); } @Test @@ -87,12 +99,12 @@ public void testRejectsRawVariantJoinKeysForEveryStrategy() { for (JoinNode.JoinStrategy strategy : JoinNode.JoinStrategy.values()) { JoinNode leftVariant = join(strategy, VARIANT_SCHEMA, TYPED_EXTRACTION_SCHEMA); QueryException exception = Assert.expectThrows(QueryException.class, - () -> leftVariant.visit(VariantTypeValidationVisitor.INSTANCE, null)); + () -> leftVariant.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("JOIN keys")); JoinNode rightVariant = join(strategy, TYPED_EXTRACTION_SCHEMA, VARIANT_SCHEMA); exception = Assert.expectThrows(QueryException.class, - () -> rightVariant.visit(VariantTypeValidationVisitor.INSTANCE, null)); + () -> rightVariant.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("JOIN keys")); } } @@ -104,7 +116,7 @@ public void testRejectsRawVariantAsofMatchKeys() { DataSchema rightTyped = new DataSchema(new String[]{"key", "match"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.LONG}); QueryException exception = Assert.expectThrows(QueryException.class, - () -> asofJoin(leftVariant, rightTyped).visit(VariantTypeValidationVisitor.INSTANCE, null)); + () -> asofJoin(leftVariant, rightTyped).visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); DataSchema leftTyped = new DataSchema(new String[]{"key", "match"}, @@ -112,7 +124,7 @@ public void testRejectsRawVariantAsofMatchKeys() { DataSchema rightVariant = new DataSchema(new String[]{"key", "match"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.VARIANT}); exception = Assert.expectThrows(QueryException.class, - () -> asofJoin(leftTyped, rightVariant).visit(VariantTypeValidationVisitor.INSTANCE, null)); + () -> asofJoin(leftTyped, rightVariant).visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); } @@ -121,7 +133,7 @@ public void testRejectsAggregatesThatConsumeRawVariant() { for (String functionName : List.of("SUM", "ANYVALUE", "DISTINCTCOUNTHLL")) { AggregateNode node = aggregate(functionName, false, VARIANT_SCHEMA); QueryException exception = - Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.expectThrows(QueryException.class, () -> node.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Aggregate function " + functionName)); Assert.assertTrue(exception.getMessage().contains("variantGet")); } @@ -132,13 +144,13 @@ public void testRejectsDistinctCountOfRawVariant() { AggregateNode node = aggregate("COUNT", true, VARIANT_SCHEMA); QueryException exception = - Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.expectThrows(QueryException.class, () -> node.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Aggregate function COUNT")); } @Test public void testAllowsRawVariantCount() { - aggregate("COUNT", false, VARIANT_SCHEMA).visit(VariantTypeValidationVisitor.INSTANCE, null); + aggregate("COUNT", false, VARIANT_SCHEMA).visit(TypeCapabilityValidationVisitor.INSTANCE, null); } @Test @@ -146,7 +158,7 @@ public void testRejectsRawVariantGroupByKey() { // COUNT is raw-variant-independent, so the rejection must come from the VARIANT GROUP BY key itself. AggregateNode node = aggregateGroupBy("COUNT", VARIANT_SCHEMA, List.of(0)); QueryException exception = - Assert.expectThrows(QueryException.class, () -> node.visit(VariantTypeValidationVisitor.INSTANCE, null)); + Assert.expectThrows(QueryException.class, () -> node.visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("GROUP BY")); Assert.assertTrue(exception.getMessage().contains("raw VARIANT")); } @@ -155,7 +167,7 @@ public void testRejectsRawVariantGroupByKey() { public void testUsesLogicalTypeInsteadOfVariantStorageType() { DataSchema bytesSchema = new DataSchema(new String[]{"bytes"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BYTES}); - aggregate("ANYVALUE", false, bytesSchema).visit(VariantTypeValidationVisitor.INSTANCE, null); + aggregate("ANYVALUE", false, bytesSchema).visit(TypeCapabilityValidationVisitor.INSTANCE, null); } @Test @@ -165,35 +177,35 @@ public void testAllowsAggregateOverTypedVariantExtraction() { List.of(new RexExpression.InputRef(0))); AggregateNode node = aggregate("MINSTRING", false, VARIANT_SCHEMA, typedExtraction); - node.visit(VariantTypeValidationVisitor.INSTANCE, null); + node.visit(TypeCapabilityValidationVisitor.INSTANCE, null); } @Test public void testValidatesWindowAggregateInputs() { QueryException exception = Assert.expectThrows(QueryException.class, - () -> window("SUM").visit(VariantTypeValidationVisitor.INSTANCE, null)); + () -> window("SUM").visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Aggregate function SUM")); - window("COUNT").visit(VariantTypeValidationVisitor.INSTANCE, null); + window("COUNT").visit(TypeCapabilityValidationVisitor.INSTANCE, null); } @Test public void testRejectsRawVariantWindowKeys() { QueryException exception = Assert.expectThrows(QueryException.class, () -> window("COUNT", VARIANT_SCHEMA, List.of(0), List.of()) - .visit(VariantTypeValidationVisitor.INSTANCE, null)); + .visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Window PARTITION BY")); exception = Assert.expectThrows(QueryException.class, () -> window("COUNT", VARIANT_SCHEMA, List.of(), List.of(new RelFieldCollation(0))) - .visit(VariantTypeValidationVisitor.INSTANCE, null)); + .visit(TypeCapabilityValidationVisitor.INSTANCE, null)); Assert.assertTrue(exception.getMessage().contains("Window ORDER BY")); } @Test public void testAllowsWindowKeysOverTypedVariantExtraction() { window("COUNT", TYPED_EXTRACTION_SCHEMA, List.of(0), List.of(new RelFieldCollation(0))) - .visit(VariantTypeValidationVisitor.INSTANCE, null); + .visit(TypeCapabilityValidationVisitor.INSTANCE, null); } private static SetOpNode setOp(SetOpNode.SetOpType setOpType, boolean all) { diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java index 96c2a7bf397c..50dfc475f4cb 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AggregateOperator.java @@ -48,7 +48,7 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.PlanNode; -import org.apache.pinot.query.planner.validation.VariantTypeValidationVisitor; +import org.apache.pinot.query.planner.validation.TypeCapabilityValidationVisitor; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; import org.apache.pinot.query.runtime.operator.utils.SortUtils; @@ -92,8 +92,10 @@ public class AggregateOperator extends MultiStageOperator { public AggregateOperator(OpChainExecutionContext context, MultiStageOperator input, AggregateNode node) { super(context); - if (node.getInputs().size() == 1) { - VariantTypeValidationVisitor.validateAggregateInputs(node, node.getInputs().get(0).getDataSchema()); + DataSchema inputSchema = input instanceof RepeatOperator ? ((RepeatOperator) input).getResultSchema() + : node.getInputs().size() == 1 ? node.getInputs().get(0).getDataSchema() : null; + if (inputSchema != null) { + TypeCapabilityValidationVisitor.validateAggregateInputs(node, inputSchema); } _resultSchema = node.getDataSchema(); _aggFunctions = getAggFunctions(node.getAggCalls()); @@ -147,7 +149,7 @@ public AggregateOperator(OpChainExecutionContext context, MultiStageOperator inp } else { _groupByExecutor = new MultistageGroupByExecutor(getGroupKeyIds(groupKeys), _aggFunctions, filterArgIds, maxFilterArgId, aggType, - leafReturnFinalResult, _resultSchema, context.getOpChainMetadata(), node.getNodeHint()); + leafReturnFinalResult, inputSchema, _resultSchema, context.getOpChainMetadata(), node.getNodeHint()); _aggregationExecutor = null; } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java index 9ea55929b13c..d8216a5bc337 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultistageGroupByExecutor.java @@ -73,8 +73,8 @@ public class MultistageGroupByExecutor { private final GroupIdGenerator _groupIdGenerator; public MultistageGroupByExecutor(int[] groupKeyIds, AggregationFunction[] aggFunctions, int[] filterArgIds, - int maxFilterArgId, AggType aggType, boolean leafReturnFinalResult, DataSchema resultSchema, - Map opChainMetadata, @Nullable PlanNode.NodeHint nodeHint) { + int maxFilterArgId, AggType aggType, boolean leafReturnFinalResult, @Nullable DataSchema inputSchema, + DataSchema resultSchema, Map opChainMetadata, @Nullable PlanNode.NodeHint nodeHint) { _groupKeyIds = groupKeyIds; _aggFunctions = aggFunctions; _filterArgIds = filterArgIds; @@ -83,10 +83,12 @@ public MultistageGroupByExecutor(int[] groupKeyIds, AggregationFunction[] aggFun _leafReturnFinalResult = leafReturnFinalResult; _resultSchema = resultSchema; for (int i = 0; i < groupKeyIds.length; i++) { - ColumnDataType dataType = resultSchema.getColumnDataType(i); + ColumnDataType dataType = inputSchema != null ? inputSchema.getColumnDataType(groupKeyIds[i]) + : resultSchema.getColumnDataType(i); if (!dataType.supportsEquality() || !dataType.supportsHashing()) { - throw new IllegalArgumentException( - "Raw VARIANT values do not support GROUP BY; extract a typed path with variantGet first"); + throw new IllegalArgumentException(dataType == ColumnDataType.VARIANT + ? "Raw VARIANT values do not support GROUP BY; extract a typed path with variantGet first" + : "GROUP BY does not support " + dataType + " values"); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java index 5045b3b8d074..50a688b0fa08 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/RepeatOperator.java @@ -107,6 +107,11 @@ public Type getOperatorType() { return Type.REPEAT; } + /// Returns the expanded schema consumed by the downstream aggregate operator. + public DataSchema getResultSchema() { + return _resultSchema; + } + @Override protected Logger logger() { return LOGGER; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java index adb80f0df67c..92694e691107 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java @@ -34,7 +34,7 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.WindowNode; -import org.apache.pinot.query.planner.validation.VariantTypeValidationVisitor; +import org.apache.pinot.query.planner.validation.TypeCapabilityValidationVisitor; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.blocks.RowHeapDataBlock; import org.apache.pinot.query.runtime.operator.utils.AggregationUtils; @@ -109,7 +109,7 @@ public class WindowAggregateOperator extends MultiStageOperator { public WindowAggregateOperator(OpChainExecutionContext context, MultiStageOperator input, DataSchema inputSchema, WindowNode node) { super(context); - VariantTypeValidationVisitor.validateWindowInputs(node, inputSchema); + TypeCapabilityValidationVisitor.validateWindowInputs(node, inputSchema); _input = input; _resultSchema = node.getDataSchema(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java index 71f7c8e0425e..c21765cc06b1 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java @@ -132,7 +132,7 @@ public Object apply(List row) { externalResult = VariantUtils.variantTypeOf(variant, _path, _reusableResult); break; case TO_JSON: - externalResult = VariantUtils.variantToJson(variant); + externalResult = VariantUtils.variantToJson(variant, _reusableResult); break; default: throw new IllegalStateException("Unhandled Variant operation: " + _operation); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java index d462cf080f2c..bb3d5dc9f8f7 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/BinarySetOperator.java @@ -45,7 +45,7 @@ public BinarySetOperator(OpChainExecutionContext opChainExecutionContext, DataSchema dataSchema) { super(opChainExecutionContext, inputOperators, dataSchema); Preconditions.checkArgument(inputOperators.size() == 2, "Binary set operator should have 2 inputs"); - validateEqualitySupported(dataSchema, "INTERSECT/EXCEPT"); + validateEqualityAndHashingSupported(dataSchema, "INTERSECT/EXCEPT"); _leftChildOperator = inputOperators.get(0); _rightChildOperator = inputOperators.get(1); _rightRowSet = HashMultiset.create(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java index 84f64d1fca35..19d930426ce9 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/set/SetOperator.java @@ -18,7 +18,6 @@ */ package org.apache.pinot.query.runtime.operator.set; -import com.google.common.base.Preconditions; import java.util.List; import org.apache.pinot.common.datatable.StatMap; import org.apache.pinot.common.utils.DataSchema; @@ -41,10 +40,13 @@ public SetOperator(OpChainExecutionContext opChainExecutionContext, List inputOperators, DataSchema dataSchema) { super(opChainExecutionContext, inputOperators, dataSchema); - validateEqualitySupported(dataSchema, "UNION DISTINCT"); + validateEqualityAndHashingSupported(dataSchema, "UNION DISTINCT"); } @Override diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java index 1b298c967f99..0d5f0286fda4 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/AggregateOperatorTest.java @@ -134,6 +134,28 @@ public void testAllowsRawVariantCountAtRuntime() { Assert.assertNotNull(new AggregateOperator(OperatorTestUtil.getTracingContext(), _input, aggregateNode)); } + @Test + public void testGroupingSetValidationUsesExpandedInputSchema() { + DataSchema inputSchema = + new DataSchema(new String[]{"dimension", "metric"}, new ColumnDataType[]{STRING, INT}); + ValueNode inputPlanNode = new ValueNode(0, inputSchema, PlanNode.NodeHint.EMPTY, List.of(), List.of()); + DataSchema expandedInputSchema = new DataSchema( + new String[]{"dimension", "metric", "$groupingSetKey$0", "$groupingId"}, + new ColumnDataType[]{STRING, INT, STRING, INT}); + RepeatOperator repeatOperator = new RepeatOperator(OperatorTestUtil.getTracingContext(), _input, new int[]{0}, + List.of(List.of(0), List.of()), expandedInputSchema); + RexExpression.FunctionCall sum = + new RexExpression.FunctionCall(ColumnDataType.DOUBLE, SqlKind.SUM.name(), + List.of(new RexExpression.InputRef(1))); + DataSchema resultSchema = + new DataSchema(new String[]{"dimension", "sum"}, new ColumnDataType[]{STRING, DOUBLE}); + AggregateNode rewrittenAggregateNode = new AggregateNode(0, resultSchema, PlanNode.NodeHint.EMPTY, + List.of(inputPlanNode), List.of(sum), List.of(-1), List.of(2, 3), AggType.DIRECT, false, List.of(), 0); + + Assert.assertNotNull( + new AggregateOperator(OperatorTestUtil.getTracingContext(), repeatOperator, rewrittenAggregateNode)); + } + @Test public void shouldHandleEndOfStreamBlockWithNoOtherInputs() { // Given: diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java index 6ce55dc48e78..128cd6e84f74 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortOperatorTest.java @@ -511,6 +511,17 @@ public void shouldRejectRawVariantCollation() { assertTrue(exception.getMessage().contains("ORDER BY does not support raw VARIANT")); } + @Test + public void shouldNameUnsupportedNonVariantCollationType() { + DataSchema schema = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.OBJECT}); + List collations = List.of(new RelFieldCollation(0)); + + IllegalArgumentException exception = + expectThrows(IllegalArgumentException.class, () -> getOperator(schema, collations)); + assertTrue(exception.getMessage().contains("ORDER BY does not support OBJECT values")); + } + private SortOperator getOperator(DataSchema schema, List collations, int fetch, int offset) { return new SortOperator(OperatorTestUtil.getTracingContext(), _input, new SortNode(-1, schema, PlanNode.NodeHint.EMPTY, List.of(), collations, fetch, offset)); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java index 83fd87cc6a00..0656bb837b66 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/SortedMailboxReceiveOperatorTest.java @@ -126,6 +126,18 @@ public void shouldRejectRawVariantCollation() { assertTrue(exception.getMessage().contains("ORDER BY does not support raw VARIANT")); } + @Test + public void shouldNameUnsupportedNonVariantCollationType() { + when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); + DataSchema objectSchema = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.OBJECT}); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, + () -> getOperator(_stageMetadata1, RelDistribution.Type.SINGLETON, objectSchema, FIELD_COLLATIONS, + Long.MAX_VALUE)); + assertTrue(exception.getMessage().contains("ORDER BY does not support OBJECT values")); + } + @Test public void shouldTimeout() { when(_mailboxService.getReceivingMailbox(eq(MAILBOX_ID_1))).thenReturn(_mailbox1); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java index 4f295b7754c7..6e7ddb2790c2 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/set/UnionOperatorTest.java @@ -126,5 +126,11 @@ public void testVariantSetOperationValidation() { Assert.assertTrue(exception.getMessage().contains("UNION DISTINCT does not support raw VARIANT")); new UnionAllOperator(OperatorTestUtil.getTracingContext(), inputs, schema); + + DataSchema objectSchema = + new DataSchema(new String[]{"payload"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.OBJECT}); + exception = Assert.expectThrows(IllegalArgumentException.class, + () -> new UnionOperator(OperatorTestUtil.getTracingContext(), inputs, objectSchema)); + Assert.assertTrue(exception.getMessage().contains("UNION DISTINCT does not support OBJECT values")); } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java index 0334267bd699..83671d9261f8 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java @@ -2049,7 +2049,9 @@ private static void validateIndexingConfigAndFieldConfigListCompatibility(Indexi /// `timestampIndexColumns` holds the TIMESTAMP-index derived columns (e.g. `$ts$DAY`) declared via /// [TimestampConfig#getGranularities()]. These are materialized as dictionary-encoded single-value TIMESTAMP /// columns at segment generation time (see [TimestampIndexUtils#applyTimestampIndex(TableConfig, Schema)]), so - /// they are absent from the schema at config-validation time and are accepted here without a schema lookup. + /// they can be absent from the schema at config-validation time and are accepted without a schema lookup. When a + /// declared derived name is already present in the schema, it is validated normally so a user column with a + /// colliding name cannot bypass type, encoding, or cardinality checks. private static void validateStarTreeIndexConfigs(List starTreeIndexConfigs, Map indexConfigsMap, Schema schema, Set timestampIndexColumns) { Set dimensionColumns = new HashSet<>(); @@ -2058,7 +2060,7 @@ private static void validateStarTreeIndexConfigs(List starT List dimensionsSplitOrder = starTreeIndexConfig.getDimensionsSplitOrder(); assert CollectionUtils.isNotEmpty(dimensionsSplitOrder); for (String dimension : dimensionsSplitOrder) { - if (timestampIndexColumns.contains(dimension)) { + if (timestampIndexColumns.contains(dimension) && schema.getFieldSpecFor(dimension) == null) { dimensionColumns.add(dimension); continue; } @@ -2148,7 +2150,7 @@ private static void validateStarTreeIndexConfigs(List starT } for (String column : Iterables.concat(dimensionColumns, aggregatedColumns)) { - if (timestampIndexColumns.contains(column)) { + if (timestampIndexColumns.contains(column) && schema.getFieldSpecFor(column) == null) { continue; } FieldSpec fieldSpec = schema.getFieldSpecFor(column); @@ -2161,7 +2163,7 @@ private static void validateStarTreeIndexConfigs(List starT } for (String column : dimensionColumns) { - if (timestampIndexColumns.contains(column)) { + if (timestampIndexColumns.contains(column) && schema.getFieldSpecFor(column) == null) { continue; } FieldSpec fieldSpec = schema.getFieldSpecFor(column); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java index 240ff896b8fd..a8f0067deb28 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java @@ -31,6 +31,8 @@ import org.apache.pinot.spi.config.table.StarTreeIndexConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.TimestampConfig; +import org.apache.pinot.spi.config.table.TimestampIndexGranularity; import org.apache.pinot.spi.config.table.UpsertConfig; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; @@ -311,6 +313,34 @@ public void testExplicitStarTreeIsRejected() { assertInvalid(tableConfig, SCHEMA, "Star-tree index cannot be created on VARIANT column"); } + @Test + public void testTimestampDerivedNameCannotHideVariantFromStarTreeValidation() { + String timestampColumn = "eventTime"; + String collidingVariantColumn = "$eventTime$DAY"; + Schema schema = new Schema.SchemaBuilder() + .setSchemaName(TABLE_NAME) + .addSingleValueDimension(ID_COLUMN, DataType.STRING) + .addDateTime(timestampColumn, DataType.TIMESTAMP, "TIMESTAMP", "1:MILLISECONDS") + .addSingleValueDimension(collidingVariantColumn, DataType.VARIANT) + .build(); + FieldConfig timestampFieldConfig = new FieldConfig.Builder(timestampColumn) + .withTimestampConfig(new TimestampConfig(List.of(TimestampIndexGranularity.DAY))) + .build(); + FieldConfig variantFieldConfig = new FieldConfig.Builder(collidingVariantColumn) + .withEncodingType(EncodingType.RAW) + .build(); + StarTreeIndexConfig starTreeIndexConfig = + new StarTreeIndexConfig(List.of(ID_COLUMN), null, List.of("SUM__" + collidingVariantColumn), null, 1); + TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNullHandlingEnabled(true) + .setFieldConfigList(List.of(timestampFieldConfig, variantFieldConfig)) + .setStarTreeIndexConfigs(List.of(starTreeIndexConfig)) + .build(); + + assertInvalid(tableConfig, schema, "Star-tree index cannot be created on VARIANT column"); + } + private static FieldConfig rawVariantFieldConfig() { return new FieldConfig.Builder(VARIANT_COLUMN) .withEncodingType(EncodingType.RAW) From 496b466c599cf9c4f7114e9cd77a0370468fb453 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Aug 2026 02:43:39 -0700 Subject: [PATCH 8/8] Fix VARIANT object navigation interoperability Honor producer-defined physical value ordering when object field offsets are non-monotonic, and preserve exact subtree envelopes through self-delimiting size parsing. Use allocation-free lookup for wide objects with a cross-producer ordering fallback, clarify materializing cursor APIs, and add focused, cross-engine, Parquet, and JMH coverage. --- .../evaluator/InbuiltFunctionEvaluator.java | 2 +- .../pinot/common/utils/VariantUtils.java | 308 +++++++++++++- .../pinot/common/utils/VariantUtilsTest.java | 377 +++++++++++++++++- .../tests/custom/VariantTypeTest.java | 154 ++++++- .../pinot/perf/BenchmarkVariantGet.java | 144 +++++++ .../ParquetVariantRecordReaderTest.java | 34 ++ .../operator/operands/VariantOperand.java | 2 +- 7 files changed, 996 insertions(+), 25 deletions(-) create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkVariantGet.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java index a9374c1cf147..9262504c6716 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/evaluator/InbuiltFunctionEvaluator.java @@ -259,7 +259,7 @@ private Object extract(@Nullable byte[] variant, boolean tolerant) { if (!present) { return null; } - return _reusableResult.getExternalValue(targetType); + return _reusableResult.toExternalValue(targetType); } @Nullable diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java index eb393401e6be..c20aca8378c0 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java @@ -103,6 +103,8 @@ public final class VariantUtils { private static final int VARIANT_UUID = 20; private static final int VARIANT_METADATA_VERSION_MASK = 0x0F; private static final int VARIANT_METADATA_VERSION = 1; + private static final int OBJECT_BINARY_SEARCH_THRESHOLD = 32; + private static final int INVALID_UTF8_COMPARISON = Integer.MIN_VALUE; private static final VariantPath ROOT_PATH = new VariantPath(new PathElement[0]); private VariantUtils() { @@ -222,7 +224,7 @@ public UUID getUuidValue() { /// ///

For BYTES and VARIANT, the returned {@code byte[]} may be retained after this result is reused. It must be /// treated as immutable and copied before mutation. - public Object getExternalValue(ResultType resultType) { + public Object toExternalValue(ResultType resultType) { switch (resultType) { case BOOLEAN: return _intValue != 0; @@ -257,7 +259,7 @@ public Object getExternalValue(ResultType resultType) { /// external-object round trip in the multi-stage engine. For BYTES, UUID, and VARIANT, the returned /// {@link ByteArray} wraps a newly materialized array that may be retained after this result is reused. Neither the /// wrapper nor its array may be mutated; callers must copy the array before mutation. - public Object getInternalValue(ResultType resultType) { + public Object toInternalValue(ResultType resultType) { switch (resultType) { case BOOLEAN: return _intValue; @@ -359,7 +361,7 @@ public static Object variantGet(@Nullable byte[] envelope, String path, String t @Nullable public static Object variantGet(@Nullable byte[] envelope, VariantPath path, ResultType targetType) { ReusableResult result = new ReusableResult(); - return extractInto(envelope, path, targetType, result) ? result.getExternalValue(targetType) : null; + return extractInto(envelope, path, targetType, result) ? result.toExternalValue(targetType) : null; } /// Strictly extracts into a reusable, unboxed result. @@ -405,7 +407,7 @@ public static Object tryVariantGet(@Nullable byte[] envelope, String path, Strin public static Object tryVariantGet(@Nullable byte[] envelope, VariantPath path, ResultType targetType) { try { ReusableResult result = new ReusableResult(); - return tryExtractInto(envelope, path, targetType, result) ? result.getExternalValue(targetType) : null; + return tryExtractInto(envelope, path, targetType, result) ? result.toExternalValue(targetType) : null; } catch (RuntimeException e) { return null; } @@ -1188,7 +1190,7 @@ private boolean navigate(byte[] envelope, VariantPath path) { reset(envelope); for (PathElement element : path._elements) { if (element._field != null) { - if (getType() != Variant.Type.OBJECT || !selectObjectField(element._fieldUtf8)) { + if (getType() != Variant.Type.OBJECT || !selectObjectField(element)) { return false; } } else if (getType() != Variant.Type.ARRAY || !selectArrayElement(element._index)) { @@ -1392,7 +1394,7 @@ private Variant asVariant() { return new Variant(_envelope, _selectedOffset, _selectedLength, _envelope, _metadataOffset, _metadataLength); } - private boolean selectObjectField(byte[] fieldUtf8) { + private boolean selectObjectField(PathElement field) { int header = getHeader(); int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; @@ -1409,23 +1411,71 @@ private boolean selectObjectField(byte[] fieldUtf8) { int totalDataLength = readUnsignedLittleEndian(finalOffsetPosition, offsetSize, selectedLimit()); requireRange(dataStart, totalDataLength, selectedLimit(), "Variant object data"); - for (int i = 0; i < numElements; i++) { + if (numElements < OBJECT_BINARY_SEARCH_THRESHOLD || !field._binarySearchSafe) { + return selectObjectFieldLinear(field._fieldUtf8, 0, numElements, idStart, idSize, offsetStart, offsetSize, + dataStart, totalDataLength); + } + + // Preserve the linear lookup's constant-time first-field case before paying the binary-search cost for the rest + // of a wide object. + int firstId = readUnsignedLittleEndian(idStart, idSize, offsetStart); + if (metadataKeyEquals(firstId, field._fieldUtf8)) { + selectObjectFieldAtIndex(0, offsetStart, offsetSize, dataStart, totalDataLength); + return true; + } + + int low = 1; + int high = numElements - 1; + while (low <= high) { + int index = (low + high) >>> 1; + int id = readUnsignedLittleEndian(idStart + index * idSize, idSize, offsetStart); + int comparison = metadataKeyCompare(id, field._field); + if (comparison == INVALID_UTF8_COMPARISON) { + // Malformed UTF-8 is not a valid Variant key, but retain the old byte-equality behavior for tolerant + // callers instead of relying on an ordering that is no longer defined. + return selectObjectFieldLinear(field._fieldUtf8, 1, numElements, idStart, idSize, offsetStart, offsetSize, + dataStart, totalDataLength); + } + if (comparison < 0) { + low = index + 1; + } else if (comparison > 0) { + high = index - 1; + } else { + selectObjectFieldAtIndex(index, offsetStart, offsetSize, dataStart, totalDataLength); + return true; + } + } + // parquet-java orders object keys with String.compareTo (UTF-16 code units), while other conforming producers + // such as Arrow Rust use Unicode scalar/UTF-8 order. Those orders differ for supplementary characters relative + // to U+E000..U+FFFF, so a miss under the parquet-java ordering is not authoritative for an external envelope. + // Preserve the allocation-free binary fast path for hits, then use byte equality as the interoperability-safe + // fallback for misses and unknown producer orderings. + return selectObjectFieldLinear(field._fieldUtf8, 1, numElements, idStart, idSize, offsetStart, offsetSize, + dataStart, totalDataLength); + } + + private boolean selectObjectFieldLinear(byte[] fieldUtf8, int startIndex, int numElements, int idStart, + int idSize, int offsetStart, int offsetSize, int dataStart, int totalDataLength) { + for (int i = startIndex; i < numElements; i++) { int id = readUnsignedLittleEndian(idStart + i * idSize, idSize, offsetStart); if (metadataKeyEquals(id, fieldUtf8)) { - int offset = readUnsignedLittleEndian(offsetStart + i * offsetSize, offsetSize, dataStart); - int nextOffset = readUnsignedLittleEndian(offsetStart + (i + 1) * offsetSize, offsetSize, dataStart); - if (offset > nextOffset || nextOffset > totalDataLength) { - throw new IllegalStateException( - "Invalid Variant object offsets: " + offset + ", " + nextOffset + ", total=" + totalDataLength); - } - _selectedOffset = dataStart + offset; - _selectedLength = nextOffset - offset; + selectObjectFieldAtIndex(i, offsetStart, offsetSize, dataStart, totalDataLength); return true; } } return false; } + private void selectObjectFieldAtIndex(int index, int offsetStart, int offsetSize, int dataStart, + int totalDataLength) { + int offset = readUnsignedLittleEndian(offsetStart + index * offsetSize, offsetSize, dataStart); + int valueOffset = checkedPosition((long) dataStart + offset, dataStart + totalDataLength, + "Variant object field"); + int valueLength = encodedValueLength(valueOffset, dataStart + totalDataLength); + _selectedOffset = valueOffset; + _selectedLength = valueLength; + } + private boolean selectArrayElement(int index) { int header = getHeader(); int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; @@ -1454,6 +1504,208 @@ private boolean selectArrayElement(int index) { return true; } + /// Returns the exact byte length of the self-delimiting Variant value at {@code valueOffset}. + /// + ///

Object field offsets are sorted by key, while their encoded values may appear in any physical order. The + /// adjacent field offset is therefore not necessarily the end of the selected value. Reading the selected value's + /// own header preserves exact subtree envelopes without scanning every object offset. + private int encodedValueLength(int valueOffset, int limit) { + requireRange(valueOffset, 1, limit, "Variant value"); + int header = Byte.toUnsignedInt(_envelope[valueOffset]); + int basicType = header & VARIANT_BASIC_TYPE_MASK; + int typeInfo = (header >>> 2) & VARIANT_PRIMITIVE_TYPE_MASK; + int length; + switch (basicType) { + case VARIANT_SHORT_STRING: + length = 1 + typeInfo; + break; + case VARIANT_OBJECT: { + int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(valueOffset + 1, sizeBytes, limit); + int idSize = ((typeInfo >>> 2) & 3) + 1; + int offsetSize = (typeInfo & 3) + 1; + int offsetStart = checkedPosition((long) valueOffset + 1 + sizeBytes + (long) numElements * idSize, + limit, "Variant object offsets"); + int dataStart = checkedPosition((long) offsetStart + ((long) numElements + 1) * offsetSize, limit, + "Variant object data"); + int finalOffsetPosition = checkedPosition((long) offsetStart + (long) numElements * offsetSize, limit, + "Variant object final offset"); + int dataLength = readUnsignedLittleEndian(finalOffsetPosition, offsetSize, limit); + int valueEnd = checkedPosition((long) dataStart + dataLength, limit, "Variant object value"); + return valueEnd - valueOffset; + } + case VARIANT_ARRAY: { + int sizeBytes = ((typeInfo >>> 2) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(valueOffset + 1, sizeBytes, limit); + int offsetSize = (typeInfo & 3) + 1; + int offsetStart = checkedPosition((long) valueOffset + 1 + sizeBytes, limit, "Variant array offsets"); + int dataStart = checkedPosition((long) offsetStart + ((long) numElements + 1) * offsetSize, limit, + "Variant array data"); + int finalOffsetPosition = checkedPosition((long) offsetStart + (long) numElements * offsetSize, limit, + "Variant array final offset"); + int dataLength = readUnsignedLittleEndian(finalOffsetPosition, offsetSize, limit); + int valueEnd = checkedPosition((long) dataStart + dataLength, limit, "Variant array value"); + return valueEnd - valueOffset; + } + case VARIANT_PRIMITIVE: + switch (typeInfo) { + case VARIANT_NULL: + case VARIANT_TRUE: + case VARIANT_FALSE: + length = 1; + break; + case VARIANT_INT8: + length = 2; + break; + case VARIANT_INT16: + length = 3; + break; + case VARIANT_INT32: + case VARIANT_DATE: + case VARIANT_FLOAT: + length = 5; + break; + case VARIANT_INT64: + case VARIANT_DOUBLE: + case VARIANT_TIMESTAMP_TZ: + case VARIANT_TIMESTAMP_NTZ: + case VARIANT_TIME: + case VARIANT_TIMESTAMP_NANOS_TZ: + case VARIANT_TIMESTAMP_NANOS_NTZ: + length = 9; + break; + case VARIANT_DECIMAL4: + length = 6; + break; + case VARIANT_DECIMAL8: + length = 10; + break; + case VARIANT_DECIMAL16: + length = 18; + break; + case VARIANT_BINARY: + case VARIANT_LONG_STRING: + length = checkedPosition((long) valueOffset + 1 + Integer.BYTES + + readUnsignedLittleEndian(valueOffset + 1, Integer.BYTES, limit), limit, "Variant binary value") + - valueOffset; + break; + case VARIANT_UUID: + length = 1 + UuidUtils.UUID_NUM_BYTES; + break; + default: + throw new UnsupportedOperationException("Unknown type in Variant. primitive type: " + typeInfo); + } + break; + default: + throw new IllegalStateException("Unhandled Variant basic type: " + basicType); + } + requireRange(valueOffset, length, limit, "Variant value"); + return length; + } + + /// Compares one encoded metadata key with {@code expected} using Java String's UTF-16 code-unit order without + /// materializing the encoded key. Returns {@link #INVALID_UTF8_COMPARISON} when the metadata key is not valid + /// UTF-8, in which case callers must not rely on object-key ordering. + private int metadataKeyCompare(int id, String expected) { + ensureMetadataParsed(); + if (id < 0 || id >= _metadataDictSize) { + throw new IllegalArgumentException( + "Invalid dictionary id: " + id + ". dictionary size: " + _metadataDictSize); + } + int offset = readUnsignedLittleEndian(_metadataOffsetListOffset + id * _metadataOffsetSize, + _metadataOffsetSize, _metadataDataOffset); + int nextOffset = readUnsignedLittleEndian(_metadataOffsetListOffset + (id + 1) * _metadataOffsetSize, + _metadataOffsetSize, _metadataDataOffset); + if (offset > nextOffset || nextOffset > _metadataDataLength) { + throw new IllegalStateException( + "Invalid Variant metadata offsets: " + offset + ", " + nextOffset + ", total=" + _metadataDataLength); + } + + int byteIndex = _metadataDataOffset + offset; + int byteLimit = _metadataDataOffset + nextOffset; + int charIndex = 0; + while (byteIndex < byteLimit) { + int decoded = decodeUtf8CodePoint(byteIndex, byteLimit); + if (decoded < 0) { + return INVALID_UTF8_COMPARISON; + } + byteIndex += decoded >>> 24; + int codePoint = decoded & 0x1F_FFFF; + if (codePoint < Character.MIN_SUPPLEMENTARY_CODE_POINT) { + if (charIndex == expected.length()) { + return 1; + } + int comparison = Character.compare((char) codePoint, expected.charAt(charIndex++)); + if (comparison != 0) { + return comparison; + } + } else { + if (charIndex == expected.length()) { + return 1; + } + int comparison = Character.compare(Character.highSurrogate(codePoint), expected.charAt(charIndex++)); + if (comparison != 0) { + return comparison; + } + if (charIndex == expected.length()) { + return 1; + } + comparison = Character.compare(Character.lowSurrogate(codePoint), expected.charAt(charIndex++)); + if (comparison != 0) { + return comparison; + } + } + } + return charIndex == expected.length() ? 0 : -1; + } + + /// Returns the decoded code point with its encoded width in the top byte, or {@code -1} for malformed UTF-8. + private int decodeUtf8CodePoint(int index, int limit) { + int first = Byte.toUnsignedInt(_envelope[index]); + if (first <= 0x7F) { + return 1 << 24 | first; + } + if (first >= 0xC2 && first <= 0xDF) { + if (index + 1 >= limit) { + return -1; + } + int second = Byte.toUnsignedInt(_envelope[index + 1]); + if ((second & 0xC0) != 0x80) { + return -1; + } + return 2 << 24 | (first & 0x1F) << 6 | second & 0x3F; + } + if (first >= 0xE0 && first <= 0xEF) { + if (index + 2 >= limit) { + return -1; + } + int second = Byte.toUnsignedInt(_envelope[index + 1]); + int third = Byte.toUnsignedInt(_envelope[index + 2]); + if ((third & 0xC0) != 0x80 + || (first == 0xE0 ? second < 0xA0 || second > 0xBF + : first == 0xED ? second < 0x80 || second > 0x9F : (second & 0xC0) != 0x80)) { + return -1; + } + return 3 << 24 | (first & 0x0F) << 12 | (second & 0x3F) << 6 | third & 0x3F; + } + if (first >= 0xF0 && first <= 0xF4) { + if (index + 3 >= limit) { + return -1; + } + int second = Byte.toUnsignedInt(_envelope[index + 1]); + int third = Byte.toUnsignedInt(_envelope[index + 2]); + int fourth = Byte.toUnsignedInt(_envelope[index + 3]); + if ((third & 0xC0) != 0x80 || (fourth & 0xC0) != 0x80 + || (first == 0xF0 ? second < 0x90 || second > 0xBF + : first == 0xF4 ? second < 0x80 || second > 0x8F : (second & 0xC0) != 0x80)) { + return -1; + } + return 4 << 24 | (first & 0x07) << 18 | (second & 0x3F) << 12 | (third & 0x3F) << 6 + | fourth & 0x3F; + } + return -1; + } + private boolean metadataKeyEquals(int id, byte[] expected) { ensureMetadataParsed(); if (id < 0 || id >= _metadataDictSize) { @@ -1577,19 +1829,39 @@ private static final class PathElement { private final String _field; private final byte[] _fieldUtf8; private final int _index; + private final boolean _binarySearchSafe; - private PathElement(String field, byte[] fieldUtf8, int index) { + private PathElement(String field, byte[] fieldUtf8, int index, boolean binarySearchSafe) { _field = field; _fieldUtf8 = fieldUtf8; _index = index; + _binarySearchSafe = binarySearchSafe; } private static PathElement forField(String field) { - return new PathElement(field, field.getBytes(StandardCharsets.UTF_8), -1); + return new PathElement(field, field.getBytes(StandardCharsets.UTF_8), -1, hasWellFormedUtf16(field)); } private static PathElement forIndex(int index) { - return new PathElement(null, null, index); + return new PathElement(null, null, index, false); + } + + private static boolean hasWellFormedUtf16(String value) { + int i = 0; + while (i < value.length()) { + char current = value.charAt(i); + if (Character.isHighSurrogate(current)) { + if (i + 1 == value.length() || !Character.isLowSurrogate(value.charAt(i + 1))) { + return false; + } + i += 2; + } else if (Character.isLowSurrogate(current)) { + return false; + } else { + i++; + } + } + return true; } } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java index 02774a9616f3..2fbfaae755ec 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java @@ -20,6 +20,7 @@ import java.math.BigDecimal; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.sql.Timestamp; import java.util.Arrays; import java.util.EnumSet; @@ -31,6 +32,7 @@ import org.apache.calcite.sql.type.SqlTypeName; import org.apache.parquet.variant.Variant; import org.apache.parquet.variant.VariantBuilder; +import org.apache.parquet.variant.VariantObjectBuilder; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.common.utils.VariantUtils.ResultType; import org.apache.pinot.common.utils.VariantUtils.ReusableResult; @@ -326,6 +328,202 @@ public void testReusableNestedNavigationAndCompiledPredicates() { assertNull(VariantUtils.variantTypeOf(new byte[0], valuePath, result)); } + @DataProvider(name = "wideObjectSizes") + public Object[][] wideObjectSizes() { + return new Object[][]{{8}, {31}, {32}, {33}, {100}}; + } + + @Test(dataProvider = "wideObjectSizes") + public void testWideObjectLookupMatchesParquetJava(int numFields) { + VariantBuilder builder = new VariantBuilder(); + VariantObjectBuilder objectBuilder = builder.startObject(); + // Append in reverse order so the test relies on the encoded object's lexicographic field ordering, not insertion + // order. Parquet switches from linear to binary lookup at 32 fields, which Pinot deliberately mirrors. + for (int i = numFields - 1; i >= 0; i--) { + objectBuilder.appendKey(wideField(i)); + objectBuilder.appendInt(i); + } + builder.endObject(); + Variant variant = builder.build(); + byte[] envelope = encode(variant); + + for (String key : List.of(wideField(0), wideField(numFields / 2), wideField(numFields - 1), "missing")) { + Variant parquetValue = variant.getFieldByKey(key); + Object expected = parquetValue != null ? parquetValue.getInt() : null; + assertEquals(VariantUtils.variantGet(envelope, "$." + key, "INT"), expected); + assertEquals(VariantUtils.tryVariantGet(envelope, "$." + key, "INT"), expected); + } + } + + @Test + public void testWideObjectLookupUsesJavaUtf16Ordering() { + String supplementary = "\uD800\uDC00"; + String supplementarySuccessor = "\uD800\uDC01"; + String privateUse = "\uE000"; + VariantBuilder builder = new VariantBuilder(); + VariantObjectBuilder objectBuilder = builder.startObject(); + for (int i = 29; i >= 0; i--) { + objectBuilder.appendKey(wideField(i)); + objectBuilder.appendInt(i); + } + objectBuilder.appendKey(privateUse); + objectBuilder.appendInt(2_000); + objectBuilder.appendKey(supplementary); + objectBuilder.appendInt(1_000); + builder.endObject(); + Variant variant = builder.build(); + byte[] envelope = encode(variant); + + // Java String order places the supplementary key before U+E000, while unsigned UTF-8 byte order does the + // opposite. Cross-checking parquet-java prevents an allocation-free byte comparator from changing semantics. + for (String key : List.of(supplementary, supplementarySuccessor, privateUse)) { + Variant parquetValue = variant.getFieldByKey(key); + Object expected = parquetValue != null ? parquetValue.getInt() : null; + assertEquals(VariantUtils.variantGet(envelope, "$." + key, "INT"), expected); + } + } + + @Test + public void testWideObjectLookupSupportsArrowRustOrdering() { + byte[] envelope = arrowRustOrderedWideObjectEnvelope(); + String supplementary = "\uD800\uDC00"; + String privateUse = "\uE000"; + + assertEquals(VariantUtils.variantGet(envelope, "$." + supplementary, "INT"), 1_000); + assertEquals(VariantUtils.tryVariantGet(envelope, "$." + supplementary, "INT"), 1_000); + assertTrue(VariantUtils.variantExists(envelope, "$." + supplementary)); + assertEquals(VariantUtils.variantGet(envelope, "$." + privateUse, "INT"), 2_000); + assertEquals(VariantUtils.tryVariantGet(envelope, "$." + privateUse, "INT"), 2_000); + assertTrue(VariantUtils.variantExists(envelope, "$." + privateUse)); + assertNull(VariantUtils.variantGet(envelope, "$.missing", "INT")); + assertNull(VariantUtils.tryVariantGet(envelope, "$.missing", "INT")); + assertFalse(VariantUtils.variantExists(envelope, "$.missing")); + } + + @Test + public void testNonMonotonicObjectOffsets() { + byte[] envelope = nonMonotonicObjectEnvelope(); + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + Variant parquetVariant = new Variant(decoded.getValue(), decoded.getMetadata()); + + for (int i = 0; i < 3; i++) { + String key = String.valueOf((char) ('a' + i)); + int expected = i + 1; + assertEquals(parquetVariant.getFieldByKey(key).getInt(), expected); + assertEquals(VariantUtils.variantGet(envelope, "$." + key, "INT"), expected); + assertEquals(VariantUtils.tryVariantGet(envelope, "$." + key, "INT"), expected); + + byte[] subtree = VariantUtils.variantGet(envelope, "$." + key); + VariantEnvelope.Decoded decodedSubtree = VariantEnvelope.decode(subtree); + assertEquals(decodedSubtree.getValue().remaining(), 2); + assertEquals(new Variant(decodedSubtree.getValue(), decodedSubtree.getMetadata()).getInt(), expected); + assertEquals(VariantUtils.variantToJson(subtree), Integer.toString(expected)); + } + assertEquals(VariantUtils.variantToJson(envelope), "{\"a\":1,\"b\":2,\"c\":3}"); + assertFalse(VariantUtils.variantExists(envelope, "$.missing")); + } + + @DataProvider(name = "nonMonotonicSubtreeEncodingCases") + public Object[][] nonMonotonicSubtreeEncodingCases() { + return new Object[][]{ + nonMonotonicSubtreeEncodingCase("object", Variant.Type.OBJECT, builder -> { + VariantObjectBuilder nested = builder.startObject(); + nested.appendKey("inner"); + nested.appendInt(7); + builder.endObject(); + }), + nonMonotonicSubtreeEncodingCase("array", Variant.Type.ARRAY, builder -> { + VariantBuilder nested = builder.startArray(); + nested.appendInt(7); + nested.appendBoolean(true); + builder.endArray(); + }), + nonMonotonicSubtreeEncodingCase("null", Variant.Type.NULL, VariantBuilder::appendNull), + nonMonotonicSubtreeEncodingCase("true", Variant.Type.BOOLEAN, builder -> builder.appendBoolean(true)), + nonMonotonicSubtreeEncodingCase("false", Variant.Type.BOOLEAN, builder -> builder.appendBoolean(false)), + nonMonotonicSubtreeEncodingCase("int8", Variant.Type.BYTE, builder -> builder.appendByte((byte) -8)), + nonMonotonicSubtreeEncodingCase("int16", Variant.Type.SHORT, + builder -> builder.appendShort((short) 32_000)), + nonMonotonicSubtreeEncodingCase("int32", Variant.Type.INT, builder -> builder.appendInt(-123_456)), + nonMonotonicSubtreeEncodingCase("int64", Variant.Type.LONG, + builder -> builder.appendLong(9_876_543_210L)), + nonMonotonicSubtreeEncodingCase("short-string", Variant.Type.STRING, + builder -> builder.appendString("short")), + nonMonotonicSubtreeEncodingCase("long-string", Variant.Type.STRING, + builder -> builder.appendString("x".repeat(128))), + nonMonotonicSubtreeEncodingCase("double", Variant.Type.DOUBLE, + builder -> builder.appendDouble(-123.5D)), + nonMonotonicSubtreeEncodingCase("decimal4", Variant.Type.DECIMAL4, + builder -> builder.appendDecimal(new BigDecimal("12.34"))), + nonMonotonicSubtreeEncodingCase("decimal8", Variant.Type.DECIMAL8, + builder -> builder.appendDecimal(new BigDecimal("1234567890.12"))), + nonMonotonicSubtreeEncodingCase("decimal16", Variant.Type.DECIMAL16, + builder -> builder.appendDecimal(new BigDecimal("12345678901234567890.1234"))), + nonMonotonicSubtreeEncodingCase("date", Variant.Type.DATE, builder -> builder.appendDate(1)), + nonMonotonicSubtreeEncodingCase("timestamp-tz", Variant.Type.TIMESTAMP_TZ, + builder -> builder.appendTimestampTz(1_234_567L)), + nonMonotonicSubtreeEncodingCase("timestamp-ntz", Variant.Type.TIMESTAMP_NTZ, + builder -> builder.appendTimestampNtz(1_234_567L)), + nonMonotonicSubtreeEncodingCase("float", Variant.Type.FLOAT, builder -> builder.appendFloat(1.25F)), + nonMonotonicSubtreeEncodingCase("binary", Variant.Type.BINARY, + builder -> builder.appendBinary(ByteBuffer.wrap(new byte[]{0, 1, -1, 42}))), + nonMonotonicSubtreeEncodingCase("time", Variant.Type.TIME, + builder -> builder.appendTime(3_723_004_005L)), + nonMonotonicSubtreeEncodingCase("timestamp-nanos-tz", Variant.Type.TIMESTAMP_NANOS_TZ, + builder -> builder.appendTimestampNanosTz(1_234_567_891L)), + nonMonotonicSubtreeEncodingCase("timestamp-nanos-ntz", Variant.Type.TIMESTAMP_NANOS_NTZ, + builder -> builder.appendTimestampNanosNtz(1_234_567_891L)), + nonMonotonicSubtreeEncodingCase("uuid", Variant.Type.UUID, + builder -> builder.appendUUID(UUID.fromString("00112233-4455-6677-8899-aabbccddeeff"))) + }; + } + + @Test(dataProvider = "nonMonotonicSubtreeEncodingCases") + public void testNonMonotonicSubtreeForEveryEncoding(String description, Variant.Type expectedType, + int expectedLength, byte[] envelope) { + assertTrue(hasNonMonotonicObjectOffsets(envelope), description); + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + Variant expected = new Variant(decoded.getValue(), decoded.getMetadata()).getFieldByKey("a"); + assertEquals(expected.getType(), expectedType, description); + + byte[] subtree = VariantUtils.variantGet(envelope, "$.a"); + byte[] tolerantSubtree = VariantUtils.tryVariantGet(envelope, "$.a"); + VariantEnvelope.Decoded decodedSubtree = VariantEnvelope.decode(subtree); + byte[] actualValue = toByteArray(decodedSubtree.getValue()); + byte[] parquetValueAndTrailingData = toByteArray(expected.getValueBuffer()); + assertEquals(actualValue.length, expectedLength, description); + assertTrue(Arrays.equals(actualValue, Arrays.copyOf(parquetValueAndTrailingData, expectedLength)), description); + assertTrue(Arrays.equals(tolerantSubtree, subtree), description); + + byte[] expectedEnvelope = VariantEnvelope.encode(expected.getMetadataBuffer(), expected.getValueBuffer()); + assertEquals(VariantUtils.variantToJson(subtree), VariantUtils.variantToJson(expectedEnvelope), description); + } + + @Test + public void testNonMonotonicSubtreeRejectsTruncatedVariableLengthValue() { + byte[] envelope = nonMonotonicSubtreeEnvelope( + builder -> builder.appendBinary(ByteBuffer.wrap(new byte[]{1, 2, 3}))); + int metadataLength = readBigEndianInt(envelope, 8); + int valueStart = VariantEnvelope.HEADER_SIZE + metadataLength; + int valueHeader = Byte.toUnsignedInt(envelope[valueStart]); + int typeInfo = valueHeader >>> 2 & 0x3F; + int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(envelope, valueStart + 1, sizeBytes); + int idSize = ((typeInfo >>> 2) & 3) + 1; + int offsetSize = (typeInfo & 3) + 1; + int offsetStart = valueStart + 1 + sizeBytes + numElements * idSize; + int dataStart = offsetStart + (numElements + 1) * offsetSize; + int targetOffset = readUnsignedLittleEndian(envelope, offsetStart, offsetSize); + int targetStart = dataStart + targetOffset; + envelope[targetStart + 1] = 0x7F; + envelope[targetStart + 2] = 0; + envelope[targetStart + 3] = 0; + envelope[targetStart + 4] = 0; + + assertThrows(IllegalArgumentException.class, () -> VariantUtils.variantGet(envelope, "$.a")); + assertNull(VariantUtils.tryVariantGet(envelope, "$.a")); + } + @Test public void testReusableResultMalformedInputAndReuse() { VariantPath path = VariantUtils.compilePath("$.items[0]"); @@ -541,8 +739,8 @@ private static void assertReusableParity(byte[] envelope, VariantPath path, Resu ReusableResult result = new ReusableResult(); assertTrue(VariantUtils.extractInto(envelope, path, resultType, result)); assertTrue(VariantUtils.tryExtractInto(envelope, path, resultType, result)); - Object externalValue = result.getExternalValue(resultType); - Object internalValue = result.getInternalValue(resultType); + Object externalValue = result.toExternalValue(resultType); + Object internalValue = result.toInternalValue(resultType); switch (resultType) { case BOOLEAN: assertEquals(result.getIntValue() != 0, expected); @@ -610,10 +808,183 @@ private static void assertReusableParity(byte[] envelope, VariantPath path, Resu } private static byte[] encode(VariantBuilder builder) { - Variant variant = builder.build(); + return encode(builder.build()); + } + + private static byte[] encode(Variant variant) { return VariantEnvelope.encode(variant.getMetadataBuffer(), variant.getValueBuffer()); } + private static String wideField(int index) { + return String.format("field%03d", index); + } + + /// Reorders a parquet-java object into Arrow Rust's Unicode scalar/UTF-8 key order while preserving each field's + /// dictionary id, physical value offset, and shared metadata. The resulting 32-field envelope exercises a valid + /// external ordering for which Java UTF-16 binary search alone is not authoritative. + private static byte[] arrowRustOrderedWideObjectEnvelope() { + String supplementary = "\uD800\uDC00"; + String privateUse = "\uE000"; + VariantBuilder builder = new VariantBuilder(); + VariantObjectBuilder objectBuilder = builder.startObject(); + for (int i = 0; i < 30; i++) { + objectBuilder.appendKey(wideField(i)); + objectBuilder.appendInt(i); + } + objectBuilder.appendKey(supplementary); + objectBuilder.appendInt(1_000); + objectBuilder.appendKey(privateUse); + objectBuilder.appendInt(2_000); + builder.endObject(); + + byte[] envelope = encode(builder); + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + Variant variant = new Variant(decoded.getValue(), decoded.getMetadata()); + byte[] value = toByteArray(decoded.getValue()); + int typeInfo = Byte.toUnsignedInt(value[0]) >>> 2 & 0x3F; + int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(value, 1, sizeBytes); + int idSize = ((typeInfo >>> 2) & 3) + 1; + int offsetSize = (typeInfo & 3) + 1; + int idStart = 1 + sizeBytes; + int offsetStart = idStart + numElements * idSize; + + String[] keys = new String[numElements]; + Integer[] arrowOrder = new Integer[numElements]; + for (int i = 0; i < numElements; i++) { + keys[i] = variant.getFieldAtIndex(i).key; + arrowOrder[i] = i; + } + Arrays.sort(arrowOrder, (left, right) -> compareUnsignedUtf8(keys[left], keys[right])); + + byte[] reorderedValue = Arrays.copyOf(value, value.length); + for (int i = 0; i < numElements; i++) { + int sourceIndex = arrowOrder[i]; + System.arraycopy(value, idStart + sourceIndex * idSize, reorderedValue, idStart + i * idSize, idSize); + System.arraycopy(value, offsetStart + sourceIndex * offsetSize, reorderedValue, + offsetStart + i * offsetSize, offsetSize); + } + return VariantEnvelope.encode(decoded.getMetadata(), ByteBuffer.wrap(reorderedValue)); + } + + private static int compareUnsignedUtf8(String left, String right) { + byte[] leftBytes = left.getBytes(StandardCharsets.UTF_8); + byte[] rightBytes = right.getBytes(StandardCharsets.UTF_8); + int length = Math.min(leftBytes.length, rightBytes.length); + for (int i = 0; i < length; i++) { + int comparison = Integer.compare(Byte.toUnsignedInt(leftBytes[i]), Byte.toUnsignedInt(rightBytes[i])); + if (comparison != 0) { + return comparison; + } + } + return Integer.compare(leftBytes.length, rightBytes.length); + } + + private static Object[] nonMonotonicSubtreeEncodingCase(String description, Variant.Type expectedType, + Consumer appender) { + VariantBuilder standaloneBuilder = new VariantBuilder(); + appender.accept(standaloneBuilder); + int expectedLength = standaloneBuilder.build().getValueBuffer().remaining(); + return new Object[]{description, expectedType, expectedLength, nonMonotonicSubtreeEnvelope(appender)}; + } + + private static byte[] nonMonotonicSubtreeEnvelope(Consumer appender) { + VariantBuilder builder = new VariantBuilder(); + VariantObjectBuilder objectBuilder = builder.startObject(); + // Write z first and a second. The object index is key-sorted as a/z, while their physical values remain z/a. + objectBuilder.appendKey("z"); + objectBuilder.appendByte((byte) 42); + objectBuilder.appendKey("a"); + appender.accept(objectBuilder); + builder.endObject(); + return reverseObjectPhysicalValues(encode(builder)); + } + + /// Reverses the physical value region of a canonical object while retaining its key-sorted ids and offset entries. + /// This models conforming producers that choose a physical value order independent of the logical key order. + private static byte[] reverseObjectPhysicalValues(byte[] envelope) { + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + byte[] value = toByteArray(decoded.getValue()); + int typeInfo = Byte.toUnsignedInt(value[0]) >>> 2 & 0x3F; + int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(value, 1, sizeBytes); + int idSize = ((typeInfo >>> 2) & 3) + 1; + int offsetSize = (typeInfo & 3) + 1; + int offsetStart = 1 + sizeBytes + numElements * idSize; + int dataStart = offsetStart + (numElements + 1) * offsetSize; + int totalDataLength = readUnsignedLittleEndian(value, offsetStart + numElements * offsetSize, offsetSize); + + int[] offsets = new int[numElements + 1]; + for (int i = 0; i <= numElements; i++) { + offsets[i] = readUnsignedLittleEndian(value, offsetStart + i * offsetSize, offsetSize); + } + byte[] reversedData = new byte[totalDataLength]; + int writeOffset = 0; + for (int i = numElements - 1; i >= 0; i--) { + int length = offsets[i + 1] - offsets[i]; + System.arraycopy(value, dataStart + offsets[i], reversedData, writeOffset, length); + writeUnsignedLittleEndian(value, offsetStart + i * offsetSize, offsetSize, writeOffset); + writeOffset += length; + } + System.arraycopy(reversedData, 0, value, dataStart, totalDataLength); + return VariantEnvelope.encode(decoded.getMetadata(), ByteBuffer.wrap(value)); + } + + private static boolean hasNonMonotonicObjectOffsets(byte[] envelope) { + byte[] value = toByteArray(VariantEnvelope.decode(envelope).getValue()); + int typeInfo = Byte.toUnsignedInt(value[0]) >>> 2 & 0x3F; + int sizeBytes = ((typeInfo >>> 4) & 1) == 0 ? 1 : Integer.BYTES; + int numElements = readUnsignedLittleEndian(value, 1, sizeBytes); + int idSize = ((typeInfo >>> 2) & 3) + 1; + int offsetSize = (typeInfo & 3) + 1; + int offsetStart = 1 + sizeBytes + numElements * idSize; + int previous = readUnsignedLittleEndian(value, offsetStart, offsetSize); + for (int i = 1; i < numElements; i++) { + int current = readUnsignedLittleEndian(value, offsetStart + i * offsetSize, offsetSize); + if (current < previous) { + return true; + } + previous = current; + } + return false; + } + + private static byte[] toByteArray(ByteBuffer buffer) { + ByteBuffer copy = buffer.duplicate(); + byte[] bytes = new byte[copy.remaining()]; + copy.get(bytes); + return bytes; + } + + private static int readUnsignedLittleEndian(byte[] bytes, int offset, int numBytes) { + int value = 0; + for (int i = 0; i < numBytes; i++) { + value |= Byte.toUnsignedInt(bytes[offset + i]) << Byte.SIZE * i; + } + return value; + } + + private static void writeUnsignedLittleEndian(byte[] bytes, int offset, int numBytes, int value) { + for (int i = 0; i < numBytes; i++) { + bytes[offset + i] = (byte) (value >>> Byte.SIZE * i); + } + } + + /// Returns a valid Variant object whose lexicographic fields a/b/c point to physically reversed int8 values. + /// Object offsets are `[4, 2, 0, 6]`, which the Variant specification explicitly permits. + private static byte[] nonMonotonicObjectEnvelope() { + byte[] metadata = {0x11, 0x03, 0x00, 0x01, 0x02, 0x03, 'a', 'b', 'c'}; + byte[] value = { + 0x02, 0x03, // object header, three fields + 0x00, 0x01, 0x02, // metadata dictionary ids for a, b, c + 0x04, 0x02, 0x00, 0x06, // physical offsets for a, b, c, then total data length + 0x0C, 0x03, // c = int8(3), at physical offset 0 + 0x0C, 0x02, // b = int8(2), at physical offset 2 + 0x0C, 0x01 // a = int8(1), at physical offset 4 + }; + return VariantEnvelope.encode(metadata, 0, metadata.length, value, 0, value.length); + } + private static Object[] jsonRenderingCase(Variant.Type type, String expectedJson, Consumer appender) { VariantBuilder builder = new VariantBuilder(); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java index 0edf5a46abdb..1bee8ffdff04 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java @@ -24,13 +24,30 @@ import java.io.InputStream; import java.nio.file.Files; import java.nio.file.StandardCopyOption; +import java.util.Arrays; +import java.util.HashSet; import java.util.List; import java.util.Locale; +import java.util.Set; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; import org.apache.pinot.integration.tests.ClusterIntegrationTestUtils; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.spi.ImmutableSegment; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.FileFormat; import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.VariantEnvelope; import org.testng.Assert; import org.testng.annotations.Test; @@ -44,7 +61,7 @@ public class VariantTypeTest extends CustomDataQueryClusterIntegrationTest { private static final String EVENT_ID = "eventId"; private static final String EVENT_TYPE = "eventType"; private static final String PAYLOAD = "payload"; - private static final int NUM_DOCS = 5; + private static final int NUM_DOCS = 6; @Override public String getTableName() { @@ -88,6 +105,14 @@ protected void setUpTable() } ClusterIntegrationTestUtils.buildSegmentFromFile(parquetFile, tableConfig, schema, "0", _segmentDir, _tarDir, FileFormat.PARQUET); + + Set existingSegments = segmentNames(_segmentDir); + byte[] nonMonotonicEnvelope = nonMonotonicObjectEnvelope(); + File interoperabilityFile = writeNonMonotonicVariantFile(nonMonotonicEnvelope); + ClusterIntegrationTestUtils.buildSegmentFromFile(interoperabilityFile, tableConfig, schema, "1", _segmentDir, + _tarDir, FileFormat.PARQUET); + File interoperabilitySegment = findNewSegment(_segmentDir, existingSegments); + assertSegmentPreservesVariantEnvelope(interoperabilitySegment, nonMonotonicEnvelope); uploadSegments(getTableName(), _tarDir); } @@ -138,6 +163,33 @@ public void testMaterializedPathAndDirectExtraction(boolean useMultiStageQueryEn } } + @Test(dataProvider = "useBothQueryEngines") + public void testNonMonotonicObjectOffsetsFromExternalParquet(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + + JsonNode response = postVariantQuery( + "SELECT variant_get(" + PAYLOAD + ", '$.a', 'INT'), " + + "try_variant_get(" + PAYLOAD + ", '$.b', 'INT'), " + + "variant_get(" + PAYLOAD + ", '$.c'), variant_to_json(" + PAYLOAD + ") " + + "FROM " + TABLE_NAME + " WHERE " + EVENT_ID + " = 'evt-006'"); + assertNoExceptions(response); + Assert.assertEquals( + response.get("resultTable").get("dataSchema").get("columnDataTypes").toString(), + "[\"INT\",\"INT\",\"VARIANT\",\"STRING\"]"); + JsonNode row = response.get("resultTable").get("rows").get(0); + Assert.assertEquals(row.get(0).asInt(), 1); + Assert.assertEquals(row.get(1).asInt(), 2); + Assert.assertEquals(row.get(2).asText(), "3"); + Assert.assertEquals(row.get(3).asText(), "{\"a\":1,\"b\":2,\"c\":3,\"eventType\":\"interop\"}"); + + response = postVariantQuery( + "SELECT try_variant_get(" + PAYLOAD + ", '$.missing', 'INT') FROM " + TABLE_NAME + + " WHERE " + EVENT_ID + " = 'evt-006'"); + assertNoExceptions(response); + Assert.assertTrue(response.get("resultTable").get("rows").get(0).get(0).isNull()); + } + @Test(dataProvider = "useBothQueryEngines") public void testJsonProjectionAndNullStates(boolean useMultiStageQueryEngine) throws Exception { @@ -305,7 +357,7 @@ public void testRawVariantAggregatesAreRejected(boolean useMultiStageQueryEngine JsonNode response = postVariantQuery("SELECT COUNT(" + PAYLOAD + ") FROM " + TABLE_NAME); assertNoExceptions(response); - Assert.assertEquals(response.get("resultTable").get("rows").get(0).get(0).asLong(), 4L, + Assert.assertEquals(response.get("resultTable").get("rows").get(0).get(0).asLong(), 5L, "COUNT is raw-value-independent and must retain SQL-null semantics"); } @@ -394,6 +446,9 @@ public void testRawVariantWindowKeysAreRejectedButTypedPathWorks() Assert.assertEquals(rows.get(4).get(0).asText(), "evt-005"); Assert.assertTrue(rows.get(4).get(1).isNull(), "SQL null must remain in the SQL-null partition"); Assert.assertEquals(rows.get(4).get(2).asLong(), 2L); + Assert.assertEquals(rows.get(5).get(0).asText(), "evt-006"); + Assert.assertEquals(rows.get(5).get(1).asText(), "interop"); + Assert.assertEquals(rows.get(5).get(2).asLong(), 1L); } @Test @@ -415,6 +470,101 @@ public void testRawVariantSetOperationsAreRejected() Assert.assertEquals(response.get("resultTable").get("rows").size(), 2); } + private File writeNonMonotonicVariantFile(byte[] envelope) + throws IOException { + MessageType parquetSchema = MessageTypeParser.parseMessageType( + "message variant_interoperability {" + + " required binary eventId (STRING);" + + " required int64 eventTime;" + + " optional group payload (VARIANT(1)) {" + + " required binary metadata;" + + " required binary value;" + + " }" + + "}"); + File parquetFile = new File(_tempDir, "variant_non_monotonic_offsets.parquet"); + Files.deleteIfExists(parquetFile.toPath()); + VariantEnvelope.Decoded decoded = VariantEnvelope.decode(envelope); + try (ParquetWriter writer = ExampleParquetWriter.builder(new Path(parquetFile.getAbsolutePath())) + .withType(parquetSchema).build()) { + Group row = new SimpleGroupFactory(parquetSchema).newGroup() + .append(EVENT_ID, "evt-006") + .append("eventTime", 1_700_000_005_000L); + row.addGroup(PAYLOAD) + .append("metadata", Binary.fromConstantByteBuffer(decoded.getMetadata())) + .append("value", Binary.fromConstantByteBuffer(decoded.getValue())); + writer.write(row); + } + return parquetFile; + } + + private static Set segmentNames(File segmentDirectory) { + Set names = new HashSet<>(); + File[] segments = segmentDirectory.listFiles(File::isDirectory); + if (segments != null) { + for (File segment : segments) { + names.add(segment.getName()); + } + } + return names; + } + + private static File findNewSegment(File segmentDirectory, Set existingNames) { + File[] segments = segmentDirectory.listFiles(File::isDirectory); + Assert.assertNotNull(segments, "Failed to list generated segments in " + segmentDirectory); + File newSegment = null; + for (File segment : segments) { + if (!existingNames.contains(segment.getName())) { + Assert.assertNull(newSegment, "Expected exactly one new interoperability segment"); + newSegment = segment; + } + } + Assert.assertNotNull(newSegment, "Missing generated interoperability segment"); + return newSegment; + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + private static void assertSegmentPreservesVariantEnvelope(File segmentDirectory, byte[] expected) { + ImmutableSegment segment = null; + ForwardIndexReaderContext context = null; + try { + segment = ImmutableSegmentLoader.load(segmentDirectory, ReadMode.heap); + ForwardIndexReader reader = segment.getForwardIndex(PAYLOAD); + context = reader.createContext(); + byte[] actual = reader.getBytes(0, context); + Assert.assertTrue(Arrays.equals(actual, expected), + "Segment generation must preserve the unshredded producer envelope exactly"); + } catch (Exception e) { + throw new IllegalStateException("Failed to validate generated VARIANT segment", e); + } finally { + if (context != null) { + context.close(); + } + if (segment != null) { + segment.destroy(); + } + } + } + + /// Returns a valid object whose a/b/c/eventType offsets are `[12, 10, 8, 0, 14]`; its values are physically + /// encoded in eventType/c/b/a order instead of lexicographic key order. + private static byte[] nonMonotonicObjectEnvelope() { + byte[] metadata = { + 0x11, 0x04, + 0x00, 0x01, 0x02, 0x03, 0x0C, + 'a', 'b', 'c', 'e', 'v', 'e', 'n', 't', 'T', 'y', 'p', 'e' + }; + byte[] value = { + 0x02, 0x04, + 0x00, 0x01, 0x02, 0x03, + 0x0C, 0x0A, 0x08, 0x00, 0x0E, + 0x1D, 'i', 'n', 't', 'e', 'r', 'o', 'p', + 0x0C, 0x03, + 0x0C, 0x02, + 0x0C, 0x01 + }; + return VariantEnvelope.encode(metadata, 0, metadata.length, value, 0, value.length); + } + private static InputStream openResource(String relativePath) { String resourcePath = RESOURCE_DIRECTORY + relativePath; InputStream inputStream = VariantTypeTest.class.getClassLoader().getResourceAsStream(resourcePath); diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkVariantGet.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkVariantGet.java new file mode 100644 index 000000000000..a7670257f725 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkVariantGet.java @@ -0,0 +1,144 @@ +/** + * 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.perf; + +import java.util.concurrent.TimeUnit; +import org.apache.parquet.variant.Variant; +import org.apache.parquet.variant.VariantBuilder; +import org.apache.parquet.variant.VariantObjectBuilder; +import org.apache.pinot.common.utils.VariantUtils; +import org.apache.pinot.common.utils.VariantUtils.ResultType; +import org.apache.pinot.common.utils.VariantUtils.ReusableResult; +import org.apache.pinot.common.utils.VariantUtils.VariantPath; +import org.apache.pinot.spi.utils.VariantEnvelope; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.ChainedOptionsBuilder; +import org.openjdk.jmh.runner.options.CommandLineOptions; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// Measures Pinot's reusable zero-copy Variant cursor against parquet-java object navigation. +/// +/// The 31/32-field cases pin the shared linear-to-binary lookup threshold. First, middle, last, and missing fields +/// expose position-dependent scans; nested values include a second object lookup. Run with `-prof gc` to compare +/// allocation rates in addition to latency. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.NANOSECONDS) +@Fork(value = 2) +@Warmup(iterations = 3, time = 1) +@Measurement(iterations = 5, time = 1) +@State(Scope.Thread) +public class BenchmarkVariantGet { + + public static void main(String[] args) + throws Exception { + ChainedOptionsBuilder options = new OptionsBuilder().parent(new CommandLineOptions(args)) + .include(BenchmarkVariantGet.class.getSimpleName()); + new Runner(options.build()).run(); + } + + @Param({"8", "31", "32", "100"}) + private int _numFields; + + @Param({"first", "middle", "last", "missing"}) + private String _targetPosition; + + @Param({"flat", "nested"}) + private String _valueShape; + + private byte[] _envelope; + private Variant _parquetVariant; + private String _targetKey; + private VariantPath _path; + private ReusableResult _result; + private boolean _nested; + + @Setup(Level.Trial) + public void setUp() { + _nested = "nested".equals(_valueShape); + VariantBuilder builder = new VariantBuilder(); + VariantObjectBuilder objectBuilder = builder.startObject(); + for (int i = _numFields - 1; i >= 0; i--) { + objectBuilder.appendKey(field(i)); + if (_nested) { + VariantObjectBuilder nestedBuilder = objectBuilder.startObject(); + nestedBuilder.appendKey("value"); + nestedBuilder.appendInt(i); + objectBuilder.endObject(); + } else { + objectBuilder.appendInt(i); + } + } + builder.endObject(); + _parquetVariant = builder.build(); + _envelope = VariantEnvelope.encode(_parquetVariant.getMetadataBuffer(), _parquetVariant.getValueBuffer()); + + switch (_targetPosition) { + case "first": + _targetKey = field(0); + break; + case "middle": + _targetKey = field(_numFields / 2); + break; + case "last": + _targetKey = field(_numFields - 1); + break; + case "missing": + _targetKey = "missing"; + break; + default: + throw new IllegalStateException("Unhandled target position: " + _targetPosition); + } + _path = VariantUtils.compilePath("$." + _targetKey + (_nested ? ".value" : "")); + _result = new ReusableResult(); + } + + @Benchmark + public int pinotReusableCursor() { + return VariantUtils.extractInto(_envelope, _path, ResultType.INT, _result) ? _result.getIntValue() : -1; + } + + @Benchmark + public int parquetJava() { + Variant value = _parquetVariant.getFieldByKey(_targetKey); + if (value == null) { + return -1; + } + if (_nested) { + value = value.getFieldByKey("value"); + } + return value != null ? value.getInt() : -1; + } + + private static String field(int index) { + return String.format("field%03d", index); + } +} diff --git a/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java index 450f31cf5a1d..e398dd68bae1 100644 --- a/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java @@ -53,6 +53,7 @@ import org.apache.parquet.variant.VariantArrayBuilder; import org.apache.parquet.variant.VariantBuilder; import org.apache.parquet.variant.VariantObjectBuilder; +import org.apache.pinot.common.utils.VariantUtils; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.data.readers.RecordReader; import org.apache.pinot.spi.utils.VariantEnvelope; @@ -133,6 +134,25 @@ public void testUnshreddedValuePreservesEncodedBuffers() { assertEquals(remainingBytes(decoded.getValue()), expectedValue); } + @Test + public void testUnshreddedNonMonotonicObjectOffsetsArePreserved() + throws Exception { + Variant expected = nonMonotonicObjectVariant(); + File dataFile = writeScalarVariantFile(expected, "non-monotonic-object-offsets.parquet", Map.of()); + try (ParquetNativeRecordReader reader = new ParquetNativeRecordReader()) { + reader.init(dataFile, null, null); + List rows = readAll(reader); + byte[] actualEnvelope = (byte[]) rows.get(2).getValue(VARIANT_FIELD); + byte[] expectedEnvelope = VariantEnvelope.encode(expected.getMetadataBuffer(), expected.getValueBuffer()); + + assertTrue(Arrays.equals(actualEnvelope, expectedEnvelope), + "Unshredded ingestion must preserve producer-owned metadata/value bytes exactly"); + assertEquals(VariantUtils.variantGet(actualEnvelope, "$.a", "INT"), 1); + assertEquals(VariantUtils.variantGet(actualEnvelope, "$.b", "INT"), 2); + assertEquals(VariantUtils.variantGet(actualEnvelope, "$.c", "INT"), 3); + } + } + @Test public void testUnshreddedValueSupportsDirectAndReadOnlyBuffers() { Variant expected = objectVariant("name", "pinot"); @@ -823,6 +843,20 @@ private static byte[] fixedLengthBytes(BigInteger value, int length) { return result; } + /// Returns a valid object whose a/b/c offsets are `[4, 2, 0, 6]` because its int8 values are physically reversed. + private static Variant nonMonotonicObjectVariant() { + byte[] metadata = {0x11, 0x03, 0x00, 0x01, 0x02, 0x03, 'a', 'b', 'c'}; + byte[] value = { + 0x02, 0x03, + 0x00, 0x01, 0x02, + 0x04, 0x02, 0x00, 0x06, + 0x0C, 0x03, + 0x0C, 0x02, + 0x0C, 0x01 + }; + return new Variant(value, metadata); + } + private static Variant objectVariant(Object... keysAndValues) { VariantBuilder builder = new VariantBuilder(); VariantObjectBuilder object = builder.startObject(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java index c21765cc06b1..db0be2171199 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java @@ -146,7 +146,7 @@ private Object extractInternal(@Nullable byte[] variant, boolean tolerant) { boolean present = tolerant ? VariantUtils.tryExtractInto(variant, _path, targetType, _reusableResult) : VariantUtils.extractInto(variant, _path, targetType, _reusableResult); - return present ? _reusableResult.getInternalValue(targetType) : null; + return present ? _reusableResult.toInternalValue(targetType) : null; } private static Operation operation(String canonicalName) {