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/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-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 b806670e1ecf..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,9 +57,16 @@ 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); + // 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 +133,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/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/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-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/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..630265903502 --- /dev/null +++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java @@ -0,0 +1,99 @@ +/** + * 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..88510efd4d88 100644 --- a/pinot-common/pom.xml +++ b/pinot-common/pom.xml @@ -200,6 +200,18 @@ org.apache.pinot pinot-timeseries-spi + + org.apache.parquet + parquet-variant + + + + org.apache.parquet + parquet-column + + + 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..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 @@ -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,146 @@ 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.toExternalValue(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..668fb0cc402d --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/VariantFunctions.java @@ -0,0 +1,100 @@ +/** + * 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..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 @@ -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,8 +329,16 @@ public RelDataType toType(RelDataTypeFactory typeFactory) { return typeFactory.createSqlType(SqlTypeName.VARBINARY); } }, - // 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. + // 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. 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) { @@ -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..c20aca8378c0 --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/VariantUtils.java @@ -0,0 +1,1867 @@ +/** + * 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 { + 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); + 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 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() { + } + + /// 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), + 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. 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; + 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. + /// + ///

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; + } + + 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. + public Object toExternalValue(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. 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 toInternalValue(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.toExternalValue(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.toExternalValue(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) { + 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; + } + 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)) { + 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(PathElement field) { + 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"); + + 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)) { + 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; + 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; + } + + /// 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) { + 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 final boolean _binarySearchSafe; + + 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, hasWellFormedUtf16(field)); + } + + private static PathElement forIndex(int 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/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/function/scalar/VariantFunctionsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java new file mode 100644 index 000000000000..e3d99235a0c9 --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/function/scalar/VariantFunctionsTest.java @@ -0,0 +1,55 @@ +/** + * 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-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..2fbfaae755ec --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/VariantUtilsTest.java @@ -0,0 +1,1005 @@ +/** + * 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.nio.charset.StandardCharsets; +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; +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; +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.DataProvider; +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 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); + 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)); + } + + @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]"); + 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\"}"); + 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")); + + 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\"}"); + } + + @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( + "{\"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.toExternalValue(resultType); + Object internalValue = result.toInternalValue(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) { + 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(); + 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 + | 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-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..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 @@ -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,45 @@ 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..f3d156f7d0d5 --- /dev/null +++ b/pinot-compatibility-verifier/src/main/java/org/apache/pinot/compat/FileContainsOp.java @@ -0,0 +1,129 @@ +/** + * 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..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 @@ -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,28 @@ 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 +100,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 +206,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 +271,68 @@ 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..b3e41a81d24e --- /dev/null +++ b/pinot-compatibility-verifier/src/test/java/org/apache/pinot/compat/BaseOpTest.java @@ -0,0 +1,78 @@ +/** + * 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-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/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/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..f62971f03d25 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/BaseVariantTransformFunction.java @@ -0,0 +1,149 @@ +/** + * 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..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,6 +108,14 @@ public void init(List arguments, Map c _rightTransformFunction = arguments.get(1); DataType leftDataType = _leftTransformFunction.getResultMetadata().getDataType(); DataType rightDataType = _rightTransformFunction.getResultMetadata().getDataType(); + // 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; + return; + } _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..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 @@ -64,6 +64,12 @@ 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) { + // Reject raw VARIANT operands only; preserve existing IN behavior for all other types. + DataType dataType = argument.getResultMetadata().getDataType(); + Preconditions.checkArgument(dataType != DataType.VARIANT, + "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..fc54703ca191 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/IsVariantNullTransformFunction.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.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..03d9892a66a8 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ParseJsonToVariantTransformFunction.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.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..af99b272ac7a --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantExistsTransformFunction.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.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..592bba90a77c --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunction.java @@ -0,0 +1,249 @@ +/** + * 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..c6843f06134f --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/VariantTypeOfTransformFunction.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.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..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 @@ -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,21 @@ public AggregationInfo(AggregationFunction[] functions, BaseProjectOperator p _functions = functions; _projectOperator = projectOperator; _useStarTree = useStarTree; + if (!useStarTree) { + validateRawVariantAggregationInputs(functions, projectOperator); + } + } + + 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() { @@ -382,6 +443,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()) { @@ -417,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/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..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,6 +52,12 @@ public static Comparator getComparator(List throw new BadQueryRequestException("MV expression: " + orderByExpressions.get(i) + " should not be included in the ORDER-BY clause"); } + 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"); + } } 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 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-core/src/test/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProviderTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProviderTest.java index 76b5fa07b7f4..2e3b07f922e0 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProviderTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/filter/predicate/PredicateEvaluatorProviderTest.java @@ -18,8 +18,13 @@ */ package org.apache.pinot.core.operator.filter.predicate; +import java.util.List; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.request.context.predicate.EqPredicate; +import org.apache.pinot.common.request.context.predicate.InPredicate; +import org.apache.pinot.common.request.context.predicate.NotEqPredicate; +import org.apache.pinot.common.request.context.predicate.NotInPredicate; +import org.apache.pinot.common.request.context.predicate.Predicate; import org.apache.pinot.common.request.context.predicate.RangePredicate; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.datasource.DataSource; @@ -31,7 +36,9 @@ import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.data.FieldSpec; 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.assertFalse; @@ -42,6 +49,23 @@ public class PredicateEvaluatorProviderTest { + @Test + public void rawVariantPredicatesProvideActionableErrors() { + ExpressionContext column = ExpressionContext.forIdentifier("payload"); + for (Predicate predicate : List.of( + new EqPredicate(column, "{}"), + new NotEqPredicate(column, "{}"), + new RangePredicate(column, "[1\t\t2]"), + new InPredicate(column, List.of("{}", "[]")), + new NotInPredicate(column, List.of("{}", "[]")))) { + BadQueryRequestException exception = Assert.expectThrows(BadQueryRequestException.class, + () -> 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/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-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..ba8d7781d99a --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/VariantGetTransformFunctionTest.java @@ -0,0 +1,558 @@ +/** + * 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.CommonConstants.NullValuePlaceHolder; +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 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\"}"); + 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) { + 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), path, 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/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). 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..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 @@ -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,32 @@ 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")); + } + + @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-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..1bee8ffdff04 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VariantTypeTest.java @@ -0,0 +1,596 @@ +/** + * 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.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; + + +/// 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 = 6; + + @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); + + 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); + } + + @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 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 { + 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"); + 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"); + 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(), 5L, + "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(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 { + 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 + ", 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("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); + 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 + 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 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); + 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-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/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..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 @@ -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,94 @@ /// - 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 = _extractAll + ? ParquetVariantConverter.createTopLevelVariantConverters(schema) + : ParquetVariantConverter.createTopLevelVariantConverters(schema, _fields::contains); + 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 +369,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..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 @@ -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,18 @@ 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..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 @@ -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,30 @@ 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..0b624cf60c1e --- /dev/null +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/main/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantConverter.java @@ -0,0 +1,324 @@ +/** + * 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 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; +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) { + 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()); + } 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); + 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); + 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..e398dd68bae1 --- /dev/null +++ b/pinot-plugins/pinot-input-format/pinot-parquet/src/test/java/org/apache/pinot/plugin/inputformat/parquet/ParquetVariantRecordReaderTest.java @@ -0,0 +1,939 @@ +/** + * 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.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; +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 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"); + 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.*")); + } + + @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) { + 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; + } + + /// 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(); + 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..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,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.TypeCapabilityValidationVisitor; import org.apache.pinot.query.routing.WorkerManager; import org.apache.pinot.query.routing.WorkerMetadata; @@ -152,12 +153,17 @@ 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(); + if (rootPlanNode.getStageId() == 0) { + TypeCapabilityValidationVisitor.validateResultSchema(rootPlanNode.getDataSchema(), + context.getPlannerContext().getEnvConfig().isNullHandlingEnabled()); + } boolean isIntermediateStage = context.getDispatchablePlanMetadataMap().get(rootPlanNode.getStageId()).getScannedTables().isEmpty(); rootPlanNode.visit(ArrayToMvValidationVisitor.INSTANCE, isIntermediateStage); + 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/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/JoinKeyTypeValidator.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java new file mode 100644 index 000000000000..9939aa8f8388 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/JoinKeyTypeValidator.java @@ -0,0 +1,72 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.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/TypeCapabilityValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitor.java new file mode 100644 index 000000000000..b6b216265ab9 --- /dev/null +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitor.java @@ -0,0 +1,199 @@ +/** + * 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.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; +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; + + +/// 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 TypeCapabilityValidationVisitor() { + } + + @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) { + 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); + } + + /// 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()) { + DataSchema.ColumnDataType dataType = getLogicalType(operand, inputSchema); + if (!dataType.supportsDirectAggregation()) { + throw unsupported("Aggregate function " + aggCall.getFunctionName(), dataType); + } + } + } + } + + /// 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(); + for (RelFieldCollation collation : node.getCollations()) { + int fieldIndex = collation.getFieldIndex(); + DataSchema.ColumnDataType dataType = dataSchema.getColumnDataType(fieldIndex); + if (!dataType.supportsOrdering()) { + throw unsupported("ORDER BY", dataType); + } + } + 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('_', ' '), dataType); + } + } + } + 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) { + try { + JoinKeyTypeValidator.validate(node, leftSchema, rightSchema); + } catch (IllegalArgumentException e) { + throw new QueryException(QueryErrorCode.QUERY_PLANNING, e.getMessage(), e); + } + } + + @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", dataType); + } + } + for (RelFieldCollation collation : node.getCollations()) { + DataSchema.ColumnDataType dataType = inputSchema.getColumnDataType(collation.getFieldIndex()); + if (!dataType.supportsOrdering()) { + throw unsupported("Window ORDER BY", dataType); + } + } + 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, 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/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/TypeCapabilityValidationVisitorTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitorTest.java new file mode 100644 index 000000000000..9a891f20727f --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/validation/TypeCapabilityValidationVisitorTest.java @@ -0,0 +1,276 @@ +/** + * 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 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 = + 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(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, + () -> TypeCapabilityValidationVisitor.validateResultSchema(VARIANT_SCHEMA, false)); + Assert.assertTrue(exception.getMessage().contains("requires query null handling")); + + TypeCapabilityValidationVisitor.validateResultSchema(VARIANT_SCHEMA, true); + TypeCapabilityValidationVisitor.validateResultSchema(TYPED_EXTRACTION_SCHEMA, false); + } + + @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(TypeCapabilityValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("raw VARIANT")); + } + } + + @Test + public void testAllowsVariantUnionAll() { + setOp(SetOpNode.SetOpType.UNION, true).visit(TypeCapabilityValidationVisitor.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(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(TypeCapabilityValidationVisitor.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(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"}, + 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(TypeCapabilityValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Raw VARIANT values do not support ASOF JOIN match keys")); + } + + @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(TypeCapabilityValidationVisitor.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(TypeCapabilityValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Aggregate function COUNT")); + } + + @Test + public void testAllowsRawVariantCount() { + aggregate("COUNT", false, VARIANT_SCHEMA).visit(TypeCapabilityValidationVisitor.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(TypeCapabilityValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("GROUP BY")); + Assert.assertTrue(exception.getMessage().contains("raw VARIANT")); + } + + @Test + public void testUsesLogicalTypeInsteadOfVariantStorageType() { + DataSchema bytesSchema = + new DataSchema(new String[]{"bytes"}, new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.BYTES}); + aggregate("ANYVALUE", false, bytesSchema).visit(TypeCapabilityValidationVisitor.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(TypeCapabilityValidationVisitor.INSTANCE, null); + } + + @Test + public void testValidatesWindowAggregateInputs() { + QueryException exception = Assert.expectThrows(QueryException.class, + () -> window("SUM").visit(TypeCapabilityValidationVisitor.INSTANCE, null)); + Assert.assertTrue(exception.getMessage().contains("Aggregate function SUM")); + + 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(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(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(TypeCapabilityValidationVisitor.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 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()); + } + + 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..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,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.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; @@ -91,6 +92,11 @@ public class AggregateOperator extends MultiStageOperator { public AggregateOperator(OpChainExecutionContext context, MultiStageOperator input, AggregateNode node) { super(context); + 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()); int numFunctions = _aggFunctions.length; @@ -143,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/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java index 1a6ea42c2f72..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 @@ -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,8 @@ 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.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; @@ -66,11 +69,28 @@ 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"); + JoinKeyTypeValidator.validate(node, leftSchema, rightSchema); _leftKeySelector = KeySelectorFactory.getKeySelector(leftKeys); _rightKeySelector = KeySelectorFactory.getKeySelector(node.getRightKeys()); _rightTable = createLookupTable(leftKeys, leftSchema); @@ -79,18 +99,55 @@ 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"); + JoinKeyTypeValidator.validate(node, 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(); 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..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; @@ -82,6 +82,15 @@ public MultistageGroupByExecutor(int[] groupKeyIds, AggregationFunction[] aggFun _aggType = aggType; _leafReturnFinalResult = leafReturnFinalResult; _resultSchema = resultSchema; + for (int i = 0; i < groupKeyIds.length; i++) { + ColumnDataType dataType = inputSchema != null ? inputSchema.getColumnDataType(groupKeyIds[i]) + : resultSchema.getColumnDataType(i); + if (!dataType.supportsEquality() || !dataType.supportsHashing()) { + 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"); + } + } int maxInitialResultHolderCapacity = getResolvedMaxInitialResultHolderCapacity(opChainMetadata, nodeHint); 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/SortOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/SortOperator.java index 82db0410b46f..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 @@ -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; @@ -27,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; @@ -73,6 +75,12 @@ public SortOperator(OpChainExecutionContext context, MultiStageOperator input, S // - There is no collation // - Input is already sorted List collations = node.getCollations(); + for (RelFieldCollation collation : collations) { + 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; _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..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 @@ -56,6 +56,12 @@ 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) { + 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"); + } } @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..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,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.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; @@ -108,6 +109,7 @@ public class WindowAggregateOperator extends MultiStageOperator { public WindowAggregateOperator(OpChainExecutionContext context, MultiStageOperator input, DataSchema inputSchema, WindowNode node) { super(context); + TypeCapabilityValidationVisitor.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..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 @@ -22,6 +22,7 @@ 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; @@ -37,13 +38,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(); + JoinKeyTypeValidator.validate(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); 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..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 @@ -115,7 +115,11 @@ 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); + // 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); } _isNotIn = isNotIn; } @@ -193,7 +197,12 @@ public Predicate(List operands, DataSchema dataSchema, IntPredica ColumnDataType lhsType = _lhs.getResultType(); ColumnDataType rhsType = _rhs.getResultType(); - if (lhsType == rhsType) { + // 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; _commonCastType = null; } else { 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 new file mode 100644 index 000000000000..2075d0afe29e --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/LiteralParseJsonOperand.java @@ -0,0 +1,93 @@ +/** + * 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..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 @@ -20,11 +20,18 @@ 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() { } @@ -43,6 +50,24 @@ 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); + } + 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) { + // 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); + } switch (functionCall.getFunctionName()) { case "AND": Preconditions.checkState(numOperands >= 2, "AND takes >=2 arguments, got: %s", numOperands); @@ -81,4 +106,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 new file mode 100644 index 000000000000..db0be2171199 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/operands/VariantOperand.java @@ -0,0 +1,219 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.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, _reusableResult); + 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.toInternalValue(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..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,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"); + 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 de49e3c53443..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 @@ -40,6 +40,16 @@ public SetOperator(OpChainExecutionContext opChainExecutionContext, List inputOperators, DataSchema dataSchema) { super(opChainExecutionContext, inputOperators, dataSchema); + 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 60ab3e503564..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 @@ -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,63 @@ 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 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/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..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 @@ -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,27 @@ 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")); + } + + @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 59aeccd6d5d5..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 @@ -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,29 @@ 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 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/WindowAggregateOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperatorTest.java index 02fa104d6578..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 @@ -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,58 @@ 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[]{"eventType", "eventId"}, new ColumnDataType[]{STRING, STRING}); + DataSchema resultSchema = + 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(), aggCalls, RANGE, + Integer.MIN_VALUE, Integer.MAX_VALUE, input); + + 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 public void testShouldWindowAggregateOverSingleInputBlock() { // Given: @@ -3524,6 +3577,14 @@ 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 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-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..5b1e146f7cd3 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/operands/FilterOperandTest.java @@ -0,0 +1,98 @@ +/** + * 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 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 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)); + + @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 = + 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")); + } + + @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")) { + 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-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..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 @@ -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,23 @@ 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); + + 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/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..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 @@ -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); } } @@ -1429,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), @@ -1452,6 +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); + // VARIANT columns are validated up front against their effective strategy (default or explicit). if (columnStrategy == UpsertConfig.Strategy.INCREMENT) { Preconditions.checkState(fieldSpec.getDataType().getStoredType().isNumeric(), @@ -1771,6 +1799,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 +1831,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 +1877,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 +1907,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 +1922,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. /// @@ -1961,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<>(); @@ -1970,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; } @@ -2060,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); @@ -2068,10 +2158,12 @@ 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) { - 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/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..a8f0067deb28 --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/VariantTableConfigValidationTest.java @@ -0,0 +1,387 @@ +/** + * 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.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; +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 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) + .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"); + } + + @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) + .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..eedadc2e5b27 --- /dev/null +++ b/pinot-spi/VARIANT_DESIGN.md @@ -0,0 +1,338 @@ + +# 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. +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 + +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. +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 + +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, 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. 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. + +## 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..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,6 +1019,75 @@ public Object convert(Object value, PinotDataType sourceType) { } }, + /// Pinot's external representation of a VARIANT value: a validated PVAR envelope in a `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 +1908,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 +2116,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..f908cf99fabc --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/VariantEnvelope.java @@ -0,0 +1,243 @@ +/** + * 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 +/// `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 `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; + 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 `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 [#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( + "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 [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..976c71912cfd --- /dev/null +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/VariantQuickStart.java @@ -0,0 +1,113 @@ +/** + * 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 000000000000..0d901fde2b28 Binary files /dev/null and b/pinot-tools/src/main/resources/examples/batch/variantEvents/rawdata/variantEvents_data.parquet differ diff --git a/pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_offline_table_config.json b/pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_offline_table_config.json new file mode 100644 index 000000000000..99937d12be32 --- /dev/null +++ b/pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_offline_table_config.json @@ -0,0 +1,39 @@ +{ + "tableName": "variantEvents", + "tableType": "OFFLINE", + "segmentsConfig": { + "segmentPushType": "APPEND", + "replication": "1", + "timeColumnName": "eventTime" + }, + "tenants": {}, + "fieldConfigList": [ + { + "name": "payload", + "encodingType": "RAW", + "indexes": { + "forward": { + "compressionCodec": "ZSTANDARD" + } + } + } + ], + "tableIndexConfig": { + "loadMode": "MMAP", + "nullHandlingEnabled": true, + "invertedIndexColumns": [ + "eventType" + ] + }, + "ingestionConfig": { + "transformConfigs": [ + { + "columnName": "eventType", + "transformFunction": "variant_get(payload, '$.eventType', 'STRING')" + } + ] + }, + "metadata": { + "customConfigs": {} + } +} diff --git a/pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_schema.json b/pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_schema.json new file mode 100644 index 000000000000..f7bce636bbe2 --- /dev/null +++ b/pinot-tools/src/main/resources/examples/batch/variantEvents/variantEvents_schema.json @@ -0,0 +1,26 @@ +{ + "schemaName": "variantEvents", + "enableColumnBasedNullHandling": true, + "dimensionFieldSpecs": [ + { + "name": "eventId", + "dataType": "STRING" + }, + { + "name": "eventType", + "dataType": "STRING" + }, + { + "name": "payload", + "dataType": "VARIANT" + } + ], + "dateTimeFieldSpecs": [ + { + "name": "eventTime", + "dataType": "LONG", + "format": "1:MILLISECONDS:EPOCH", + "granularity": "1:MILLISECONDS" + } + ] +} diff --git a/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/QuickStartTest.java b/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/QuickStartTest.java index fe5b386861ab..1bc8c3f10bfa 100644 --- a/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/QuickStartTest.java +++ b/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/QuickStartTest.java @@ -26,6 +26,7 @@ import org.apache.pinot.tools.Quickstart; import org.apache.pinot.tools.RealtimeQuickStart; import org.apache.pinot.tools.RealtimeQuickStartWithMinion; +import org.apache.pinot.tools.VariantQuickStart; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; @@ -72,6 +73,12 @@ public void testMatchStringToCommand() assertEquals(quickStartClassFor("MULTI_CLUSTER"), MultiClusterQuickstart.class); assertEquals(quickStartClassFor("MULTICLUSTER"), MultiClusterQuickstart.class); + + assertEquals(quickStartClassFor("VARIANT"), VariantQuickStart.class); + assertEquals(quickStartClassFor("OFFLINE_VARIANT"), VariantQuickStart.class); + assertEquals(quickStartClassFor("OFFLINE-VARIANT"), VariantQuickStart.class); + assertEquals(quickStartClassFor("BATCH_VARIANT"), VariantQuickStart.class); + assertEquals(quickStartClassFor("BATCH-VARIANT"), VariantQuickStart.class); } /// The batch-side quickstarts that only differed by their sample queries were merged into [Quickstart]; their types diff --git a/pom.xml b/pom.xml index ba86e35c380a..ece240bbf92b 100644 --- a/pom.xml +++ b/pom.xml @@ -841,6 +841,11 @@ parquet-hadoop ${parquet.version} + + org.apache.parquet + parquet-variant + ${parquet.version} + org.apache.orc orc-core